diff --git a/.github/workflows/backend-ci.yml b/.github/workflows/backend-ci.yml index a9cf349937..bb5cf692bc 100644 --- a/.github/workflows/backend-ci.yml +++ b/.github/workflows/backend-ci.yml @@ -20,7 +20,7 @@ jobs: cache-dependency-path: backend/go.sum - name: Verify Go version run: | - go version | grep -q 'go1.26.4' + go version | grep -q 'go1.26.5' - name: Unit tests working-directory: backend run: make test-unit @@ -60,7 +60,7 @@ jobs: cache-dependency-path: backend/go.sum - name: Verify Go version run: | - go version | grep -q 'go1.26.4' + go version | grep -q 'go1.26.5' - name: golangci-lint uses: golangci/golangci-lint-action@v9 with: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f8c5dcfbfa..2ba01833c1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -115,7 +115,7 @@ jobs: - name: Verify Go version run: | - go version | grep -q 'go1.26.4' + go version | grep -q 'go1.26.5' # Docker setup for GoReleaser - name: Set up QEMU diff --git a/.github/workflows/security-scan.yml b/.github/workflows/security-scan.yml index ab3305ab6e..96a7ae4edd 100644 --- a/.github/workflows/security-scan.yml +++ b/.github/workflows/security-scan.yml @@ -23,7 +23,7 @@ jobs: cache-dependency-path: backend/go.sum - name: Verify Go version run: | - go version | grep -q 'go1.26.4' + go version | grep -q 'go1.26.5' - name: Run govulncheck working-directory: backend run: | diff --git a/Dockerfile b/Dockerfile index bae531ac6e..17b631d5a0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,3 +1,4 @@ +# syntax=docker/dockerfile:1.7 # ============================================================================= # Sub2API Multi-Stage Dockerfile # ============================================================================= @@ -7,16 +8,18 @@ # ============================================================================= ARG NODE_IMAGE=node:24-alpine -ARG GOLANG_IMAGE=golang:1.26.4-alpine +ARG GOLANG_IMAGE=golang:1.26.5-alpine ARG ALPINE_IMAGE=alpine:3.21 ARG POSTGRES_IMAGE=postgres:18-alpine ARG GOPROXY=https://goproxy.cn,direct ARG GOSUMDB=sum.golang.google.cn +ARG NPM_CONFIG_REGISTRY= # ----------------------------------------------------------------------------- # Stage 1: Frontend Builder # ----------------------------------------------------------------------------- FROM ${NODE_IMAGE} AS frontend-builder +ARG NPM_CONFIG_REGISTRY WORKDIR /app/frontend @@ -25,7 +28,9 @@ RUN corepack enable && corepack prepare pnpm@9 --activate # Install dependencies first (better caching) COPY frontend/package.json frontend/pnpm-lock.yaml ./ -RUN pnpm install --frozen-lockfile +RUN --mount=type=cache,id=sub2api-pnpm-store,target=/root/.local/share/pnpm/store \ + if [ -n "${NPM_CONFIG_REGISTRY}" ]; then pnpm config set registry "${NPM_CONFIG_REGISTRY}"; fi && \ + pnpm install --frozen-lockfile --prefer-offline # Copy frontend source and build. # LegalDocumentView.vue (admin-compliance gate) build-time imports diff --git a/README.md b/README.md index eb24b0d3d6..661cd8e7eb 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,11 @@ Please read the following carefully before using this project: + + + + + @@ -41,6 +46,11 @@ Please read the following carefully before using this project: + + + + + @@ -145,6 +155,14 @@ Model authenticity: no content intervention or secondary filtering — experienc + + + + +
CCTK.AIThanks to CCTK.AI for sponsoring this project! CCTK.AI is an AI API gateway focused on stability and cost-effectiveness, offering fast relay services for Claude, OpenAI, Gemini, and other popular models. It works seamlessly with Claude Code, Codex, and other mainstream coding tools, delivering the same model capabilities at a fraction of the official cost. Register via this link for faster, more stable, and more affordable AI API access.
openmodel One API, every top model! OpenModel is a production-grade, high-availability AI API gateway that makes your applications truly fast and stable: automatic failover, smart routing to the best-performing channel, and a production-grade SLA. An SLA that far surpasses any single provider — making stability your core competitive advantage. Works directly with Claude Code, Codex, and Gemini CLI. Register via this link to get started.Thanks to ETok.ai for sponsoring this project! ETok.ai is dedicated to building a one-stop AI programming tool service platform. We offer professional Claude Code packages and technical community services, with support for Google Gemini and OpenAI Codex. Through carefully designed plans and a professional tech community, we provide developers with reliable service guarantees and continuous technical support, making AI-assisted programming a true productivity tool. Click here to register!
APIKEY.FUNThanks to APIKEY.FUN for sponsoring this project! APIKEY.FUN is one of the core contributors to the sub2api open-source project, dedicated to providing open, stable, and cost-effective AI API access. The platform supports API relay services for Claude, OpenAI, Gemini, and other popular models, with pricing starting from as low as 7% of the original rate. Register via the exclusive link: APIKEY to enjoy a permanent 5% discount on all recharges.
AIGoCode Thanks to AIGoCode for sponsoring this project! AIGoCode is an all-in-one platform that integrates Claude Code, Codex, and the latest Gemini models, providing you with stable, efficient, and highly cost-effective AI coding services. The platform offers flexible subscription plans, zero risk of account suspension, direct access with no VPN required, and lightning-fast responses. AIGoCode has prepared a special benefit for sub2api users: if you register via this link, you'll receive an extra 10% bonus credit on your first top-up!
fastaitoken🎉 Thanks to FastAIToken for sponsoring this project! FastAIToken is an AI API aggregation platform for developers, supporting mainstream large models such as OpenAI, Claude, and Gemini. Top-up at 1:1 — 1 CNY = 1 USD of API credit — letting developers use the world's leading large model services at lower cost and with greater convenience.
+ +🚀 The platform offers a variety of channels to choose from: an ultra-low-price 0.02x OpenAI promotional group (limited time), groups as low as 0.25x OpenAI, 0.7x Claude with 95% fixed cache, and a 1.2x Claude Max channel. It also provides a public status page showing real-time availability, latency, and operating status of each group for transparent and reliable service, plus 7×24 human technical support (not bots) with fast responses to developer needs. +
## Overview diff --git a/README_CN.md b/README_CN.md index a93056282b..88c8ce11b1 100644 --- a/README_CN.md +++ b/README_CN.md @@ -32,6 +32,11 @@ + + + + + @@ -42,6 +47,11 @@ + + + + + @@ -148,6 +158,14 @@ + + + + +
CCTK.AI感谢 CCTK.AI 赞助了本项目!CCTK.AI 是一个专注于稳定与性价比的 AI API 网关平台,提供 Claude、OpenAI、Gemini 等主流模型的高速中转服务,无缝兼容 Claude Code、Codex 等主流编程工具,以远低于官方的成本获得同等的模型能力。点击此链接注册,即刻体验更快、更稳、更省的 AI API 接入。
openmodel 一个API,顶级模型随便用!OpenModel 专注于生产级、高可用的 AI API 网关,让你的应用真正做到高速稳定:自动故障转移、智能选最优渠道、生产级 SLA 保障。远超单一供应商的 SLA,让稳定性成为您的核心竞争力。感谢 ETok.ai 赞助了本项目!ETok.ai 致力于打造一站式 AI 编程工具服务平台。我们提供 Claude Code 专业套餐及技术社群服务,同时支持 Google Gemini 和 OpenAI Codex。通过精心设计的套餐方案和专业的技术社群,为开发者提供稳定的服务保障和持续的技术支持,让 AI 辅助编程真正成为开发者的生产力工具。点击这里注册!
APIKEY.FUN感谢 APIKEY.FUN 赞助了本项目!APIKEY.FUN 是 sub2api 开源项目的核心贡献者之一,致力于提供开放、稳定、高性价比的 AI API 接入服务。平台支持 Claude、OpenAI、Gemini 等热门模型的 API 中转服务,价格低至官方原价的 7%。通过专属链接 APIKEY 注册,可享受所有充值永久 95 折优惠。
AIGoCode 感谢 AIGoCode 赞助了本项目!AIGoCode 是一站式集成 Claude Code、Codex 以及最新 Gemini 模型的综合平台,为您提供稳定、高效、高性价比的 AI 编程服务。平台提供灵活的订阅方案,零封号风险,免 VPN 直连,响应极速。AIGoCode 为 sub2api 用户准备了专属福利:通过此链接注册,首次充值可额外获得 10% 赠送额度!
fastaitoken🎉 感谢 FastAIToken 对本项目的赞助! FastAIToken 是面向开发者的 AI API 聚合平台,支持 OpenAI、Claude、Gemini 等主流大模型,充值 1:1,1 元 = 1 美元 API 额度,让开发者以更低成本、更便捷地使用全球领先的大模型服务。
+ +🚀 平台提供多种渠道自由选择:超级低价的0.02x OpenAI 福利分组(限时)、低至 0.25x OpenAI 分组、0.7x Claude 95%固定缓存、1.2x Claude Max 渠道;同时提供公开状态页,实时展示各分组的可用率、延迟及运行状态,服务透明可靠,并提供 7×24 小时真人技术支持(非机器人),快速响应开发者需求。 +
diff --git a/README_JA.md b/README_JA.md index bd154a9a6b..21d070e397 100644 --- a/README_JA.md +++ b/README_JA.md @@ -31,6 +31,11 @@ + + + + + @@ -41,6 +46,11 @@ + + + + + @@ -143,6 +153,14 @@ + + + + +
CCTK.AICCTK.AI のご支援に感謝します!CCTK.AI は安定性とコストパフォーマンスにこだわった AI API ゲートウェイで、Claude、OpenAI、Gemini など主要モデルの高速中継サービスを提供しています。Claude Code や Codex などの主要なコーディングツールにシームレスに対応し、公式価格を大きく下回るコストで同等のモデル能力を利用できます。こちらのリンクから登録して、より速く、より安定した、よりお得な AI API 接続をお試しください。
openmodel 1つの API で、トップモデルを使い放題!OpenModel は本番環境グレードで高可用性の AI API ゲートウェイに特化し、アプリを真に高速・安定させます:自動フェイルオーバー、最適なチャネルへのスマートルーティング、本番グレードの SLA 保証。単一プロバイダーをはるかに上回る SLA で、安定性をあなたの核心的な競争力にします。ETok.ai のご支援に感謝します!ETok.ai はワンストップ AI プログラミングツールサービスプラットフォームの構築に取り組んでいます。Claude Code の専用プランと技術コミュニティサービスを提供し、Google Gemini や OpenAI Codex もサポートしています。丁寧に設計されたプランと専門的な技術コミュニティを通じて、開発者に安定したサービス保証と継続的な技術サポートを提供し、AI アシスト プログラミングを真の生産性向上ツールにします。こちらから登録!
APIKEY.FUNAPIKEY.FUN のご支援に感謝します!APIKEY.FUN は sub2api オープンソースプロジェクトのコアコントリビューターの一つであり、オープンで安定した、コストパフォーマンスに優れた AI API アクセスサービスの提供に取り組んでいます。プラットフォームは Claude、OpenAI、Gemini など人気モデルの API 中継サービスをサポートし、価格は公式料金のわずか 7% から。専用リンク APIKEY から登録すると、すべてのチャージで永久 5% 割引をご利用いただけます。
AIGoCode AIGoCode のご支援に感謝します!AIGoCode は Claude Code、Codex、最新の Gemini モデルを統合したオールインワンプラットフォームで、安定的かつ効率的でコストパフォーマンスに優れた AI コーディングサービスを提供します。柔軟なサブスクリプションプラン、アカウント停止リスクゼロ、VPN 不要の直接アクセス、超高速レスポンスが特長です。AIGoCode は sub2api ユーザー向けに特別特典を用意しています:こちらのリンクから登録すると、初回チャージ時に 10% のボーナスクレジットを追加プレゼント!
fastaitoken🎉 FastAIToken のご支援に感謝します!FastAIToken は開発者向けの AI API アグリゲーションプラットフォームで、OpenAI、Claude、Gemini などの主要な大規模モデルに対応しています。チャージは 1:1(1 元 = 1 米ドル分の API クレジット)で、開発者がより低コスト・より手軽に世界トップクラスの大規模モデルサービスを利用できます。
+ +🚀 プラットフォームでは多彩なチャネルを自由に選択できます:超低価格の 0.02x OpenAI 特典グループ(期間限定)、最低 0.25x の OpenAI グループ、0.7x Claude(95% 固定キャッシュ)、1.2x Claude Max チャネル。さらに、各グループの可用率・レイテンシ・稼働状況をリアルタイムで表示する公開ステータスページを提供し、透明で信頼できるサービスを実現。7×24 時間の有人テクニカルサポート(ボットではありません)により、開発者のニーズに迅速に対応します。 +
## 概要 diff --git a/REDIS_SCAN_ARCHITECTURE_OPTIMIZATION.md b/REDIS_SCAN_ARCHITECTURE_OPTIMIZATION.md deleted file mode 100644 index 7d95419a39..0000000000 --- a/REDIS_SCAN_ARCHITECTURE_OPTIMIZATION.md +++ /dev/null @@ -1,568 +0,0 @@ -# Redis SCAN 架构优化执行文档 - -本文是架构级执行文档,只覆盖三个目标: - -1. 账号并发活跃负载查询不再通过 Redis keyspace `SCAN` 发现账号。 -2. 账号/用户并发槽过期清理、启动遗留槽清理不再通过 Redis keyspace `SCAN` 发现 key。 -3. 用户消息队列孤儿锁清理不再通过 Redis keyspace `SCAN` 发现 lock key。 - -不覆盖旁路录制、OpenAI failover、日志量、业务限流配置调参。不要把本文扩展成短期止血方案。 - -## 成功标准 - -实现完成后必须同时满足: - -- `backend/internal/repository/concurrency_cache.go` 中不得再调用 `c.rdb.Scan(...)`。 -- `backend/internal/repository/user_msg_queue_cache.go` 中不得再调用 `c.rdb.Scan(...)`。 -- `backend/internal/service/user_msg_queue_service.go` 中不得再出现 `ScanLockKeys` 接口调用。 -- `GetActiveAccountLoadMap` 只读显式维护的 Redis 索引,不扫描 Redis keyspace。 -- `CleanupExpiredAccountSlotKeys` 只处理显式索引中的候选账号,不扫描 Redis keyspace。 -- `CleanupStaleProcessSlots` 不扫描 Redis keyspace;它必须基于显式索引清理,或只依赖 TTL/score 自然过期。 -- UMQ cleanup worker 只读 `umq:lock:index` 候选项,不扫描 `umq:{*}:lock`。 -- 主业务并发限制仍以原账号/用户 slot key 为准,不能以索引为准。索引只能用于发现候选对象、监控和清理。 - -执行完必须用下面命令确认生产代码没有遗留扫描: - -```powershell -rg -n "\.Scan\(" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go -rg -n "ScanLockKeys|scanAccountIDsByPrefix|cleanupSlotsByPattern|deleteKeysByPattern|umqScanPattern" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go -``` - -第一条必须无输出。第二条必须无生产函数残留;测试文件不在此检查范围。 - -## 不可违反的约束 - -- 不要用 `KEYS` 替代 `SCAN`。 -- 不要把全量 Redis keyspace 扫描移动到另一个函数、goroutine、启动流程或管理接口里。 -- 不要在请求路径、后台 worker、启动流程中做 Redis keyspace pattern enumeration。 -- 不要在 Redis Lua 脚本里同时操作“全局索引 key”和“账号/用户局部 key”。项目代码当前有 Redis Cluster 兼容要求,这种写法会在 Cluster 下触发 CROSSSLOT。 -- 索引更新失败不得改变主业务 acquire/release 的成功结果。索引是 best-effort discovery structure,不是并发正确性的来源。 -- 不能因为索引缺失而拒绝用户请求。索引缺失最多影响 Ops 实时视图和后台提前清理;原 slot/wait key 的 TTL 必须保证最终自愈。 - -## 新增 Redis Key - -### 并发索引 - -在 `backend/internal/repository/concurrency_cache.go` 增加常量: - -```go -const ( - accountActiveIndexKey = "concurrency:account:active_index" // ZSET member=accountID, score=expireAtUnixSeconds - userActiveIndexKey = "concurrency:user:active_index" // ZSET member=userID, score=expireAtUnixSeconds - - activeIndexCleanupBatchSize = 1000 - activeIndexPipelineChunkSize = 500 -) -``` - -语义: - -- `accountActiveIndexKey` 记录“可能有账号槽位或账号等待计数”的账号 ID。 -- `userActiveIndexKey` 记录“可能有用户槽位或用户等待计数”的用户 ID。 -- ZSET score 是候选对象的保守过期时间,单位为 Unix 秒。 -- member 必须是十进制 ID 字符串,不要存完整 Redis key。 -- 索引允许短暂 stale;读索引后必须二次查询真实 slot/wait key。 - -score 规则: - -- 成功获取账号槽位:score = Redis 当前秒 + `slotTTLSeconds`。 -- 成功增加账号等待计数:score = Redis 当前秒 + `waitQueueTTLSeconds`。 -- 成功获取用户槽位:score = Redis 当前秒 + `slotTTLSeconds`。 -- 成功增加用户等待计数:score = Redis 当前秒 + `waitQueueTTLSeconds`。 -- release/decrement 后如果真实 slot count 和 wait count 都为 0,则从索引 `ZREM`。 -- release/decrement 后如果仍有 slot 或 wait,则重新 `ZADD` 一个新的保守过期时间。 - -### UMQ 锁索引 - -在 `backend/internal/repository/user_msg_queue_cache.go` 增加常量: - -```go -const ( - umqLockIndexKey = "umq:lock:index" // ZSET member=accountID, score=lockExpireAtUnixMs - umqLockIndexCleanupBatchSize = 1000 -) -``` - -语义: - -- `umqLockIndexKey` 记录“可能存在 UMQ lock”的账号 ID。 -- ZSET score 是 lock 的预计过期时间,单位为 Unix 毫秒。 -- member 必须是十进制 accountID 字符串。 -- 索引只用于 cleanup worker 找候选 lock。锁是否存在、是否孤儿,必须再查 `umq:{accountID}:lock`。 - -## 第一部分:并发活跃索引 - -修改文件:`backend/internal/repository/concurrency_cache.go`。 - -### 1.1 增加 Redis 时间 helper - -新增 helper,所有索引 score 使用 Redis server time,不用本机时间: - -```go -func (c *concurrencyCache) redisUnixSeconds(ctx context.Context) (int64, error) { - now, err := c.rdb.Time(ctx).Result() - if err != nil { - return 0, fmt.Errorf("redis TIME: %w", err) - } - return now.Unix(), nil -} -``` - -不要在 Lua 脚本里写全局索引,避免 CROSSSLOT。 - -### 1.2 增加索引 touch/remove/refresh helper - -新增以下 helper。名字可以微调,但行为不能改。 - -```go -func (c *concurrencyCache) touchAccountActiveIndex(ctx context.Context, accountID int64, ttlSeconds int) { - c.touchActiveIndex(ctx, accountActiveIndexKey, accountID, ttlSeconds) -} - -func (c *concurrencyCache) touchUserActiveIndex(ctx context.Context, userID int64, ttlSeconds int) { - c.touchActiveIndex(ctx, userActiveIndexKey, userID, ttlSeconds) -} - -func (c *concurrencyCache) touchActiveIndex(ctx context.Context, indexKey string, id int64, ttlSeconds int) { - if c == nil || c.rdb == nil || id <= 0 || ttlSeconds <= 0 { - return - } - now, err := c.redisUnixSeconds(ctx) - if err != nil { - return - } - _ = c.rdb.ZAdd(ctx, indexKey, redis.Z{ - Score: float64(now + int64(ttlSeconds)), - Member: strconv.FormatInt(id, 10), - }).Err() -} -``` - -索引维护是 best-effort,所以 helper 内部吞掉错误。不要把索引错误返回给 acquire/release 调用方。 - -再新增 refresh helper: - -```go -func (c *concurrencyCache) refreshAccountActiveIndex(ctx context.Context, accountID int64) { - // 真实状态以 accountSlotKey(accountID) 和 accountWaitKey(accountID) 为准。 - // 先清理该账号 slot 中过期成员,再读 ZCARD 和 GET wait。 - // 如果 slotCount == 0 && waitCount <= 0:ZREM accountActiveIndexKey accountID。 - // 否则:ZADD accountActiveIndexKey accountID,score = now + maxRelevantTTL。 -} - -func (c *concurrencyCache) refreshUserActiveIndex(ctx context.Context, userID int64) { - // 真实状态以 userSlotKey(userID) 和 waitQueueKey(userID) 为准。 - // 行为同 refreshAccountActiveIndex。 -} -``` - -实现要求: - -- `refresh*` 必须 best-effort,不能向 release/decrement 返回索引错误。 -- `waitCount` 读取 `redis.Nil` 时按 0 处理。 -- `waitCount < 0` 必须按 0 处理。 -- `slotCount > 0` 时 score 至少延长 `slotTTLSeconds`。 -- `waitCount > 0` 时 score 至少延长 `waitQueueTTLSeconds`。 -- 两者都存在时使用更大的 TTL。 - -### 1.3 修改账号写路径 - -修改 `AcquireAccountSlot`: - -```go -result, err := acquireScript.Run(...).Int() -if err != nil { return false, err } -if result == 1 { - c.touchAccountActiveIndex(ctx, accountID, c.slotTTLSeconds) -} -return result == 1, nil -``` - -修改 `ReleaseAccountSlot`: - -```go -if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil { - return err -} -c.refreshAccountActiveIndex(ctx, accountID) -return nil -``` - -修改 `IncrementAccountWaitCount`: - -```go -result, err := incrementAccountWaitScript.Run(...).Int() -if err != nil { return false, err } -if result == 1 { - c.touchAccountActiveIndex(ctx, accountID, c.waitQueueTTLSeconds) -} -return result == 1, nil -``` - -修改 `DecrementAccountWaitCount`: - -```go -_, err := decrementWaitScript.Run(...).Result() -if err == nil { - c.refreshAccountActiveIndex(ctx, accountID) -} -return err -``` - -### 1.4 修改用户写路径 - -同账号路径,修改: - -- `AcquireUserSlot` -- `ReleaseUserSlot` -- `IncrementWaitCount` -- `DecrementWaitCount` - -用户索引使用 `userActiveIndexKey`。 - -### 1.5 重写 GetActiveAccountLoadMap - -删除 `scanAccountIDsByPrefix` 和 `parseAccountIDFromPrefixedKey` 的生产调用。`GetActiveAccountLoadMap` 必须改成: - -1. 获取 Redis 当前秒。 -2. `ZRemRangeByScore(accountActiveIndexKey, "-inf", strconv.FormatInt(now, 10))` 删除过期候选。 -3. `ZRangeByScore(accountActiveIndexKey, &redis.ZRangeBy{Min: strconv.FormatInt(now+1, 10), Max: "+inf"})` 获取候选账号 ID。 -4. 解析 member 为 `int64`,非法 member 记录到待删除列表。 -5. 分块 pipeline,块大小 `activeIndexPipelineChunkSize`。 -6. 对每个候选账号执行: - - `ZRemRangeByScore(accountSlotKey(id), "-inf", cutoffUnixSeconds)` - - `ZCard(accountSlotKey(id))` - - `Get(accountWaitKey(id))` -7. 构造结果时只返回 `currentConcurrency > 0 || waitingCount > 0` 的账号。 -8. 对真实状态为空或 member 非法的账号执行 `ZREM accountActiveIndexKey member`。 -9. 对真实状态仍活跃但 index score 已接近过期的账号,调用 `touchAccountActiveIndex` 刷新。 - -禁止: - -- 禁止再扫 `concurrency:account:*`。 -- 禁止再扫 `wait:account:*`。 -- 禁止用索引里的 score 直接判断并发数。 - -## 第二部分:并发槽清理和启动清理 - -修改文件:`backend/internal/repository/concurrency_cache.go`。 - -### 2.1 重写 CleanupExpiredAccountSlotKeys - -当前实现调用 `cleanupExpiredSlotKeysByPattern(ctx, accountSlotKeyPrefix+"*")`,必须删除。 - -新行为: - -1. 获取 Redis 当前秒 `now`。 -2. 从 `accountActiveIndexKey` 读取过期候选: - -```go -ids, err := c.rdb.ZRangeByScore(ctx, accountActiveIndexKey, &redis.ZRangeBy{ - Min: "-inf", - Max: strconv.FormatInt(now, 10), - Count: activeIndexCleanupBatchSize, -}).Result() -``` - -3. 对每个候选账号清理该账号 slot 过期成员并读真实状态。 -4. 如果真实 `slotCount == 0 && waitCount <= 0`,从 `accountActiveIndexKey` 删除该账号。 -5. 如果真实仍活跃,刷新 `accountActiveIndexKey` score。 -6. 不需要处理不在索引中的账号;其 slot key 自身有 `EXPIRE`,并且 acquire/get-load 会惰性清理过期成员。 - -这个函数不再表示“遍历所有账号槽位 key”,而是“处理索引中到期的账号候选”。保留原函数名是为了少改接口。 - -### 2.2 重写 CleanupStaleProcessSlots - -当前实现会扫描: - -- `concurrency:account:*` -- `concurrency:user:*` -- `wait:account:*` -- `concurrency:wait:*` - -必须去掉这些扫描。 - -新行为必须基于索引: - -1. 从 `accountActiveIndexKey` 读取所有未过期候选账号。 -2. 对每个账号: - - 对 `accountSlotKey(id)` 运行“单 key 清理脚本”,删除 requestID 前缀不是当前 `activeRequestPrefix` 的成员。 - - 删除 `accountWaitKey(id)`,因为等待者属于旧进程,重启后不能继续等待。 - - 调用 `refreshAccountActiveIndex(ctx, id)`。 -3. 从 `userActiveIndexKey` 读取所有未过期候选用户。 -4. 对每个用户: - - 对 `userSlotKey(id)` 运行同一个“单 key 清理脚本”。 - - 删除 `waitQueueKey(id)`。 - - 调用 `refreshUserActiveIndex(ctx, id)`。 - -新增单 key Lua 脚本,替代当前 `startupCleanupScript` 的多 key 版本: - -```lua -local key = KEYS[1] -local activePrefix = ARGV[1] -local slotTTL = tonumber(ARGV[2]) -local removed = 0 -local members = redis.call('ZRANGE', key, 0, -1) -for _, member in ipairs(members) do - if string.sub(member, 1, string.len(activePrefix)) ~= activePrefix then - removed = removed + redis.call('ZREM', key, member) - end -end -if redis.call('ZCARD', key) == 0 then - redis.call('DEL', key) -else - redis.call('EXPIRE', key, slotTTL) -end -return removed -``` - -该脚本只接受一个 slot key,避免 Redis Cluster CROSSSLOT。 - -如果索引不存在或为空: - -- `CleanupStaleProcessSlots` 直接返回 nil。 -- 不要 fallback 到 `SCAN`。 -- 旧版本遗留 key 依赖 Redis TTL 自然过期。不要在 app 启动时做兼容性 keyspace backfill。 - -### 2.3 删除旧扫描函数 - -删除以下生产函数: - -- `scanAccountIDsByPrefix` -- `parseAccountIDFromPrefixedKey`,如果没有其他生产调用 -- `cleanupExpiredSlotKeysByPattern` -- `cleanupSlotsByPattern` -- `deleteKeysByPattern` - -如果测试需要解析 key,测试内自建 helper,不要保留生产 helper。 - -## 第三部分:UMQ 锁索引 - -修改文件: - -- `backend/internal/repository/user_msg_queue_cache.go` -- `backend/internal/service/user_msg_queue_service.go` - -### 3.1 修改 service 接口 - -在 `backend/internal/service/user_msg_queue_service.go` 的 `UserMsgQueueCache` 接口中删除: - -```go -ScanLockKeys(ctx context.Context, maxCount int) ([]int64, error) -ForceReleaseLock(ctx context.Context, accountID int64) error -``` - -替换为: - -```go -ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (cleaned int, err error) -``` - -原因:cleanup worker 不应该知道 Redis lock key 的枚举方式,也不应该先枚举再逐个 `ForceReleaseLock`。候选读取、PTTL 校验、索引刷新应该封装在 cache 层。 - -### 3.2 修改 acquireLockScript 返回值 - -当前脚本只返回 0/1。改成返回数组: - -```lua -redis.replicate_commands() -local cur = redis.call('GET', KEYS[1]) -local ttl = tonumber(ARGV[2]) -if cur == ARGV[1] then - redis.call('PEXPIRE', KEYS[1], ttl) - local t = redis.call('TIME') - local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000) - return {1, ms + ttl} -end -if cur ~= false then - return {0, 0} -end -redis.call('SET', KEYS[1], ARGV[1], 'PX', ttl) -local t = redis.call('TIME') -local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000) -return {1, ms + ttl} -``` - -Go 侧解析: - -- 第一个元素是 acquired,1 表示拿到锁。 -- 第二个元素是 Redis 时间计算出的 `expireAtUnixMs`。 -- acquired 为 1 时,best-effort 写 `ZADD umqLockIndexKey expireAtMs accountID`。 -- `ZADD` 失败不能让 `AcquireLock` 返回失败。 - -### 3.3 修改 ReleaseLock - -`ReleaseLock` 主逻辑保持原子释放锁和写 last key。 - -释放成功时: - -```go -if result == 1 { - _ = c.rdb.ZRem(ctx, umqLockIndexKey, strconv.FormatInt(accountID, 10)).Err() -} -``` - -释放失败时不要删除索引。失败可能是 requestID 不匹配或 lock 已过期;cleanup worker 会处理 stale index。 - -### 3.4 新增 reconcile 脚本 - -删除 `forceReleaseLockScript` 的外部使用。新增脚本: - -```lua -local pttl = redis.call('PTTL', KEYS[1]) -if pttl == -2 then - return {-2, 0} -end -if pttl == -1 then - redis.call('DEL', KEYS[1]) - return {-1, 0} -end -return {1, pttl} -``` - -返回语义: - -- `-2`:lock key 不存在。Go 侧 `ZREM umqLockIndexKey accountID`。 -- `-1`:lock key 存在但无 TTL,脚本已删除。Go 侧 `ZREM umqLockIndexKey accountID`,cleaned++。 -- `1`:lock key 仍有 TTL。Go 侧用 Redis 当前毫秒 + pttl 刷新 `umqLockIndexKey` score。 - -### 3.5 实现 ReconcileExpiredLockCandidates - -实现步骤: - -1. 用 `c.rdb.Time(ctx)` 获取 Redis 当前毫秒 `nowMs`。 -2. 从 `umqLockIndexKey` 取到期候选: - -```go -members, err := c.rdb.ZRangeByScore(ctx, umqLockIndexKey, &redis.ZRangeBy{ - Min: "-inf", - Max: strconv.FormatInt(nowMs, 10), - Count: int64(maxCount), -}).Result() -``` - -3. 逐个解析 accountID。非法 member 直接 `ZREM`。 -4. 对合法 accountID 运行 reconcile 脚本,key 为 `umqLockKey(accountID)`。 -5. 根据返回值删除索引、刷新索引或累计 cleaned。 -6. 函数返回 cleaned 数。 - -禁止: - -- 禁止 fallback 到 `SCAN umq:{*}:lock`。 -- 禁止用 `KEYS umq:*`。 -- 禁止 cleanup worker 自己解析 lock key。 - -### 3.6 修改 StartCleanupWorker - -当前 worker 先 `ScanLockKeys` 再逐个 `ForceReleaseLock`。改成: - -```go -cleaned, err := s.cache.ReconcileExpiredLockCandidates(ctx, 1000) -if err != nil { - logger.LegacyPrintf("service.umq", "Cleanup reconcile failed: %v", err) - return -} -if cleaned > 0 { - logger.LegacyPrintf("service.umq", "Cleanup completed: released %d orphaned locks", cleaned) -} -``` - -worker 不再知道扫描、PTTL、索引等细节。 - -### 3.7 删除旧 UMQ 扫描函数 - -删除: - -- `umqScanPattern` -- `ScanLockKeys` -- `ForceReleaseLock`,如果无生产调用 - -如果测试仍需要强造 PTTL == -1 的 key,只在测试里直接写 Redis。 - -## 测试要求 - -### 并发缓存测试 - -新增或修改 `backend/internal/repository/concurrency_cache_*_test.go`。 - -必须覆盖: - -1. `AcquireAccountSlot` 成功后 `GetActiveAccountLoadMap` 能看到该账号。 -2. `ReleaseAccountSlot` 后 `GetActiveAccountLoadMap` 不再返回该账号。 -3. `IncrementAccountWaitCount` 成功后 `GetActiveAccountLoadMap` 能看到 waiting count。 -4. `DecrementAccountWaitCount` 后如果无 slot,则索引被移除。 -5. `CleanupExpiredAccountSlotKeys` 不依赖 keyspace scan:测试里只创建索引成员和对应 slot key,然后确认会清理;再创建未索引 slot key,确认不会被该函数主动发现。 -6. `CleanupStaleProcessSlots` 只处理索引中的 account/user,删除旧 request prefix 成员,保留当前 prefix 成员,删除 account/user wait key。 -7. 索引中存在非法 member 时,`GetActiveAccountLoadMap` 不报错,并移除非法 member。 - -### UMQ 测试 - -新增或修改 `backend/internal/repository/user_msg_queue_cache*_test.go` 和 `backend/internal/service/user_msg_queue_service*_test.go`。 - -必须覆盖: - -1. `AcquireLock` 成功后写入 `umq:lock:index`,score 大于 Redis 当前毫秒。 -2. `ReleaseLock` 成功后删除 `umq:lock:index` member。 -3. lock 已自然过期时,`ReconcileExpiredLockCandidates` 删除 stale index member。 -4. lock 仍有 TTL 但 index score 到期时,`ReconcileExpiredLockCandidates` 刷新 index score,不删除 lock。 -5. lock 存在且 `PTTL == -1` 时,`ReconcileExpiredLockCandidates` 删除 lock,删除 index member,并返回 cleaned=1。 -6. index 中非法 member 不导致错误,并被删除。 -7. `StartCleanupWorker` 调用 `ReconcileExpiredLockCandidates`,不再调用 `ScanLockKeys` 或 `ForceReleaseLock`。 - -### 禁止项测试 - -实现完成后运行: - -```powershell -rg -n "\.Scan\(" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go -rg -n "ScanLockKeys|umqScanPattern|scanAccountIDsByPrefix|cleanupExpiredSlotKeysByPattern|cleanupSlotsByPattern|deleteKeysByPattern" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go -``` - -上述命令必须无输出。 - -再运行相关测试。按项目约定,编译很慢时先把代码复制到 WSL 文件系统再跑: - -```bash -cd backend -go test ./internal/repository ./internal/service -``` - -如果全量包太慢,至少先跑: - -```bash -cd backend -go test ./internal/repository -run 'Concurrency|UserMsgQueue|Redis' -go test ./internal/service -run 'Concurrency|UserMessageQueue' -``` - -## 迁移和兼容 - -不要在应用启动时扫描旧 key 回填索引。 - -原因: - -- 这会把问题从运行期 `SCAN` 搬到启动期 `SCAN`。 -- 生产实例重启时 Redis 已经高 CPU,启动扫描会放大抖动。 -- 并发 slot key 和 wait key 都有 TTL,新版本写路径会为新流量维护索引,旧 key 可自然过期。 - -兼容策略: - -- 新版本上线后,新请求会逐步填充 `concurrency:*:active_index` 和 `umq:lock:index`。 -- 旧并发 slot key 没有索引时,不影响并发限制本身;对应账号下一次 acquire/get-load 会清理自己的 slot。 -- 旧 UMQ lock 如果有 TTL,会自然过期。 -- 极少数历史 `PTTL == -1` UMQ lock 且没有 index 的情况,不由应用自动发现。需要人工离线维护时,单独写一次性脚本,维护窗口运行,不要放进服务启动或后台 worker。 - -## 代码审查检查表 - -提交前逐项确认: - -- [ ] 没有新增 `KEYS`。 -- [ ] 没有新增生产路径 `SCAN`。 -- [ ] 没有在 Lua 脚本中同时操作全局索引 key 和账号/用户局部 key。 -- [ ] 索引维护失败不会让 acquire/release/decrement 的主结果失败。 -- [ ] `GetActiveAccountLoadMap` 对 stale index、非法 member、Redis nil 都能正常返回。 -- [ ] `CleanupExpiredAccountSlotKeys` 不再遍历 keyspace。 -- [ ] `CleanupStaleProcessSlots` 不再遍历 keyspace。 -- [ ] UMQ cleanup worker 不再知道 lock key pattern。 -- [ ] 所有旧扫描 helper 已删除或仅存在于测试文件。 -- [ ] 新测试覆盖成功路径、stale index、非法 member、PTTL -1、自然过期。 diff --git a/assets/partners/logos/cctk.jpg b/assets/partners/logos/cctk.jpg new file mode 100644 index 0000000000..3e50a07ae4 Binary files /dev/null and b/assets/partners/logos/cctk.jpg differ diff --git a/assets/partners/logos/fastaitoken.jpg b/assets/partners/logos/fastaitoken.jpg new file mode 100644 index 0000000000..afcf2acf51 Binary files /dev/null and b/assets/partners/logos/fastaitoken.jpg differ diff --git a/backend/Dockerfile b/backend/Dockerfile index d4adb5564b..9976abe46b 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.4-alpine +FROM golang:1.26.5-alpine WORKDIR /app diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index a0e8ec1d4e..010fbfb884 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -0.1.145 +0.1.149 diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index b9a9a3e80e..496473bc88 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -85,6 +85,8 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + batchImageCleanup *service.BatchImageCleanupService, + batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, emailQueue *service.EmailQueueService, billingCache *service.BillingCacheService, @@ -167,6 +169,18 @@ func provideCleanup( } return nil }}, + {"BatchImageCleanupService", func() error { + if batchImageCleanup != nil { + batchImageCleanup.Stop() + } + return nil + }}, + {"BatchImageWorkerRuntime", func() error { + if batchImageWorker != nil { + batchImageWorker.Stop() + } + return nil + }}, {"TokenRefreshService", func() error { tokenRefresh.Stop() return nil diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index a412563a6c..e148c5c363 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -96,6 +96,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { usageLogRepository := repository.NewUsageLogRepository(client, db) usageService := service.NewUsageService(usageLogRepository, userRepository, client, apiKeyAuthCacheInvalidator) opsRepository := repository.NewOpsRepository(db) + batchImageRepository := repository.NewBatchImageRepository(db) + batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig) + batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig) usageBillingRepository := repository.NewUsageBillingRepository(client, db) gatewayCache := repository.NewGatewayCache(redisClient) schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db) @@ -134,6 +137,11 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { channelRepository := repository.NewChannelRepository(db) channelService := service.NewChannelService(channelRepository, groupRepository, apiKeyAuthCacheInvalidator, pricingService) modelPricingResolver := service.NewModelPricingResolver(channelService, billingService) + batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver) + batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig) + batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig) + batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig) + batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig) notificationEmailService := service.NewNotificationEmailService(settingRepository, emailService) balanceNotifyService := service.ProvideBalanceNotifyService(emailService, settingRepository, accountRepository, notificationEmailService) gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, rpmCache, digestSessionStore, settingService, tlsFingerprintProfileService, channelService, modelPricingResolver, balanceNotifyService, serviceUserPlatformQuotaRepository) @@ -259,9 +267,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService, channelService) paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry) availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService) + batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService) idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig) idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig) - handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, idempotencyCoordinator, idempotencyCleanupService) + handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService) jwtAuthMiddleware := middleware.NewJWTAuthMiddleware(authService, userService) adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService) apiKeyAuthMiddleware := middleware.NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig) @@ -280,7 +289,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db) channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService) userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher) application := &Application{ Server: httpServer, Cleanup: v, @@ -322,6 +331,8 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + batchImageCleanup *service.BatchImageCleanupService, + batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, emailQueue *service.EmailQueueService, billingCache *service.BillingCacheService, @@ -403,6 +414,18 @@ func provideCleanup( } return nil }}, + {"BatchImageCleanupService", func() error { + if batchImageCleanup != nil { + batchImageCleanup.Stop() + } + return nil + }}, + {"BatchImageWorkerRuntime", func() error { + if batchImageWorker != nil { + batchImageWorker.Stop() + } + return nil + }}, {"TokenRefreshService", func() error { tokenRefresh.Stop() return nil diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index ef74cb4a2d..27707bc8c6 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -65,6 +65,8 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { subscriptionExpirySvc, &service.UsageCleanupService{}, idempotencyCleanupSvc, + &service.BatchImageCleanupService{}, + nil, // batchImageWorker pricingSvc, emailQueueSvc, billingCacheSvc, diff --git a/backend/ent/batchimageevent.go b/backend/ent/batchimageevent.go new file mode 100644 index 0000000000..3f95616e81 --- /dev/null +++ b/backend/ent/batchimageevent.go @@ -0,0 +1,158 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "encoding/json" + "fmt" + "strings" + "time" + + "entgo.io/ent" + "entgo.io/ent/dialect/sql" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" +) + +// BatchImageEvent is the model entity for the BatchImageEvent schema. +type BatchImageEvent struct { + config `json:"-"` + // ID of the ent. + ID int64 `json:"id,omitempty"` + // JobID holds the value of the "job_id" field. + JobID string `json:"job_id,omitempty"` + // EventType holds the value of the "event_type" field. + EventType string `json:"event_type,omitempty"` + // Payload holds the value of the "payload" field. + Payload map[string]interface{} `json:"payload,omitempty"` + // EventHash holds the value of the "event_hash" field. + EventHash *string `json:"event_hash,omitempty"` + // CreatedAt holds the value of the "created_at" field. + CreatedAt time.Time `json:"created_at,omitempty"` + selectValues sql.SelectValues +} + +// scanValues returns the types for scanning values from sql.Rows. +func (*BatchImageEvent) scanValues(columns []string) ([]any, error) { + values := make([]any, len(columns)) + for i := range columns { + switch columns[i] { + case batchimageevent.FieldPayload: + values[i] = new([]byte) + case batchimageevent.FieldID: + values[i] = new(sql.NullInt64) + case batchimageevent.FieldJobID, batchimageevent.FieldEventType, batchimageevent.FieldEventHash: + values[i] = new(sql.NullString) + case batchimageevent.FieldCreatedAt: + values[i] = new(sql.NullTime) + default: + values[i] = new(sql.UnknownType) + } + } + return values, nil +} + +// assignValues assigns the values that were returned from sql.Rows (after scanning) +// to the BatchImageEvent fields. +func (_m *BatchImageEvent) assignValues(columns []string, values []any) error { + if m, n := len(values), len(columns); m < n { + return fmt.Errorf("mismatch number of scan values: %d != %d", m, n) + } + for i := range columns { + switch columns[i] { + case batchimageevent.FieldID: + value, ok := values[i].(*sql.NullInt64) + if !ok { + return fmt.Errorf("unexpected type %T for field id", value) + } + _m.ID = int64(value.Int64) + case batchimageevent.FieldJobID: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field job_id", values[i]) + } else if value.Valid { + _m.JobID = value.String + } + case batchimageevent.FieldEventType: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field event_type", values[i]) + } else if value.Valid { + _m.EventType = value.String + } + case batchimageevent.FieldPayload: + if value, ok := values[i].(*[]byte); !ok { + return fmt.Errorf("unexpected type %T for field payload", values[i]) + } else if value != nil && len(*value) > 0 { + if err := json.Unmarshal(*value, &_m.Payload); err != nil { + return fmt.Errorf("unmarshal field payload: %w", err) + } + } + case batchimageevent.FieldEventHash: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field event_hash", values[i]) + } else if value.Valid { + _m.EventHash = new(string) + *_m.EventHash = value.String + } + case batchimageevent.FieldCreatedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field created_at", values[i]) + } else if value.Valid { + _m.CreatedAt = value.Time + } + default: + _m.selectValues.Set(columns[i], values[i]) + } + } + return nil +} + +// Value returns the ent.Value that was dynamically selected and assigned to the BatchImageEvent. +// This includes values selected through modifiers, order, etc. +func (_m *BatchImageEvent) Value(name string) (ent.Value, error) { + return _m.selectValues.Get(name) +} + +// Update returns a builder for updating this BatchImageEvent. +// Note that you need to call BatchImageEvent.Unwrap() before calling this method if this BatchImageEvent +// was returned from a transaction, and the transaction was committed or rolled back. +func (_m *BatchImageEvent) Update() *BatchImageEventUpdateOne { + return NewBatchImageEventClient(_m.config).UpdateOne(_m) +} + +// Unwrap unwraps the BatchImageEvent entity that was returned from a transaction after it was closed, +// so that all future queries will be executed through the driver which created the transaction. +func (_m *BatchImageEvent) Unwrap() *BatchImageEvent { + _tx, ok := _m.config.driver.(*txDriver) + if !ok { + panic("ent: BatchImageEvent is not a transactional entity") + } + _m.config.driver = _tx.drv + return _m +} + +// String implements the fmt.Stringer. +func (_m *BatchImageEvent) String() string { + var builder strings.Builder + builder.WriteString("BatchImageEvent(") + builder.WriteString(fmt.Sprintf("id=%v, ", _m.ID)) + builder.WriteString("job_id=") + builder.WriteString(_m.JobID) + builder.WriteString(", ") + builder.WriteString("event_type=") + builder.WriteString(_m.EventType) + builder.WriteString(", ") + builder.WriteString("payload=") + builder.WriteString(fmt.Sprintf("%v", _m.Payload)) + builder.WriteString(", ") + if v := _m.EventHash; v != nil { + builder.WriteString("event_hash=") + builder.WriteString(*v) + } + builder.WriteString(", ") + builder.WriteString("created_at=") + builder.WriteString(_m.CreatedAt.Format(time.ANSIC)) + builder.WriteByte(')') + return builder.String() +} + +// BatchImageEvents is a parsable slice of BatchImageEvent. +type BatchImageEvents []*BatchImageEvent diff --git a/backend/ent/batchimageevent/batchimageevent.go b/backend/ent/batchimageevent/batchimageevent.go new file mode 100644 index 0000000000..88b3dd8eca --- /dev/null +++ b/backend/ent/batchimageevent/batchimageevent.go @@ -0,0 +1,87 @@ +// Code generated by ent, DO NOT EDIT. + +package batchimageevent + +import ( + "time" + + "entgo.io/ent/dialect/sql" +) + +const ( + // Label holds the string label denoting the batchimageevent type in the database. + Label = "batch_image_event" + // FieldID holds the string denoting the id field in the database. + FieldID = "id" + // FieldJobID holds the string denoting the job_id field in the database. + FieldJobID = "job_id" + // FieldEventType holds the string denoting the event_type field in the database. + FieldEventType = "event_type" + // FieldPayload holds the string denoting the payload field in the database. + FieldPayload = "payload" + // FieldEventHash holds the string denoting the event_hash field in the database. + FieldEventHash = "event_hash" + // FieldCreatedAt holds the string denoting the created_at field in the database. + FieldCreatedAt = "created_at" + // Table holds the table name of the batchimageevent in the database. + Table = "batch_image_events" +) + +// Columns holds all SQL columns for batchimageevent fields. +var Columns = []string{ + FieldID, + FieldJobID, + FieldEventType, + FieldPayload, + FieldEventHash, + FieldCreatedAt, +} + +// ValidColumn reports if the column name is valid (part of the table columns). +func ValidColumn(column string) bool { + for i := range Columns { + if column == Columns[i] { + return true + } + } + return false +} + +var ( + // JobIDValidator is a validator for the "job_id" field. It is called by the builders before save. + JobIDValidator func(string) error + // EventTypeValidator is a validator for the "event_type" field. It is called by the builders before save. + EventTypeValidator func(string) error + // EventHashValidator is a validator for the "event_hash" field. It is called by the builders before save. + EventHashValidator func(string) error + // DefaultCreatedAt holds the default value on creation for the "created_at" field. + DefaultCreatedAt func() time.Time +) + +// OrderOption defines the ordering options for the BatchImageEvent queries. +type OrderOption func(*sql.Selector) + +// ByID orders the results by the id field. +func ByID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldID, opts...).ToFunc() +} + +// ByJobID orders the results by the job_id field. +func ByJobID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldJobID, opts...).ToFunc() +} + +// ByEventType orders the results by the event_type field. +func ByEventType(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldEventType, opts...).ToFunc() +} + +// ByEventHash orders the results by the event_hash field. +func ByEventHash(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldEventHash, opts...).ToFunc() +} + +// ByCreatedAt orders the results by the created_at field. +func ByCreatedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldCreatedAt, opts...).ToFunc() +} diff --git a/backend/ent/batchimageevent/where.go b/backend/ent/batchimageevent/where.go new file mode 100644 index 0000000000..3b5ef034f3 --- /dev/null +++ b/backend/ent/batchimageevent/where.go @@ -0,0 +1,345 @@ +// Code generated by ent, DO NOT EDIT. + +package batchimageevent + +import ( + "time" + + "entgo.io/ent/dialect/sql" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// ID filters vertices based on their ID field. +func ID(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldID, id)) +} + +// IDEQ applies the EQ predicate on the ID field. +func IDEQ(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldID, id)) +} + +// IDNEQ applies the NEQ predicate on the ID field. +func IDNEQ(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNEQ(FieldID, id)) +} + +// IDIn applies the In predicate on the ID field. +func IDIn(ids ...int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIn(FieldID, ids...)) +} + +// IDNotIn applies the NotIn predicate on the ID field. +func IDNotIn(ids ...int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotIn(FieldID, ids...)) +} + +// IDGT applies the GT predicate on the ID field. +func IDGT(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGT(FieldID, id)) +} + +// IDGTE applies the GTE predicate on the ID field. +func IDGTE(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGTE(FieldID, id)) +} + +// IDLT applies the LT predicate on the ID field. +func IDLT(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLT(FieldID, id)) +} + +// IDLTE applies the LTE predicate on the ID field. +func IDLTE(id int64) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLTE(FieldID, id)) +} + +// JobID applies equality check predicate on the "job_id" field. It's identical to JobIDEQ. +func JobID(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldJobID, v)) +} + +// EventType applies equality check predicate on the "event_type" field. It's identical to EventTypeEQ. +func EventType(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldEventType, v)) +} + +// EventHash applies equality check predicate on the "event_hash" field. It's identical to EventHashEQ. +func EventHash(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldEventHash, v)) +} + +// CreatedAt applies equality check predicate on the "created_at" field. It's identical to CreatedAtEQ. +func CreatedAt(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldCreatedAt, v)) +} + +// JobIDEQ applies the EQ predicate on the "job_id" field. +func JobIDEQ(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldJobID, v)) +} + +// JobIDNEQ applies the NEQ predicate on the "job_id" field. +func JobIDNEQ(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNEQ(FieldJobID, v)) +} + +// JobIDIn applies the In predicate on the "job_id" field. +func JobIDIn(vs ...string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIn(FieldJobID, vs...)) +} + +// JobIDNotIn applies the NotIn predicate on the "job_id" field. +func JobIDNotIn(vs ...string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotIn(FieldJobID, vs...)) +} + +// JobIDGT applies the GT predicate on the "job_id" field. +func JobIDGT(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGT(FieldJobID, v)) +} + +// JobIDGTE applies the GTE predicate on the "job_id" field. +func JobIDGTE(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGTE(FieldJobID, v)) +} + +// JobIDLT applies the LT predicate on the "job_id" field. +func JobIDLT(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLT(FieldJobID, v)) +} + +// JobIDLTE applies the LTE predicate on the "job_id" field. +func JobIDLTE(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLTE(FieldJobID, v)) +} + +// JobIDContains applies the Contains predicate on the "job_id" field. +func JobIDContains(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldContains(FieldJobID, v)) +} + +// JobIDHasPrefix applies the HasPrefix predicate on the "job_id" field. +func JobIDHasPrefix(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldHasPrefix(FieldJobID, v)) +} + +// JobIDHasSuffix applies the HasSuffix predicate on the "job_id" field. +func JobIDHasSuffix(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldHasSuffix(FieldJobID, v)) +} + +// JobIDEqualFold applies the EqualFold predicate on the "job_id" field. +func JobIDEqualFold(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEqualFold(FieldJobID, v)) +} + +// JobIDContainsFold applies the ContainsFold predicate on the "job_id" field. +func JobIDContainsFold(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldContainsFold(FieldJobID, v)) +} + +// EventTypeEQ applies the EQ predicate on the "event_type" field. +func EventTypeEQ(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldEventType, v)) +} + +// EventTypeNEQ applies the NEQ predicate on the "event_type" field. +func EventTypeNEQ(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNEQ(FieldEventType, v)) +} + +// EventTypeIn applies the In predicate on the "event_type" field. +func EventTypeIn(vs ...string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIn(FieldEventType, vs...)) +} + +// EventTypeNotIn applies the NotIn predicate on the "event_type" field. +func EventTypeNotIn(vs ...string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotIn(FieldEventType, vs...)) +} + +// EventTypeGT applies the GT predicate on the "event_type" field. +func EventTypeGT(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGT(FieldEventType, v)) +} + +// EventTypeGTE applies the GTE predicate on the "event_type" field. +func EventTypeGTE(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGTE(FieldEventType, v)) +} + +// EventTypeLT applies the LT predicate on the "event_type" field. +func EventTypeLT(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLT(FieldEventType, v)) +} + +// EventTypeLTE applies the LTE predicate on the "event_type" field. +func EventTypeLTE(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLTE(FieldEventType, v)) +} + +// EventTypeContains applies the Contains predicate on the "event_type" field. +func EventTypeContains(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldContains(FieldEventType, v)) +} + +// EventTypeHasPrefix applies the HasPrefix predicate on the "event_type" field. +func EventTypeHasPrefix(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldHasPrefix(FieldEventType, v)) +} + +// EventTypeHasSuffix applies the HasSuffix predicate on the "event_type" field. +func EventTypeHasSuffix(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldHasSuffix(FieldEventType, v)) +} + +// EventTypeEqualFold applies the EqualFold predicate on the "event_type" field. +func EventTypeEqualFold(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEqualFold(FieldEventType, v)) +} + +// EventTypeContainsFold applies the ContainsFold predicate on the "event_type" field. +func EventTypeContainsFold(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldContainsFold(FieldEventType, v)) +} + +// PayloadIsNil applies the IsNil predicate on the "payload" field. +func PayloadIsNil() predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIsNull(FieldPayload)) +} + +// PayloadNotNil applies the NotNil predicate on the "payload" field. +func PayloadNotNil() predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotNull(FieldPayload)) +} + +// EventHashEQ applies the EQ predicate on the "event_hash" field. +func EventHashEQ(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldEventHash, v)) +} + +// EventHashNEQ applies the NEQ predicate on the "event_hash" field. +func EventHashNEQ(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNEQ(FieldEventHash, v)) +} + +// EventHashIn applies the In predicate on the "event_hash" field. +func EventHashIn(vs ...string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIn(FieldEventHash, vs...)) +} + +// EventHashNotIn applies the NotIn predicate on the "event_hash" field. +func EventHashNotIn(vs ...string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotIn(FieldEventHash, vs...)) +} + +// EventHashGT applies the GT predicate on the "event_hash" field. +func EventHashGT(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGT(FieldEventHash, v)) +} + +// EventHashGTE applies the GTE predicate on the "event_hash" field. +func EventHashGTE(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGTE(FieldEventHash, v)) +} + +// EventHashLT applies the LT predicate on the "event_hash" field. +func EventHashLT(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLT(FieldEventHash, v)) +} + +// EventHashLTE applies the LTE predicate on the "event_hash" field. +func EventHashLTE(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLTE(FieldEventHash, v)) +} + +// EventHashContains applies the Contains predicate on the "event_hash" field. +func EventHashContains(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldContains(FieldEventHash, v)) +} + +// EventHashHasPrefix applies the HasPrefix predicate on the "event_hash" field. +func EventHashHasPrefix(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldHasPrefix(FieldEventHash, v)) +} + +// EventHashHasSuffix applies the HasSuffix predicate on the "event_hash" field. +func EventHashHasSuffix(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldHasSuffix(FieldEventHash, v)) +} + +// EventHashIsNil applies the IsNil predicate on the "event_hash" field. +func EventHashIsNil() predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIsNull(FieldEventHash)) +} + +// EventHashNotNil applies the NotNil predicate on the "event_hash" field. +func EventHashNotNil() predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotNull(FieldEventHash)) +} + +// EventHashEqualFold applies the EqualFold predicate on the "event_hash" field. +func EventHashEqualFold(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEqualFold(FieldEventHash, v)) +} + +// EventHashContainsFold applies the ContainsFold predicate on the "event_hash" field. +func EventHashContainsFold(v string) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldContainsFold(FieldEventHash, v)) +} + +// CreatedAtEQ applies the EQ predicate on the "created_at" field. +func CreatedAtEQ(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldEQ(FieldCreatedAt, v)) +} + +// CreatedAtNEQ applies the NEQ predicate on the "created_at" field. +func CreatedAtNEQ(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNEQ(FieldCreatedAt, v)) +} + +// CreatedAtIn applies the In predicate on the "created_at" field. +func CreatedAtIn(vs ...time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldIn(FieldCreatedAt, vs...)) +} + +// CreatedAtNotIn applies the NotIn predicate on the "created_at" field. +func CreatedAtNotIn(vs ...time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldNotIn(FieldCreatedAt, vs...)) +} + +// CreatedAtGT applies the GT predicate on the "created_at" field. +func CreatedAtGT(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGT(FieldCreatedAt, v)) +} + +// CreatedAtGTE applies the GTE predicate on the "created_at" field. +func CreatedAtGTE(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldGTE(FieldCreatedAt, v)) +} + +// CreatedAtLT applies the LT predicate on the "created_at" field. +func CreatedAtLT(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLT(FieldCreatedAt, v)) +} + +// CreatedAtLTE applies the LTE predicate on the "created_at" field. +func CreatedAtLTE(v time.Time) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.FieldLTE(FieldCreatedAt, v)) +} + +// And groups predicates with the AND operator between them. +func And(predicates ...predicate.BatchImageEvent) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.AndPredicates(predicates...)) +} + +// Or groups predicates with the OR operator between them. +func Or(predicates ...predicate.BatchImageEvent) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.OrPredicates(predicates...)) +} + +// Not applies the not operator on the given predicate. +func Not(p predicate.BatchImageEvent) predicate.BatchImageEvent { + return predicate.BatchImageEvent(sql.NotPredicates(p)) +} diff --git a/backend/ent/batchimageevent_create.go b/backend/ent/batchimageevent_create.go new file mode 100644 index 0000000000..c6ebef1dff --- /dev/null +++ b/backend/ent/batchimageevent_create.go @@ -0,0 +1,714 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "errors" + "fmt" + "time" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" +) + +// BatchImageEventCreate is the builder for creating a BatchImageEvent entity. +type BatchImageEventCreate struct { + config + mutation *BatchImageEventMutation + hooks []Hook + conflict []sql.ConflictOption +} + +// SetJobID sets the "job_id" field. +func (_c *BatchImageEventCreate) SetJobID(v string) *BatchImageEventCreate { + _c.mutation.SetJobID(v) + return _c +} + +// SetEventType sets the "event_type" field. +func (_c *BatchImageEventCreate) SetEventType(v string) *BatchImageEventCreate { + _c.mutation.SetEventType(v) + return _c +} + +// SetPayload sets the "payload" field. +func (_c *BatchImageEventCreate) SetPayload(v map[string]interface{}) *BatchImageEventCreate { + _c.mutation.SetPayload(v) + return _c +} + +// SetEventHash sets the "event_hash" field. +func (_c *BatchImageEventCreate) SetEventHash(v string) *BatchImageEventCreate { + _c.mutation.SetEventHash(v) + return _c +} + +// SetNillableEventHash sets the "event_hash" field if the given value is not nil. +func (_c *BatchImageEventCreate) SetNillableEventHash(v *string) *BatchImageEventCreate { + if v != nil { + _c.SetEventHash(*v) + } + return _c +} + +// SetCreatedAt sets the "created_at" field. +func (_c *BatchImageEventCreate) SetCreatedAt(v time.Time) *BatchImageEventCreate { + _c.mutation.SetCreatedAt(v) + return _c +} + +// SetNillableCreatedAt sets the "created_at" field if the given value is not nil. +func (_c *BatchImageEventCreate) SetNillableCreatedAt(v *time.Time) *BatchImageEventCreate { + if v != nil { + _c.SetCreatedAt(*v) + } + return _c +} + +// Mutation returns the BatchImageEventMutation object of the builder. +func (_c *BatchImageEventCreate) Mutation() *BatchImageEventMutation { + return _c.mutation +} + +// Save creates the BatchImageEvent in the database. +func (_c *BatchImageEventCreate) Save(ctx context.Context) (*BatchImageEvent, error) { + _c.defaults() + return withHooks(ctx, _c.sqlSave, _c.mutation, _c.hooks) +} + +// SaveX calls Save and panics if Save returns an error. +func (_c *BatchImageEventCreate) SaveX(ctx context.Context) *BatchImageEvent { + v, err := _c.Save(ctx) + if err != nil { + panic(err) + } + return v +} + +// Exec executes the query. +func (_c *BatchImageEventCreate) Exec(ctx context.Context) error { + _, err := _c.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_c *BatchImageEventCreate) ExecX(ctx context.Context) { + if err := _c.Exec(ctx); err != nil { + panic(err) + } +} + +// defaults sets the default values of the builder before save. +func (_c *BatchImageEventCreate) defaults() { + if _, ok := _c.mutation.CreatedAt(); !ok { + v := batchimageevent.DefaultCreatedAt() + _c.mutation.SetCreatedAt(v) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_c *BatchImageEventCreate) check() error { + if _, ok := _c.mutation.JobID(); !ok { + return &ValidationError{Name: "job_id", err: errors.New(`ent: missing required field "BatchImageEvent.job_id"`)} + } + if v, ok := _c.mutation.JobID(); ok { + if err := batchimageevent.JobIDValidator(v); err != nil { + return &ValidationError{Name: "job_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.job_id": %w`, err)} + } + } + if _, ok := _c.mutation.EventType(); !ok { + return &ValidationError{Name: "event_type", err: errors.New(`ent: missing required field "BatchImageEvent.event_type"`)} + } + if v, ok := _c.mutation.EventType(); ok { + if err := batchimageevent.EventTypeValidator(v); err != nil { + return &ValidationError{Name: "event_type", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.event_type": %w`, err)} + } + } + if v, ok := _c.mutation.EventHash(); ok { + if err := batchimageevent.EventHashValidator(v); err != nil { + return &ValidationError{Name: "event_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.event_hash": %w`, err)} + } + } + if _, ok := _c.mutation.CreatedAt(); !ok { + return &ValidationError{Name: "created_at", err: errors.New(`ent: missing required field "BatchImageEvent.created_at"`)} + } + return nil +} + +func (_c *BatchImageEventCreate) sqlSave(ctx context.Context) (*BatchImageEvent, error) { + if err := _c.check(); err != nil { + return nil, err + } + _node, _spec := _c.createSpec() + if err := sqlgraph.CreateNode(ctx, _c.driver, _spec); err != nil { + if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return nil, err + } + id := _spec.ID.Value.(int64) + _node.ID = int64(id) + _c.mutation.id = &_node.ID + _c.mutation.done = true + return _node, nil +} + +func (_c *BatchImageEventCreate) createSpec() (*BatchImageEvent, *sqlgraph.CreateSpec) { + var ( + _node = &BatchImageEvent{config: _c.config} + _spec = sqlgraph.NewCreateSpec(batchimageevent.Table, sqlgraph.NewFieldSpec(batchimageevent.FieldID, field.TypeInt64)) + ) + _spec.OnConflict = _c.conflict + if value, ok := _c.mutation.JobID(); ok { + _spec.SetField(batchimageevent.FieldJobID, field.TypeString, value) + _node.JobID = value + } + if value, ok := _c.mutation.EventType(); ok { + _spec.SetField(batchimageevent.FieldEventType, field.TypeString, value) + _node.EventType = value + } + if value, ok := _c.mutation.Payload(); ok { + _spec.SetField(batchimageevent.FieldPayload, field.TypeJSON, value) + _node.Payload = value + } + if value, ok := _c.mutation.EventHash(); ok { + _spec.SetField(batchimageevent.FieldEventHash, field.TypeString, value) + _node.EventHash = &value + } + if value, ok := _c.mutation.CreatedAt(); ok { + _spec.SetField(batchimageevent.FieldCreatedAt, field.TypeTime, value) + _node.CreatedAt = value + } + return _node, _spec +} + +// OnConflict allows configuring the `ON CONFLICT` / `ON DUPLICATE KEY` clause +// of the `INSERT` statement. For example: +// +// client.BatchImageEvent.Create(). +// SetJobID(v). +// OnConflict( +// // Update the row with the new values +// // the was proposed for insertion. +// sql.ResolveWithNewValues(), +// ). +// // Override some of the fields with custom +// // update values. +// Update(func(u *ent.BatchImageEventUpsert) { +// SetJobID(v+v). +// }). +// Exec(ctx) +func (_c *BatchImageEventCreate) OnConflict(opts ...sql.ConflictOption) *BatchImageEventUpsertOne { + _c.conflict = opts + return &BatchImageEventUpsertOne{ + create: _c, + } +} + +// OnConflictColumns calls `OnConflict` and configures the columns +// as conflict target. Using this option is equivalent to using: +// +// client.BatchImageEvent.Create(). +// OnConflict(sql.ConflictColumns(columns...)). +// Exec(ctx) +func (_c *BatchImageEventCreate) OnConflictColumns(columns ...string) *BatchImageEventUpsertOne { + _c.conflict = append(_c.conflict, sql.ConflictColumns(columns...)) + return &BatchImageEventUpsertOne{ + create: _c, + } +} + +type ( + // BatchImageEventUpsertOne is the builder for "upsert"-ing + // one BatchImageEvent node. + BatchImageEventUpsertOne struct { + create *BatchImageEventCreate + } + + // BatchImageEventUpsert is the "OnConflict" setter. + BatchImageEventUpsert struct { + *sql.UpdateSet + } +) + +// SetJobID sets the "job_id" field. +func (u *BatchImageEventUpsert) SetJobID(v string) *BatchImageEventUpsert { + u.Set(batchimageevent.FieldJobID, v) + return u +} + +// UpdateJobID sets the "job_id" field to the value that was provided on create. +func (u *BatchImageEventUpsert) UpdateJobID() *BatchImageEventUpsert { + u.SetExcluded(batchimageevent.FieldJobID) + return u +} + +// SetEventType sets the "event_type" field. +func (u *BatchImageEventUpsert) SetEventType(v string) *BatchImageEventUpsert { + u.Set(batchimageevent.FieldEventType, v) + return u +} + +// UpdateEventType sets the "event_type" field to the value that was provided on create. +func (u *BatchImageEventUpsert) UpdateEventType() *BatchImageEventUpsert { + u.SetExcluded(batchimageevent.FieldEventType) + return u +} + +// SetPayload sets the "payload" field. +func (u *BatchImageEventUpsert) SetPayload(v map[string]interface{}) *BatchImageEventUpsert { + u.Set(batchimageevent.FieldPayload, v) + return u +} + +// UpdatePayload sets the "payload" field to the value that was provided on create. +func (u *BatchImageEventUpsert) UpdatePayload() *BatchImageEventUpsert { + u.SetExcluded(batchimageevent.FieldPayload) + return u +} + +// ClearPayload clears the value of the "payload" field. +func (u *BatchImageEventUpsert) ClearPayload() *BatchImageEventUpsert { + u.SetNull(batchimageevent.FieldPayload) + return u +} + +// SetEventHash sets the "event_hash" field. +func (u *BatchImageEventUpsert) SetEventHash(v string) *BatchImageEventUpsert { + u.Set(batchimageevent.FieldEventHash, v) + return u +} + +// UpdateEventHash sets the "event_hash" field to the value that was provided on create. +func (u *BatchImageEventUpsert) UpdateEventHash() *BatchImageEventUpsert { + u.SetExcluded(batchimageevent.FieldEventHash) + return u +} + +// ClearEventHash clears the value of the "event_hash" field. +func (u *BatchImageEventUpsert) ClearEventHash() *BatchImageEventUpsert { + u.SetNull(batchimageevent.FieldEventHash) + return u +} + +// UpdateNewValues updates the mutable fields using the new values that were set on create. +// Using this option is equivalent to using: +// +// client.BatchImageEvent.Create(). +// OnConflict( +// sql.ResolveWithNewValues(), +// ). +// Exec(ctx) +func (u *BatchImageEventUpsertOne) UpdateNewValues() *BatchImageEventUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWithNewValues()) + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(s *sql.UpdateSet) { + if _, exists := u.create.mutation.CreatedAt(); exists { + s.SetIgnore(batchimageevent.FieldCreatedAt) + } + })) + return u +} + +// Ignore sets each column to itself in case of conflict. +// Using this option is equivalent to using: +// +// client.BatchImageEvent.Create(). +// OnConflict(sql.ResolveWithIgnore()). +// Exec(ctx) +func (u *BatchImageEventUpsertOne) Ignore() *BatchImageEventUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWithIgnore()) + return u +} + +// DoNothing configures the conflict_action to `DO NOTHING`. +// Supported only by SQLite and PostgreSQL. +func (u *BatchImageEventUpsertOne) DoNothing() *BatchImageEventUpsertOne { + u.create.conflict = append(u.create.conflict, sql.DoNothing()) + return u +} + +// Update allows overriding fields `UPDATE` values. See the BatchImageEventCreate.OnConflict +// documentation for more info. +func (u *BatchImageEventUpsertOne) Update(set func(*BatchImageEventUpsert)) *BatchImageEventUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(update *sql.UpdateSet) { + set(&BatchImageEventUpsert{UpdateSet: update}) + })) + return u +} + +// SetJobID sets the "job_id" field. +func (u *BatchImageEventUpsertOne) SetJobID(v string) *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetJobID(v) + }) +} + +// UpdateJobID sets the "job_id" field to the value that was provided on create. +func (u *BatchImageEventUpsertOne) UpdateJobID() *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdateJobID() + }) +} + +// SetEventType sets the "event_type" field. +func (u *BatchImageEventUpsertOne) SetEventType(v string) *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetEventType(v) + }) +} + +// UpdateEventType sets the "event_type" field to the value that was provided on create. +func (u *BatchImageEventUpsertOne) UpdateEventType() *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdateEventType() + }) +} + +// SetPayload sets the "payload" field. +func (u *BatchImageEventUpsertOne) SetPayload(v map[string]interface{}) *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetPayload(v) + }) +} + +// UpdatePayload sets the "payload" field to the value that was provided on create. +func (u *BatchImageEventUpsertOne) UpdatePayload() *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdatePayload() + }) +} + +// ClearPayload clears the value of the "payload" field. +func (u *BatchImageEventUpsertOne) ClearPayload() *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.ClearPayload() + }) +} + +// SetEventHash sets the "event_hash" field. +func (u *BatchImageEventUpsertOne) SetEventHash(v string) *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetEventHash(v) + }) +} + +// UpdateEventHash sets the "event_hash" field to the value that was provided on create. +func (u *BatchImageEventUpsertOne) UpdateEventHash() *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdateEventHash() + }) +} + +// ClearEventHash clears the value of the "event_hash" field. +func (u *BatchImageEventUpsertOne) ClearEventHash() *BatchImageEventUpsertOne { + return u.Update(func(s *BatchImageEventUpsert) { + s.ClearEventHash() + }) +} + +// Exec executes the query. +func (u *BatchImageEventUpsertOne) Exec(ctx context.Context) error { + if len(u.create.conflict) == 0 { + return errors.New("ent: missing options for BatchImageEventCreate.OnConflict") + } + return u.create.Exec(ctx) +} + +// ExecX is like Exec, but panics if an error occurs. +func (u *BatchImageEventUpsertOne) ExecX(ctx context.Context) { + if err := u.create.Exec(ctx); err != nil { + panic(err) + } +} + +// Exec executes the UPSERT query and returns the inserted/updated ID. +func (u *BatchImageEventUpsertOne) ID(ctx context.Context) (id int64, err error) { + node, err := u.create.Save(ctx) + if err != nil { + return id, err + } + return node.ID, nil +} + +// IDX is like ID, but panics if an error occurs. +func (u *BatchImageEventUpsertOne) IDX(ctx context.Context) int64 { + id, err := u.ID(ctx) + if err != nil { + panic(err) + } + return id +} + +// BatchImageEventCreateBulk is the builder for creating many BatchImageEvent entities in bulk. +type BatchImageEventCreateBulk struct { + config + err error + builders []*BatchImageEventCreate + conflict []sql.ConflictOption +} + +// Save creates the BatchImageEvent entities in the database. +func (_c *BatchImageEventCreateBulk) Save(ctx context.Context) ([]*BatchImageEvent, error) { + if _c.err != nil { + return nil, _c.err + } + specs := make([]*sqlgraph.CreateSpec, len(_c.builders)) + nodes := make([]*BatchImageEvent, len(_c.builders)) + mutators := make([]Mutator, len(_c.builders)) + for i := range _c.builders { + func(i int, root context.Context) { + builder := _c.builders[i] + builder.defaults() + var mut Mutator = MutateFunc(func(ctx context.Context, m Mutation) (Value, error) { + mutation, ok := m.(*BatchImageEventMutation) + if !ok { + return nil, fmt.Errorf("unexpected mutation type %T", m) + } + if err := builder.check(); err != nil { + return nil, err + } + builder.mutation = mutation + var err error + nodes[i], specs[i] = builder.createSpec() + if i < len(mutators)-1 { + _, err = mutators[i+1].Mutate(root, _c.builders[i+1].mutation) + } else { + spec := &sqlgraph.BatchCreateSpec{Nodes: specs} + spec.OnConflict = _c.conflict + // Invoke the actual operation on the latest mutation in the chain. + if err = sqlgraph.BatchCreate(ctx, _c.driver, spec); err != nil { + if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + } + } + if err != nil { + return nil, err + } + mutation.id = &nodes[i].ID + if specs[i].ID.Value != nil { + id := specs[i].ID.Value.(int64) + nodes[i].ID = int64(id) + } + mutation.done = true + return nodes[i], nil + }) + for i := len(builder.hooks) - 1; i >= 0; i-- { + mut = builder.hooks[i](mut) + } + mutators[i] = mut + }(i, ctx) + } + if len(mutators) > 0 { + if _, err := mutators[0].Mutate(ctx, _c.builders[0].mutation); err != nil { + return nil, err + } + } + return nodes, nil +} + +// SaveX is like Save, but panics if an error occurs. +func (_c *BatchImageEventCreateBulk) SaveX(ctx context.Context) []*BatchImageEvent { + v, err := _c.Save(ctx) + if err != nil { + panic(err) + } + return v +} + +// Exec executes the query. +func (_c *BatchImageEventCreateBulk) Exec(ctx context.Context) error { + _, err := _c.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_c *BatchImageEventCreateBulk) ExecX(ctx context.Context) { + if err := _c.Exec(ctx); err != nil { + panic(err) + } +} + +// OnConflict allows configuring the `ON CONFLICT` / `ON DUPLICATE KEY` clause +// of the `INSERT` statement. For example: +// +// client.BatchImageEvent.CreateBulk(builders...). +// OnConflict( +// // Update the row with the new values +// // the was proposed for insertion. +// sql.ResolveWithNewValues(), +// ). +// // Override some of the fields with custom +// // update values. +// Update(func(u *ent.BatchImageEventUpsert) { +// SetJobID(v+v). +// }). +// Exec(ctx) +func (_c *BatchImageEventCreateBulk) OnConflict(opts ...sql.ConflictOption) *BatchImageEventUpsertBulk { + _c.conflict = opts + return &BatchImageEventUpsertBulk{ + create: _c, + } +} + +// OnConflictColumns calls `OnConflict` and configures the columns +// as conflict target. Using this option is equivalent to using: +// +// client.BatchImageEvent.Create(). +// OnConflict(sql.ConflictColumns(columns...)). +// Exec(ctx) +func (_c *BatchImageEventCreateBulk) OnConflictColumns(columns ...string) *BatchImageEventUpsertBulk { + _c.conflict = append(_c.conflict, sql.ConflictColumns(columns...)) + return &BatchImageEventUpsertBulk{ + create: _c, + } +} + +// BatchImageEventUpsertBulk is the builder for "upsert"-ing +// a bulk of BatchImageEvent nodes. +type BatchImageEventUpsertBulk struct { + create *BatchImageEventCreateBulk +} + +// UpdateNewValues updates the mutable fields using the new values that +// were set on create. Using this option is equivalent to using: +// +// client.BatchImageEvent.Create(). +// OnConflict( +// sql.ResolveWithNewValues(), +// ). +// Exec(ctx) +func (u *BatchImageEventUpsertBulk) UpdateNewValues() *BatchImageEventUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWithNewValues()) + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(s *sql.UpdateSet) { + for _, b := range u.create.builders { + if _, exists := b.mutation.CreatedAt(); exists { + s.SetIgnore(batchimageevent.FieldCreatedAt) + } + } + })) + return u +} + +// Ignore sets each column to itself in case of conflict. +// Using this option is equivalent to using: +// +// client.BatchImageEvent.Create(). +// OnConflict(sql.ResolveWithIgnore()). +// Exec(ctx) +func (u *BatchImageEventUpsertBulk) Ignore() *BatchImageEventUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWithIgnore()) + return u +} + +// DoNothing configures the conflict_action to `DO NOTHING`. +// Supported only by SQLite and PostgreSQL. +func (u *BatchImageEventUpsertBulk) DoNothing() *BatchImageEventUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.DoNothing()) + return u +} + +// Update allows overriding fields `UPDATE` values. See the BatchImageEventCreateBulk.OnConflict +// documentation for more info. +func (u *BatchImageEventUpsertBulk) Update(set func(*BatchImageEventUpsert)) *BatchImageEventUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(update *sql.UpdateSet) { + set(&BatchImageEventUpsert{UpdateSet: update}) + })) + return u +} + +// SetJobID sets the "job_id" field. +func (u *BatchImageEventUpsertBulk) SetJobID(v string) *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetJobID(v) + }) +} + +// UpdateJobID sets the "job_id" field to the value that was provided on create. +func (u *BatchImageEventUpsertBulk) UpdateJobID() *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdateJobID() + }) +} + +// SetEventType sets the "event_type" field. +func (u *BatchImageEventUpsertBulk) SetEventType(v string) *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetEventType(v) + }) +} + +// UpdateEventType sets the "event_type" field to the value that was provided on create. +func (u *BatchImageEventUpsertBulk) UpdateEventType() *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdateEventType() + }) +} + +// SetPayload sets the "payload" field. +func (u *BatchImageEventUpsertBulk) SetPayload(v map[string]interface{}) *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetPayload(v) + }) +} + +// UpdatePayload sets the "payload" field to the value that was provided on create. +func (u *BatchImageEventUpsertBulk) UpdatePayload() *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdatePayload() + }) +} + +// ClearPayload clears the value of the "payload" field. +func (u *BatchImageEventUpsertBulk) ClearPayload() *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.ClearPayload() + }) +} + +// SetEventHash sets the "event_hash" field. +func (u *BatchImageEventUpsertBulk) SetEventHash(v string) *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.SetEventHash(v) + }) +} + +// UpdateEventHash sets the "event_hash" field to the value that was provided on create. +func (u *BatchImageEventUpsertBulk) UpdateEventHash() *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.UpdateEventHash() + }) +} + +// ClearEventHash clears the value of the "event_hash" field. +func (u *BatchImageEventUpsertBulk) ClearEventHash() *BatchImageEventUpsertBulk { + return u.Update(func(s *BatchImageEventUpsert) { + s.ClearEventHash() + }) +} + +// Exec executes the query. +func (u *BatchImageEventUpsertBulk) Exec(ctx context.Context) error { + if u.create.err != nil { + return u.create.err + } + for i, b := range u.create.builders { + if len(b.conflict) != 0 { + return fmt.Errorf("ent: OnConflict was set for builder %d. Set it on the BatchImageEventCreateBulk instead", i) + } + } + if len(u.create.conflict) == 0 { + return errors.New("ent: missing options for BatchImageEventCreateBulk.OnConflict") + } + return u.create.Exec(ctx) +} + +// ExecX is like Exec, but panics if an error occurs. +func (u *BatchImageEventUpsertBulk) ExecX(ctx context.Context) { + if err := u.create.Exec(ctx); err != nil { + panic(err) + } +} diff --git a/backend/ent/batchimageevent_delete.go b/backend/ent/batchimageevent_delete.go new file mode 100644 index 0000000000..54a51bef35 --- /dev/null +++ b/backend/ent/batchimageevent_delete.go @@ -0,0 +1,88 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageEventDelete is the builder for deleting a BatchImageEvent entity. +type BatchImageEventDelete struct { + config + hooks []Hook + mutation *BatchImageEventMutation +} + +// Where appends a list predicates to the BatchImageEventDelete builder. +func (_d *BatchImageEventDelete) Where(ps ...predicate.BatchImageEvent) *BatchImageEventDelete { + _d.mutation.Where(ps...) + return _d +} + +// Exec executes the deletion query and returns how many vertices were deleted. +func (_d *BatchImageEventDelete) Exec(ctx context.Context) (int, error) { + return withHooks(ctx, _d.sqlExec, _d.mutation, _d.hooks) +} + +// ExecX is like Exec, but panics if an error occurs. +func (_d *BatchImageEventDelete) ExecX(ctx context.Context) int { + n, err := _d.Exec(ctx) + if err != nil { + panic(err) + } + return n +} + +func (_d *BatchImageEventDelete) sqlExec(ctx context.Context) (int, error) { + _spec := sqlgraph.NewDeleteSpec(batchimageevent.Table, sqlgraph.NewFieldSpec(batchimageevent.FieldID, field.TypeInt64)) + if ps := _d.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + affected, err := sqlgraph.DeleteNodes(ctx, _d.driver, _spec) + if err != nil && sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + _d.mutation.done = true + return affected, err +} + +// BatchImageEventDeleteOne is the builder for deleting a single BatchImageEvent entity. +type BatchImageEventDeleteOne struct { + _d *BatchImageEventDelete +} + +// Where appends a list predicates to the BatchImageEventDelete builder. +func (_d *BatchImageEventDeleteOne) Where(ps ...predicate.BatchImageEvent) *BatchImageEventDeleteOne { + _d._d.mutation.Where(ps...) + return _d +} + +// Exec executes the deletion query. +func (_d *BatchImageEventDeleteOne) Exec(ctx context.Context) error { + n, err := _d._d.Exec(ctx) + switch { + case err != nil: + return err + case n == 0: + return &NotFoundError{batchimageevent.Label} + default: + return nil + } +} + +// ExecX is like Exec, but panics if an error occurs. +func (_d *BatchImageEventDeleteOne) ExecX(ctx context.Context) { + if err := _d.Exec(ctx); err != nil { + panic(err) + } +} diff --git a/backend/ent/batchimageevent_query.go b/backend/ent/batchimageevent_query.go new file mode 100644 index 0000000000..26fc5189a1 --- /dev/null +++ b/backend/ent/batchimageevent_query.go @@ -0,0 +1,564 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "fmt" + "math" + + "entgo.io/ent" + "entgo.io/ent/dialect" + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageEventQuery is the builder for querying BatchImageEvent entities. +type BatchImageEventQuery struct { + config + ctx *QueryContext + order []batchimageevent.OrderOption + inters []Interceptor + predicates []predicate.BatchImageEvent + modifiers []func(*sql.Selector) + // intermediate query (i.e. traversal path). + sql *sql.Selector + path func(context.Context) (*sql.Selector, error) +} + +// Where adds a new predicate for the BatchImageEventQuery builder. +func (_q *BatchImageEventQuery) Where(ps ...predicate.BatchImageEvent) *BatchImageEventQuery { + _q.predicates = append(_q.predicates, ps...) + return _q +} + +// Limit the number of records to be returned by this query. +func (_q *BatchImageEventQuery) Limit(limit int) *BatchImageEventQuery { + _q.ctx.Limit = &limit + return _q +} + +// Offset to start from. +func (_q *BatchImageEventQuery) Offset(offset int) *BatchImageEventQuery { + _q.ctx.Offset = &offset + return _q +} + +// Unique configures the query builder to filter duplicate records on query. +// By default, unique is set to true, and can be disabled using this method. +func (_q *BatchImageEventQuery) Unique(unique bool) *BatchImageEventQuery { + _q.ctx.Unique = &unique + return _q +} + +// Order specifies how the records should be ordered. +func (_q *BatchImageEventQuery) Order(o ...batchimageevent.OrderOption) *BatchImageEventQuery { + _q.order = append(_q.order, o...) + return _q +} + +// First returns the first BatchImageEvent entity from the query. +// Returns a *NotFoundError when no BatchImageEvent was found. +func (_q *BatchImageEventQuery) First(ctx context.Context) (*BatchImageEvent, error) { + nodes, err := _q.Limit(1).All(setContextOp(ctx, _q.ctx, ent.OpQueryFirst)) + if err != nil { + return nil, err + } + if len(nodes) == 0 { + return nil, &NotFoundError{batchimageevent.Label} + } + return nodes[0], nil +} + +// FirstX is like First, but panics if an error occurs. +func (_q *BatchImageEventQuery) FirstX(ctx context.Context) *BatchImageEvent { + node, err := _q.First(ctx) + if err != nil && !IsNotFound(err) { + panic(err) + } + return node +} + +// FirstID returns the first BatchImageEvent ID from the query. +// Returns a *NotFoundError when no BatchImageEvent ID was found. +func (_q *BatchImageEventQuery) FirstID(ctx context.Context) (id int64, err error) { + var ids []int64 + if ids, err = _q.Limit(1).IDs(setContextOp(ctx, _q.ctx, ent.OpQueryFirstID)); err != nil { + return + } + if len(ids) == 0 { + err = &NotFoundError{batchimageevent.Label} + return + } + return ids[0], nil +} + +// FirstIDX is like FirstID, but panics if an error occurs. +func (_q *BatchImageEventQuery) FirstIDX(ctx context.Context) int64 { + id, err := _q.FirstID(ctx) + if err != nil && !IsNotFound(err) { + panic(err) + } + return id +} + +// Only returns a single BatchImageEvent entity found by the query, ensuring it only returns one. +// Returns a *NotSingularError when more than one BatchImageEvent entity is found. +// Returns a *NotFoundError when no BatchImageEvent entities are found. +func (_q *BatchImageEventQuery) Only(ctx context.Context) (*BatchImageEvent, error) { + nodes, err := _q.Limit(2).All(setContextOp(ctx, _q.ctx, ent.OpQueryOnly)) + if err != nil { + return nil, err + } + switch len(nodes) { + case 1: + return nodes[0], nil + case 0: + return nil, &NotFoundError{batchimageevent.Label} + default: + return nil, &NotSingularError{batchimageevent.Label} + } +} + +// OnlyX is like Only, but panics if an error occurs. +func (_q *BatchImageEventQuery) OnlyX(ctx context.Context) *BatchImageEvent { + node, err := _q.Only(ctx) + if err != nil { + panic(err) + } + return node +} + +// OnlyID is like Only, but returns the only BatchImageEvent ID in the query. +// Returns a *NotSingularError when more than one BatchImageEvent ID is found. +// Returns a *NotFoundError when no entities are found. +func (_q *BatchImageEventQuery) OnlyID(ctx context.Context) (id int64, err error) { + var ids []int64 + if ids, err = _q.Limit(2).IDs(setContextOp(ctx, _q.ctx, ent.OpQueryOnlyID)); err != nil { + return + } + switch len(ids) { + case 1: + id = ids[0] + case 0: + err = &NotFoundError{batchimageevent.Label} + default: + err = &NotSingularError{batchimageevent.Label} + } + return +} + +// OnlyIDX is like OnlyID, but panics if an error occurs. +func (_q *BatchImageEventQuery) OnlyIDX(ctx context.Context) int64 { + id, err := _q.OnlyID(ctx) + if err != nil { + panic(err) + } + return id +} + +// All executes the query and returns a list of BatchImageEvents. +func (_q *BatchImageEventQuery) All(ctx context.Context) ([]*BatchImageEvent, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryAll) + if err := _q.prepareQuery(ctx); err != nil { + return nil, err + } + qr := querierAll[[]*BatchImageEvent, *BatchImageEventQuery]() + return withInterceptors[[]*BatchImageEvent](ctx, _q, qr, _q.inters) +} + +// AllX is like All, but panics if an error occurs. +func (_q *BatchImageEventQuery) AllX(ctx context.Context) []*BatchImageEvent { + nodes, err := _q.All(ctx) + if err != nil { + panic(err) + } + return nodes +} + +// IDs executes the query and returns a list of BatchImageEvent IDs. +func (_q *BatchImageEventQuery) IDs(ctx context.Context) (ids []int64, err error) { + if _q.ctx.Unique == nil && _q.path != nil { + _q.Unique(true) + } + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryIDs) + if err = _q.Select(batchimageevent.FieldID).Scan(ctx, &ids); err != nil { + return nil, err + } + return ids, nil +} + +// IDsX is like IDs, but panics if an error occurs. +func (_q *BatchImageEventQuery) IDsX(ctx context.Context) []int64 { + ids, err := _q.IDs(ctx) + if err != nil { + panic(err) + } + return ids +} + +// Count returns the count of the given query. +func (_q *BatchImageEventQuery) Count(ctx context.Context) (int, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryCount) + if err := _q.prepareQuery(ctx); err != nil { + return 0, err + } + return withInterceptors[int](ctx, _q, querierCount[*BatchImageEventQuery](), _q.inters) +} + +// CountX is like Count, but panics if an error occurs. +func (_q *BatchImageEventQuery) CountX(ctx context.Context) int { + count, err := _q.Count(ctx) + if err != nil { + panic(err) + } + return count +} + +// Exist returns true if the query has elements in the graph. +func (_q *BatchImageEventQuery) Exist(ctx context.Context) (bool, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryExist) + switch _, err := _q.FirstID(ctx); { + case IsNotFound(err): + return false, nil + case err != nil: + return false, fmt.Errorf("ent: check existence: %w", err) + default: + return true, nil + } +} + +// ExistX is like Exist, but panics if an error occurs. +func (_q *BatchImageEventQuery) ExistX(ctx context.Context) bool { + exist, err := _q.Exist(ctx) + if err != nil { + panic(err) + } + return exist +} + +// Clone returns a duplicate of the BatchImageEventQuery builder, including all associated steps. It can be +// used to prepare common query builders and use them differently after the clone is made. +func (_q *BatchImageEventQuery) Clone() *BatchImageEventQuery { + if _q == nil { + return nil + } + return &BatchImageEventQuery{ + config: _q.config, + ctx: _q.ctx.Clone(), + order: append([]batchimageevent.OrderOption{}, _q.order...), + inters: append([]Interceptor{}, _q.inters...), + predicates: append([]predicate.BatchImageEvent{}, _q.predicates...), + // clone intermediate query. + sql: _q.sql.Clone(), + path: _q.path, + } +} + +// GroupBy is used to group vertices by one or more fields/columns. +// It is often used with aggregate functions, like: count, max, mean, min, sum. +// +// Example: +// +// var v []struct { +// JobID string `json:"job_id,omitempty"` +// Count int `json:"count,omitempty"` +// } +// +// client.BatchImageEvent.Query(). +// GroupBy(batchimageevent.FieldJobID). +// Aggregate(ent.Count()). +// Scan(ctx, &v) +func (_q *BatchImageEventQuery) GroupBy(field string, fields ...string) *BatchImageEventGroupBy { + _q.ctx.Fields = append([]string{field}, fields...) + grbuild := &BatchImageEventGroupBy{build: _q} + grbuild.flds = &_q.ctx.Fields + grbuild.label = batchimageevent.Label + grbuild.scan = grbuild.Scan + return grbuild +} + +// Select allows the selection one or more fields/columns for the given query, +// instead of selecting all fields in the entity. +// +// Example: +// +// var v []struct { +// JobID string `json:"job_id,omitempty"` +// } +// +// client.BatchImageEvent.Query(). +// Select(batchimageevent.FieldJobID). +// Scan(ctx, &v) +func (_q *BatchImageEventQuery) Select(fields ...string) *BatchImageEventSelect { + _q.ctx.Fields = append(_q.ctx.Fields, fields...) + sbuild := &BatchImageEventSelect{BatchImageEventQuery: _q} + sbuild.label = batchimageevent.Label + sbuild.flds, sbuild.scan = &_q.ctx.Fields, sbuild.Scan + return sbuild +} + +// Aggregate returns a BatchImageEventSelect configured with the given aggregations. +func (_q *BatchImageEventQuery) Aggregate(fns ...AggregateFunc) *BatchImageEventSelect { + return _q.Select().Aggregate(fns...) +} + +func (_q *BatchImageEventQuery) prepareQuery(ctx context.Context) error { + for _, inter := range _q.inters { + if inter == nil { + return fmt.Errorf("ent: uninitialized interceptor (forgotten import ent/runtime?)") + } + if trv, ok := inter.(Traverser); ok { + if err := trv.Traverse(ctx, _q); err != nil { + return err + } + } + } + for _, f := range _q.ctx.Fields { + if !batchimageevent.ValidColumn(f) { + return &ValidationError{Name: f, err: fmt.Errorf("ent: invalid field %q for query", f)} + } + } + if _q.path != nil { + prev, err := _q.path(ctx) + if err != nil { + return err + } + _q.sql = prev + } + return nil +} + +func (_q *BatchImageEventQuery) sqlAll(ctx context.Context, hooks ...queryHook) ([]*BatchImageEvent, error) { + var ( + nodes = []*BatchImageEvent{} + _spec = _q.querySpec() + ) + _spec.ScanValues = func(columns []string) ([]any, error) { + return (*BatchImageEvent).scanValues(nil, columns) + } + _spec.Assign = func(columns []string, values []any) error { + node := &BatchImageEvent{config: _q.config} + nodes = append(nodes, node) + return node.assignValues(columns, values) + } + if len(_q.modifiers) > 0 { + _spec.Modifiers = _q.modifiers + } + for i := range hooks { + hooks[i](ctx, _spec) + } + if err := sqlgraph.QueryNodes(ctx, _q.driver, _spec); err != nil { + return nil, err + } + if len(nodes) == 0 { + return nodes, nil + } + return nodes, nil +} + +func (_q *BatchImageEventQuery) sqlCount(ctx context.Context) (int, error) { + _spec := _q.querySpec() + if len(_q.modifiers) > 0 { + _spec.Modifiers = _q.modifiers + } + _spec.Node.Columns = _q.ctx.Fields + if len(_q.ctx.Fields) > 0 { + _spec.Unique = _q.ctx.Unique != nil && *_q.ctx.Unique + } + return sqlgraph.CountNodes(ctx, _q.driver, _spec) +} + +func (_q *BatchImageEventQuery) querySpec() *sqlgraph.QuerySpec { + _spec := sqlgraph.NewQuerySpec(batchimageevent.Table, batchimageevent.Columns, sqlgraph.NewFieldSpec(batchimageevent.FieldID, field.TypeInt64)) + _spec.From = _q.sql + if unique := _q.ctx.Unique; unique != nil { + _spec.Unique = *unique + } else if _q.path != nil { + _spec.Unique = true + } + if fields := _q.ctx.Fields; len(fields) > 0 { + _spec.Node.Columns = make([]string, 0, len(fields)) + _spec.Node.Columns = append(_spec.Node.Columns, batchimageevent.FieldID) + for i := range fields { + if fields[i] != batchimageevent.FieldID { + _spec.Node.Columns = append(_spec.Node.Columns, fields[i]) + } + } + } + if ps := _q.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if limit := _q.ctx.Limit; limit != nil { + _spec.Limit = *limit + } + if offset := _q.ctx.Offset; offset != nil { + _spec.Offset = *offset + } + if ps := _q.order; len(ps) > 0 { + _spec.Order = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + return _spec +} + +func (_q *BatchImageEventQuery) sqlQuery(ctx context.Context) *sql.Selector { + builder := sql.Dialect(_q.driver.Dialect()) + t1 := builder.Table(batchimageevent.Table) + columns := _q.ctx.Fields + if len(columns) == 0 { + columns = batchimageevent.Columns + } + selector := builder.Select(t1.Columns(columns...)...).From(t1) + if _q.sql != nil { + selector = _q.sql + selector.Select(selector.Columns(columns...)...) + } + if _q.ctx.Unique != nil && *_q.ctx.Unique { + selector.Distinct() + } + for _, m := range _q.modifiers { + m(selector) + } + for _, p := range _q.predicates { + p(selector) + } + for _, p := range _q.order { + p(selector) + } + if offset := _q.ctx.Offset; offset != nil { + // limit is mandatory for offset clause. We start + // with default value, and override it below if needed. + selector.Offset(*offset).Limit(math.MaxInt32) + } + if limit := _q.ctx.Limit; limit != nil { + selector.Limit(*limit) + } + return selector +} + +// ForUpdate locks the selected rows against concurrent updates, and prevent them from being +// updated, deleted or "selected ... for update" by other sessions, until the transaction is +// either committed or rolled-back. +func (_q *BatchImageEventQuery) ForUpdate(opts ...sql.LockOption) *BatchImageEventQuery { + if _q.driver.Dialect() == dialect.Postgres { + _q.Unique(false) + } + _q.modifiers = append(_q.modifiers, func(s *sql.Selector) { + s.ForUpdate(opts...) + }) + return _q +} + +// ForShare behaves similarly to ForUpdate, except that it acquires a shared mode lock +// on any rows that are read. Other sessions can read the rows, but cannot modify them +// until your transaction commits. +func (_q *BatchImageEventQuery) ForShare(opts ...sql.LockOption) *BatchImageEventQuery { + if _q.driver.Dialect() == dialect.Postgres { + _q.Unique(false) + } + _q.modifiers = append(_q.modifiers, func(s *sql.Selector) { + s.ForShare(opts...) + }) + return _q +} + +// BatchImageEventGroupBy is the group-by builder for BatchImageEvent entities. +type BatchImageEventGroupBy struct { + selector + build *BatchImageEventQuery +} + +// Aggregate adds the given aggregation functions to the group-by query. +func (_g *BatchImageEventGroupBy) Aggregate(fns ...AggregateFunc) *BatchImageEventGroupBy { + _g.fns = append(_g.fns, fns...) + return _g +} + +// Scan applies the selector query and scans the result into the given value. +func (_g *BatchImageEventGroupBy) Scan(ctx context.Context, v any) error { + ctx = setContextOp(ctx, _g.build.ctx, ent.OpQueryGroupBy) + if err := _g.build.prepareQuery(ctx); err != nil { + return err + } + return scanWithInterceptors[*BatchImageEventQuery, *BatchImageEventGroupBy](ctx, _g.build, _g, _g.build.inters, v) +} + +func (_g *BatchImageEventGroupBy) sqlScan(ctx context.Context, root *BatchImageEventQuery, v any) error { + selector := root.sqlQuery(ctx).Select() + aggregation := make([]string, 0, len(_g.fns)) + for _, fn := range _g.fns { + aggregation = append(aggregation, fn(selector)) + } + if len(selector.SelectedColumns()) == 0 { + columns := make([]string, 0, len(*_g.flds)+len(_g.fns)) + for _, f := range *_g.flds { + columns = append(columns, selector.C(f)) + } + columns = append(columns, aggregation...) + selector.Select(columns...) + } + selector.GroupBy(selector.Columns(*_g.flds...)...) + if err := selector.Err(); err != nil { + return err + } + rows := &sql.Rows{} + query, args := selector.Query() + if err := _g.build.driver.Query(ctx, query, args, rows); err != nil { + return err + } + defer rows.Close() + return sql.ScanSlice(rows, v) +} + +// BatchImageEventSelect is the builder for selecting fields of BatchImageEvent entities. +type BatchImageEventSelect struct { + *BatchImageEventQuery + selector +} + +// Aggregate adds the given aggregation functions to the selector query. +func (_s *BatchImageEventSelect) Aggregate(fns ...AggregateFunc) *BatchImageEventSelect { + _s.fns = append(_s.fns, fns...) + return _s +} + +// Scan applies the selector query and scans the result into the given value. +func (_s *BatchImageEventSelect) Scan(ctx context.Context, v any) error { + ctx = setContextOp(ctx, _s.ctx, ent.OpQuerySelect) + if err := _s.prepareQuery(ctx); err != nil { + return err + } + return scanWithInterceptors[*BatchImageEventQuery, *BatchImageEventSelect](ctx, _s.BatchImageEventQuery, _s, _s.inters, v) +} + +func (_s *BatchImageEventSelect) sqlScan(ctx context.Context, root *BatchImageEventQuery, v any) error { + selector := root.sqlQuery(ctx) + aggregation := make([]string, 0, len(_s.fns)) + for _, fn := range _s.fns { + aggregation = append(aggregation, fn(selector)) + } + switch n := len(*_s.selector.flds); { + case n == 0 && len(aggregation) > 0: + selector.Select(aggregation...) + case n != 0 && len(aggregation) > 0: + selector.AppendSelect(aggregation...) + } + rows := &sql.Rows{} + query, args := selector.Query() + if err := _s.driver.Query(ctx, query, args, rows); err != nil { + return err + } + defer rows.Close() + return sql.ScanSlice(rows, v) +} diff --git a/backend/ent/batchimageevent_update.go b/backend/ent/batchimageevent_update.go new file mode 100644 index 0000000000..39035d6691 --- /dev/null +++ b/backend/ent/batchimageevent_update.go @@ -0,0 +1,377 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "errors" + "fmt" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageEventUpdate is the builder for updating BatchImageEvent entities. +type BatchImageEventUpdate struct { + config + hooks []Hook + mutation *BatchImageEventMutation +} + +// Where appends a list predicates to the BatchImageEventUpdate builder. +func (_u *BatchImageEventUpdate) Where(ps ...predicate.BatchImageEvent) *BatchImageEventUpdate { + _u.mutation.Where(ps...) + return _u +} + +// SetJobID sets the "job_id" field. +func (_u *BatchImageEventUpdate) SetJobID(v string) *BatchImageEventUpdate { + _u.mutation.SetJobID(v) + return _u +} + +// SetNillableJobID sets the "job_id" field if the given value is not nil. +func (_u *BatchImageEventUpdate) SetNillableJobID(v *string) *BatchImageEventUpdate { + if v != nil { + _u.SetJobID(*v) + } + return _u +} + +// SetEventType sets the "event_type" field. +func (_u *BatchImageEventUpdate) SetEventType(v string) *BatchImageEventUpdate { + _u.mutation.SetEventType(v) + return _u +} + +// SetNillableEventType sets the "event_type" field if the given value is not nil. +func (_u *BatchImageEventUpdate) SetNillableEventType(v *string) *BatchImageEventUpdate { + if v != nil { + _u.SetEventType(*v) + } + return _u +} + +// SetPayload sets the "payload" field. +func (_u *BatchImageEventUpdate) SetPayload(v map[string]interface{}) *BatchImageEventUpdate { + _u.mutation.SetPayload(v) + return _u +} + +// ClearPayload clears the value of the "payload" field. +func (_u *BatchImageEventUpdate) ClearPayload() *BatchImageEventUpdate { + _u.mutation.ClearPayload() + return _u +} + +// SetEventHash sets the "event_hash" field. +func (_u *BatchImageEventUpdate) SetEventHash(v string) *BatchImageEventUpdate { + _u.mutation.SetEventHash(v) + return _u +} + +// SetNillableEventHash sets the "event_hash" field if the given value is not nil. +func (_u *BatchImageEventUpdate) SetNillableEventHash(v *string) *BatchImageEventUpdate { + if v != nil { + _u.SetEventHash(*v) + } + return _u +} + +// ClearEventHash clears the value of the "event_hash" field. +func (_u *BatchImageEventUpdate) ClearEventHash() *BatchImageEventUpdate { + _u.mutation.ClearEventHash() + return _u +} + +// Mutation returns the BatchImageEventMutation object of the builder. +func (_u *BatchImageEventUpdate) Mutation() *BatchImageEventMutation { + return _u.mutation +} + +// Save executes the query and returns the number of nodes affected by the update operation. +func (_u *BatchImageEventUpdate) Save(ctx context.Context) (int, error) { + return withHooks(ctx, _u.sqlSave, _u.mutation, _u.hooks) +} + +// SaveX is like Save, but panics if an error occurs. +func (_u *BatchImageEventUpdate) SaveX(ctx context.Context) int { + affected, err := _u.Save(ctx) + if err != nil { + panic(err) + } + return affected +} + +// Exec executes the query. +func (_u *BatchImageEventUpdate) Exec(ctx context.Context) error { + _, err := _u.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_u *BatchImageEventUpdate) ExecX(ctx context.Context) { + if err := _u.Exec(ctx); err != nil { + panic(err) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_u *BatchImageEventUpdate) check() error { + if v, ok := _u.mutation.JobID(); ok { + if err := batchimageevent.JobIDValidator(v); err != nil { + return &ValidationError{Name: "job_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.job_id": %w`, err)} + } + } + if v, ok := _u.mutation.EventType(); ok { + if err := batchimageevent.EventTypeValidator(v); err != nil { + return &ValidationError{Name: "event_type", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.event_type": %w`, err)} + } + } + if v, ok := _u.mutation.EventHash(); ok { + if err := batchimageevent.EventHashValidator(v); err != nil { + return &ValidationError{Name: "event_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.event_hash": %w`, err)} + } + } + return nil +} + +func (_u *BatchImageEventUpdate) sqlSave(ctx context.Context) (_node int, err error) { + if err := _u.check(); err != nil { + return _node, err + } + _spec := sqlgraph.NewUpdateSpec(batchimageevent.Table, batchimageevent.Columns, sqlgraph.NewFieldSpec(batchimageevent.FieldID, field.TypeInt64)) + if ps := _u.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if value, ok := _u.mutation.JobID(); ok { + _spec.SetField(batchimageevent.FieldJobID, field.TypeString, value) + } + if value, ok := _u.mutation.EventType(); ok { + _spec.SetField(batchimageevent.FieldEventType, field.TypeString, value) + } + if value, ok := _u.mutation.Payload(); ok { + _spec.SetField(batchimageevent.FieldPayload, field.TypeJSON, value) + } + if _u.mutation.PayloadCleared() { + _spec.ClearField(batchimageevent.FieldPayload, field.TypeJSON) + } + if value, ok := _u.mutation.EventHash(); ok { + _spec.SetField(batchimageevent.FieldEventHash, field.TypeString, value) + } + if _u.mutation.EventHashCleared() { + _spec.ClearField(batchimageevent.FieldEventHash, field.TypeString) + } + if _node, err = sqlgraph.UpdateNodes(ctx, _u.driver, _spec); err != nil { + if _, ok := err.(*sqlgraph.NotFoundError); ok { + err = &NotFoundError{batchimageevent.Label} + } else if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return 0, err + } + _u.mutation.done = true + return _node, nil +} + +// BatchImageEventUpdateOne is the builder for updating a single BatchImageEvent entity. +type BatchImageEventUpdateOne struct { + config + fields []string + hooks []Hook + mutation *BatchImageEventMutation +} + +// SetJobID sets the "job_id" field. +func (_u *BatchImageEventUpdateOne) SetJobID(v string) *BatchImageEventUpdateOne { + _u.mutation.SetJobID(v) + return _u +} + +// SetNillableJobID sets the "job_id" field if the given value is not nil. +func (_u *BatchImageEventUpdateOne) SetNillableJobID(v *string) *BatchImageEventUpdateOne { + if v != nil { + _u.SetJobID(*v) + } + return _u +} + +// SetEventType sets the "event_type" field. +func (_u *BatchImageEventUpdateOne) SetEventType(v string) *BatchImageEventUpdateOne { + _u.mutation.SetEventType(v) + return _u +} + +// SetNillableEventType sets the "event_type" field if the given value is not nil. +func (_u *BatchImageEventUpdateOne) SetNillableEventType(v *string) *BatchImageEventUpdateOne { + if v != nil { + _u.SetEventType(*v) + } + return _u +} + +// SetPayload sets the "payload" field. +func (_u *BatchImageEventUpdateOne) SetPayload(v map[string]interface{}) *BatchImageEventUpdateOne { + _u.mutation.SetPayload(v) + return _u +} + +// ClearPayload clears the value of the "payload" field. +func (_u *BatchImageEventUpdateOne) ClearPayload() *BatchImageEventUpdateOne { + _u.mutation.ClearPayload() + return _u +} + +// SetEventHash sets the "event_hash" field. +func (_u *BatchImageEventUpdateOne) SetEventHash(v string) *BatchImageEventUpdateOne { + _u.mutation.SetEventHash(v) + return _u +} + +// SetNillableEventHash sets the "event_hash" field if the given value is not nil. +func (_u *BatchImageEventUpdateOne) SetNillableEventHash(v *string) *BatchImageEventUpdateOne { + if v != nil { + _u.SetEventHash(*v) + } + return _u +} + +// ClearEventHash clears the value of the "event_hash" field. +func (_u *BatchImageEventUpdateOne) ClearEventHash() *BatchImageEventUpdateOne { + _u.mutation.ClearEventHash() + return _u +} + +// Mutation returns the BatchImageEventMutation object of the builder. +func (_u *BatchImageEventUpdateOne) Mutation() *BatchImageEventMutation { + return _u.mutation +} + +// Where appends a list predicates to the BatchImageEventUpdate builder. +func (_u *BatchImageEventUpdateOne) Where(ps ...predicate.BatchImageEvent) *BatchImageEventUpdateOne { + _u.mutation.Where(ps...) + return _u +} + +// Select allows selecting one or more fields (columns) of the returned entity. +// The default is selecting all fields defined in the entity schema. +func (_u *BatchImageEventUpdateOne) Select(field string, fields ...string) *BatchImageEventUpdateOne { + _u.fields = append([]string{field}, fields...) + return _u +} + +// Save executes the query and returns the updated BatchImageEvent entity. +func (_u *BatchImageEventUpdateOne) Save(ctx context.Context) (*BatchImageEvent, error) { + return withHooks(ctx, _u.sqlSave, _u.mutation, _u.hooks) +} + +// SaveX is like Save, but panics if an error occurs. +func (_u *BatchImageEventUpdateOne) SaveX(ctx context.Context) *BatchImageEvent { + node, err := _u.Save(ctx) + if err != nil { + panic(err) + } + return node +} + +// Exec executes the query on the entity. +func (_u *BatchImageEventUpdateOne) Exec(ctx context.Context) error { + _, err := _u.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_u *BatchImageEventUpdateOne) ExecX(ctx context.Context) { + if err := _u.Exec(ctx); err != nil { + panic(err) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_u *BatchImageEventUpdateOne) check() error { + if v, ok := _u.mutation.JobID(); ok { + if err := batchimageevent.JobIDValidator(v); err != nil { + return &ValidationError{Name: "job_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.job_id": %w`, err)} + } + } + if v, ok := _u.mutation.EventType(); ok { + if err := batchimageevent.EventTypeValidator(v); err != nil { + return &ValidationError{Name: "event_type", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.event_type": %w`, err)} + } + } + if v, ok := _u.mutation.EventHash(); ok { + if err := batchimageevent.EventHashValidator(v); err != nil { + return &ValidationError{Name: "event_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageEvent.event_hash": %w`, err)} + } + } + return nil +} + +func (_u *BatchImageEventUpdateOne) sqlSave(ctx context.Context) (_node *BatchImageEvent, err error) { + if err := _u.check(); err != nil { + return _node, err + } + _spec := sqlgraph.NewUpdateSpec(batchimageevent.Table, batchimageevent.Columns, sqlgraph.NewFieldSpec(batchimageevent.FieldID, field.TypeInt64)) + id, ok := _u.mutation.ID() + if !ok { + return nil, &ValidationError{Name: "id", err: errors.New(`ent: missing "BatchImageEvent.id" for update`)} + } + _spec.Node.ID.Value = id + if fields := _u.fields; len(fields) > 0 { + _spec.Node.Columns = make([]string, 0, len(fields)) + _spec.Node.Columns = append(_spec.Node.Columns, batchimageevent.FieldID) + for _, f := range fields { + if !batchimageevent.ValidColumn(f) { + return nil, &ValidationError{Name: f, err: fmt.Errorf("ent: invalid field %q for query", f)} + } + if f != batchimageevent.FieldID { + _spec.Node.Columns = append(_spec.Node.Columns, f) + } + } + } + if ps := _u.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if value, ok := _u.mutation.JobID(); ok { + _spec.SetField(batchimageevent.FieldJobID, field.TypeString, value) + } + if value, ok := _u.mutation.EventType(); ok { + _spec.SetField(batchimageevent.FieldEventType, field.TypeString, value) + } + if value, ok := _u.mutation.Payload(); ok { + _spec.SetField(batchimageevent.FieldPayload, field.TypeJSON, value) + } + if _u.mutation.PayloadCleared() { + _spec.ClearField(batchimageevent.FieldPayload, field.TypeJSON) + } + if value, ok := _u.mutation.EventHash(); ok { + _spec.SetField(batchimageevent.FieldEventHash, field.TypeString, value) + } + if _u.mutation.EventHashCleared() { + _spec.ClearField(batchimageevent.FieldEventHash, field.TypeString) + } + _node = &BatchImageEvent{config: _u.config} + _spec.Assign = _node.assignValues + _spec.ScanValues = _node.scanValues + if err = sqlgraph.UpdateNode(ctx, _u.driver, _spec); err != nil { + if _, ok := err.(*sqlgraph.NotFoundError); ok { + err = &NotFoundError{batchimageevent.Label} + } else if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return nil, err + } + _u.mutation.done = true + return _node, nil +} diff --git a/backend/ent/batchimageitem.go b/backend/ent/batchimageitem.go new file mode 100644 index 0000000000..47b876f158 --- /dev/null +++ b/backend/ent/batchimageitem.go @@ -0,0 +1,320 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "fmt" + "strings" + "time" + + "entgo.io/ent" + "entgo.io/ent/dialect/sql" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" +) + +// BatchImageItem is the model entity for the BatchImageItem schema. +type BatchImageItem struct { + config `json:"-"` + // ID of the ent. + ID int64 `json:"id,omitempty"` + // JobID holds the value of the "job_id" field. + JobID string `json:"job_id,omitempty"` + // CustomID holds the value of the "custom_id" field. + CustomID string `json:"custom_id,omitempty"` + // Status holds the value of the "status" field. + Status string `json:"status,omitempty"` + // RequestHash holds the value of the "request_hash" field. + RequestHash *string `json:"request_hash,omitempty"` + // PromptPreview holds the value of the "prompt_preview" field. + PromptPreview *string `json:"prompt_preview,omitempty"` + // ProviderSourceObject holds the value of the "provider_source_object" field. + ProviderSourceObject *string `json:"provider_source_object,omitempty"` + // SourceLineNumber holds the value of the "source_line_number" field. + SourceLineNumber *int `json:"source_line_number,omitempty"` + // SourceByteOffset holds the value of the "source_byte_offset" field. + SourceByteOffset *int64 `json:"source_byte_offset,omitempty"` + // SourceByteLength holds the value of the "source_byte_length" field. + SourceByteLength *int64 `json:"source_byte_length,omitempty"` + // MimeType holds the value of the "mime_type" field. + MimeType *string `json:"mime_type,omitempty"` + // FileExtension holds the value of the "file_extension" field. + FileExtension *string `json:"file_extension,omitempty"` + // ImageCount holds the value of the "image_count" field. + ImageCount int `json:"image_count,omitempty"` + // ErrorCode holds the value of the "error_code" field. + ErrorCode *string `json:"error_code,omitempty"` + // ErrorMessage holds the value of the "error_message" field. + ErrorMessage *string `json:"error_message,omitempty"` + // BilledAmount holds the value of the "billed_amount" field. + BilledAmount *float64 `json:"billed_amount,omitempty"` + // CreatedAt holds the value of the "created_at" field. + CreatedAt time.Time `json:"created_at,omitempty"` + // IndexedAt holds the value of the "indexed_at" field. + IndexedAt *time.Time `json:"indexed_at,omitempty"` + selectValues sql.SelectValues +} + +// scanValues returns the types for scanning values from sql.Rows. +func (*BatchImageItem) scanValues(columns []string) ([]any, error) { + values := make([]any, len(columns)) + for i := range columns { + switch columns[i] { + case batchimageitem.FieldBilledAmount: + values[i] = new(sql.NullFloat64) + case batchimageitem.FieldID, batchimageitem.FieldSourceLineNumber, batchimageitem.FieldSourceByteOffset, batchimageitem.FieldSourceByteLength, batchimageitem.FieldImageCount: + values[i] = new(sql.NullInt64) + case batchimageitem.FieldJobID, batchimageitem.FieldCustomID, batchimageitem.FieldStatus, batchimageitem.FieldRequestHash, batchimageitem.FieldPromptPreview, batchimageitem.FieldProviderSourceObject, batchimageitem.FieldMimeType, batchimageitem.FieldFileExtension, batchimageitem.FieldErrorCode, batchimageitem.FieldErrorMessage: + values[i] = new(sql.NullString) + case batchimageitem.FieldCreatedAt, batchimageitem.FieldIndexedAt: + values[i] = new(sql.NullTime) + default: + values[i] = new(sql.UnknownType) + } + } + return values, nil +} + +// assignValues assigns the values that were returned from sql.Rows (after scanning) +// to the BatchImageItem fields. +func (_m *BatchImageItem) assignValues(columns []string, values []any) error { + if m, n := len(values), len(columns); m < n { + return fmt.Errorf("mismatch number of scan values: %d != %d", m, n) + } + for i := range columns { + switch columns[i] { + case batchimageitem.FieldID: + value, ok := values[i].(*sql.NullInt64) + if !ok { + return fmt.Errorf("unexpected type %T for field id", value) + } + _m.ID = int64(value.Int64) + case batchimageitem.FieldJobID: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field job_id", values[i]) + } else if value.Valid { + _m.JobID = value.String + } + case batchimageitem.FieldCustomID: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field custom_id", values[i]) + } else if value.Valid { + _m.CustomID = value.String + } + case batchimageitem.FieldStatus: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field status", values[i]) + } else if value.Valid { + _m.Status = value.String + } + case batchimageitem.FieldRequestHash: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field request_hash", values[i]) + } else if value.Valid { + _m.RequestHash = new(string) + *_m.RequestHash = value.String + } + case batchimageitem.FieldPromptPreview: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field prompt_preview", values[i]) + } else if value.Valid { + _m.PromptPreview = new(string) + *_m.PromptPreview = value.String + } + case batchimageitem.FieldProviderSourceObject: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field provider_source_object", values[i]) + } else if value.Valid { + _m.ProviderSourceObject = new(string) + *_m.ProviderSourceObject = value.String + } + case batchimageitem.FieldSourceLineNumber: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field source_line_number", values[i]) + } else if value.Valid { + _m.SourceLineNumber = new(int) + *_m.SourceLineNumber = int(value.Int64) + } + case batchimageitem.FieldSourceByteOffset: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field source_byte_offset", values[i]) + } else if value.Valid { + _m.SourceByteOffset = new(int64) + *_m.SourceByteOffset = value.Int64 + } + case batchimageitem.FieldSourceByteLength: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field source_byte_length", values[i]) + } else if value.Valid { + _m.SourceByteLength = new(int64) + *_m.SourceByteLength = value.Int64 + } + case batchimageitem.FieldMimeType: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field mime_type", values[i]) + } else if value.Valid { + _m.MimeType = new(string) + *_m.MimeType = value.String + } + case batchimageitem.FieldFileExtension: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field file_extension", values[i]) + } else if value.Valid { + _m.FileExtension = new(string) + *_m.FileExtension = value.String + } + case batchimageitem.FieldImageCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field image_count", values[i]) + } else if value.Valid { + _m.ImageCount = int(value.Int64) + } + case batchimageitem.FieldErrorCode: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field error_code", values[i]) + } else if value.Valid { + _m.ErrorCode = new(string) + *_m.ErrorCode = value.String + } + case batchimageitem.FieldErrorMessage: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field error_message", values[i]) + } else if value.Valid { + _m.ErrorMessage = new(string) + *_m.ErrorMessage = value.String + } + case batchimageitem.FieldBilledAmount: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field billed_amount", values[i]) + } else if value.Valid { + _m.BilledAmount = new(float64) + *_m.BilledAmount = value.Float64 + } + case batchimageitem.FieldCreatedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field created_at", values[i]) + } else if value.Valid { + _m.CreatedAt = value.Time + } + case batchimageitem.FieldIndexedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field indexed_at", values[i]) + } else if value.Valid { + _m.IndexedAt = new(time.Time) + *_m.IndexedAt = value.Time + } + default: + _m.selectValues.Set(columns[i], values[i]) + } + } + return nil +} + +// Value returns the ent.Value that was dynamically selected and assigned to the BatchImageItem. +// This includes values selected through modifiers, order, etc. +func (_m *BatchImageItem) Value(name string) (ent.Value, error) { + return _m.selectValues.Get(name) +} + +// Update returns a builder for updating this BatchImageItem. +// Note that you need to call BatchImageItem.Unwrap() before calling this method if this BatchImageItem +// was returned from a transaction, and the transaction was committed or rolled back. +func (_m *BatchImageItem) Update() *BatchImageItemUpdateOne { + return NewBatchImageItemClient(_m.config).UpdateOne(_m) +} + +// Unwrap unwraps the BatchImageItem entity that was returned from a transaction after it was closed, +// so that all future queries will be executed through the driver which created the transaction. +func (_m *BatchImageItem) Unwrap() *BatchImageItem { + _tx, ok := _m.config.driver.(*txDriver) + if !ok { + panic("ent: BatchImageItem is not a transactional entity") + } + _m.config.driver = _tx.drv + return _m +} + +// String implements the fmt.Stringer. +func (_m *BatchImageItem) String() string { + var builder strings.Builder + builder.WriteString("BatchImageItem(") + builder.WriteString(fmt.Sprintf("id=%v, ", _m.ID)) + builder.WriteString("job_id=") + builder.WriteString(_m.JobID) + builder.WriteString(", ") + builder.WriteString("custom_id=") + builder.WriteString(_m.CustomID) + builder.WriteString(", ") + builder.WriteString("status=") + builder.WriteString(_m.Status) + builder.WriteString(", ") + if v := _m.RequestHash; v != nil { + builder.WriteString("request_hash=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.PromptPreview; v != nil { + builder.WriteString("prompt_preview=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.ProviderSourceObject; v != nil { + builder.WriteString("provider_source_object=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.SourceLineNumber; v != nil { + builder.WriteString("source_line_number=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.SourceByteOffset; v != nil { + builder.WriteString("source_byte_offset=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.SourceByteLength; v != nil { + builder.WriteString("source_byte_length=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.MimeType; v != nil { + builder.WriteString("mime_type=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.FileExtension; v != nil { + builder.WriteString("file_extension=") + builder.WriteString(*v) + } + builder.WriteString(", ") + builder.WriteString("image_count=") + builder.WriteString(fmt.Sprintf("%v", _m.ImageCount)) + builder.WriteString(", ") + if v := _m.ErrorCode; v != nil { + builder.WriteString("error_code=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.ErrorMessage; v != nil { + builder.WriteString("error_message=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.BilledAmount; v != nil { + builder.WriteString("billed_amount=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + builder.WriteString("created_at=") + builder.WriteString(_m.CreatedAt.Format(time.ANSIC)) + builder.WriteString(", ") + if v := _m.IndexedAt; v != nil { + builder.WriteString("indexed_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteByte(')') + return builder.String() +} + +// BatchImageItems is a parsable slice of BatchImageItem. +type BatchImageItems []*BatchImageItem diff --git a/backend/ent/batchimageitem/batchimageitem.go b/backend/ent/batchimageitem/batchimageitem.go new file mode 100644 index 0000000000..3656e31dfe --- /dev/null +++ b/backend/ent/batchimageitem/batchimageitem.go @@ -0,0 +1,200 @@ +// Code generated by ent, DO NOT EDIT. + +package batchimageitem + +import ( + "time" + + "entgo.io/ent/dialect/sql" +) + +const ( + // Label holds the string label denoting the batchimageitem type in the database. + Label = "batch_image_item" + // FieldID holds the string denoting the id field in the database. + FieldID = "id" + // FieldJobID holds the string denoting the job_id field in the database. + FieldJobID = "job_id" + // FieldCustomID holds the string denoting the custom_id field in the database. + FieldCustomID = "custom_id" + // FieldStatus holds the string denoting the status field in the database. + FieldStatus = "status" + // FieldRequestHash holds the string denoting the request_hash field in the database. + FieldRequestHash = "request_hash" + // FieldPromptPreview holds the string denoting the prompt_preview field in the database. + FieldPromptPreview = "prompt_preview" + // FieldProviderSourceObject holds the string denoting the provider_source_object field in the database. + FieldProviderSourceObject = "provider_source_object" + // FieldSourceLineNumber holds the string denoting the source_line_number field in the database. + FieldSourceLineNumber = "source_line_number" + // FieldSourceByteOffset holds the string denoting the source_byte_offset field in the database. + FieldSourceByteOffset = "source_byte_offset" + // FieldSourceByteLength holds the string denoting the source_byte_length field in the database. + FieldSourceByteLength = "source_byte_length" + // FieldMimeType holds the string denoting the mime_type field in the database. + FieldMimeType = "mime_type" + // FieldFileExtension holds the string denoting the file_extension field in the database. + FieldFileExtension = "file_extension" + // FieldImageCount holds the string denoting the image_count field in the database. + FieldImageCount = "image_count" + // FieldErrorCode holds the string denoting the error_code field in the database. + FieldErrorCode = "error_code" + // FieldErrorMessage holds the string denoting the error_message field in the database. + FieldErrorMessage = "error_message" + // FieldBilledAmount holds the string denoting the billed_amount field in the database. + FieldBilledAmount = "billed_amount" + // FieldCreatedAt holds the string denoting the created_at field in the database. + FieldCreatedAt = "created_at" + // FieldIndexedAt holds the string denoting the indexed_at field in the database. + FieldIndexedAt = "indexed_at" + // Table holds the table name of the batchimageitem in the database. + Table = "batch_image_items" +) + +// Columns holds all SQL columns for batchimageitem fields. +var Columns = []string{ + FieldID, + FieldJobID, + FieldCustomID, + FieldStatus, + FieldRequestHash, + FieldPromptPreview, + FieldProviderSourceObject, + FieldSourceLineNumber, + FieldSourceByteOffset, + FieldSourceByteLength, + FieldMimeType, + FieldFileExtension, + FieldImageCount, + FieldErrorCode, + FieldErrorMessage, + FieldBilledAmount, + FieldCreatedAt, + FieldIndexedAt, +} + +// ValidColumn reports if the column name is valid (part of the table columns). +func ValidColumn(column string) bool { + for i := range Columns { + if column == Columns[i] { + return true + } + } + return false +} + +var ( + // JobIDValidator is a validator for the "job_id" field. It is called by the builders before save. + JobIDValidator func(string) error + // CustomIDValidator is a validator for the "custom_id" field. It is called by the builders before save. + CustomIDValidator func(string) error + // StatusValidator is a validator for the "status" field. It is called by the builders before save. + StatusValidator func(string) error + // RequestHashValidator is a validator for the "request_hash" field. It is called by the builders before save. + RequestHashValidator func(string) error + // ProviderSourceObjectValidator is a validator for the "provider_source_object" field. It is called by the builders before save. + ProviderSourceObjectValidator func(string) error + // MimeTypeValidator is a validator for the "mime_type" field. It is called by the builders before save. + MimeTypeValidator func(string) error + // FileExtensionValidator is a validator for the "file_extension" field. It is called by the builders before save. + FileExtensionValidator func(string) error + // DefaultImageCount holds the default value on creation for the "image_count" field. + DefaultImageCount int + // ErrorCodeValidator is a validator for the "error_code" field. It is called by the builders before save. + ErrorCodeValidator func(string) error + // DefaultCreatedAt holds the default value on creation for the "created_at" field. + DefaultCreatedAt func() time.Time +) + +// OrderOption defines the ordering options for the BatchImageItem queries. +type OrderOption func(*sql.Selector) + +// ByID orders the results by the id field. +func ByID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldID, opts...).ToFunc() +} + +// ByJobID orders the results by the job_id field. +func ByJobID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldJobID, opts...).ToFunc() +} + +// ByCustomID orders the results by the custom_id field. +func ByCustomID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldCustomID, opts...).ToFunc() +} + +// ByStatus orders the results by the status field. +func ByStatus(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldStatus, opts...).ToFunc() +} + +// ByRequestHash orders the results by the request_hash field. +func ByRequestHash(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldRequestHash, opts...).ToFunc() +} + +// ByPromptPreview orders the results by the prompt_preview field. +func ByPromptPreview(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldPromptPreview, opts...).ToFunc() +} + +// ByProviderSourceObject orders the results by the provider_source_object field. +func ByProviderSourceObject(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldProviderSourceObject, opts...).ToFunc() +} + +// BySourceLineNumber orders the results by the source_line_number field. +func BySourceLineNumber(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldSourceLineNumber, opts...).ToFunc() +} + +// BySourceByteOffset orders the results by the source_byte_offset field. +func BySourceByteOffset(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldSourceByteOffset, opts...).ToFunc() +} + +// BySourceByteLength orders the results by the source_byte_length field. +func BySourceByteLength(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldSourceByteLength, opts...).ToFunc() +} + +// ByMimeType orders the results by the mime_type field. +func ByMimeType(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldMimeType, opts...).ToFunc() +} + +// ByFileExtension orders the results by the file_extension field. +func ByFileExtension(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldFileExtension, opts...).ToFunc() +} + +// ByImageCount orders the results by the image_count field. +func ByImageCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldImageCount, opts...).ToFunc() +} + +// ByErrorCode orders the results by the error_code field. +func ByErrorCode(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldErrorCode, opts...).ToFunc() +} + +// ByErrorMessage orders the results by the error_message field. +func ByErrorMessage(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldErrorMessage, opts...).ToFunc() +} + +// ByBilledAmount orders the results by the billed_amount field. +func ByBilledAmount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldBilledAmount, opts...).ToFunc() +} + +// ByCreatedAt orders the results by the created_at field. +func ByCreatedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldCreatedAt, opts...).ToFunc() +} + +// ByIndexedAt orders the results by the indexed_at field. +func ByIndexedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldIndexedAt, opts...).ToFunc() +} diff --git a/backend/ent/batchimageitem/where.go b/backend/ent/batchimageitem/where.go new file mode 100644 index 0000000000..55dc32fde7 --- /dev/null +++ b/backend/ent/batchimageitem/where.go @@ -0,0 +1,1205 @@ +// Code generated by ent, DO NOT EDIT. + +package batchimageitem + +import ( + "time" + + "entgo.io/ent/dialect/sql" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// ID filters vertices based on their ID field. +func ID(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldID, id)) +} + +// IDEQ applies the EQ predicate on the ID field. +func IDEQ(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldID, id)) +} + +// IDNEQ applies the NEQ predicate on the ID field. +func IDNEQ(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldID, id)) +} + +// IDIn applies the In predicate on the ID field. +func IDIn(ids ...int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldID, ids...)) +} + +// IDNotIn applies the NotIn predicate on the ID field. +func IDNotIn(ids ...int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldID, ids...)) +} + +// IDGT applies the GT predicate on the ID field. +func IDGT(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldID, id)) +} + +// IDGTE applies the GTE predicate on the ID field. +func IDGTE(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldID, id)) +} + +// IDLT applies the LT predicate on the ID field. +func IDLT(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldID, id)) +} + +// IDLTE applies the LTE predicate on the ID field. +func IDLTE(id int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldID, id)) +} + +// JobID applies equality check predicate on the "job_id" field. It's identical to JobIDEQ. +func JobID(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldJobID, v)) +} + +// CustomID applies equality check predicate on the "custom_id" field. It's identical to CustomIDEQ. +func CustomID(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldCustomID, v)) +} + +// Status applies equality check predicate on the "status" field. It's identical to StatusEQ. +func Status(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldStatus, v)) +} + +// RequestHash applies equality check predicate on the "request_hash" field. It's identical to RequestHashEQ. +func RequestHash(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldRequestHash, v)) +} + +// PromptPreview applies equality check predicate on the "prompt_preview" field. It's identical to PromptPreviewEQ. +func PromptPreview(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldPromptPreview, v)) +} + +// ProviderSourceObject applies equality check predicate on the "provider_source_object" field. It's identical to ProviderSourceObjectEQ. +func ProviderSourceObject(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldProviderSourceObject, v)) +} + +// SourceLineNumber applies equality check predicate on the "source_line_number" field. It's identical to SourceLineNumberEQ. +func SourceLineNumber(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldSourceLineNumber, v)) +} + +// SourceByteOffset applies equality check predicate on the "source_byte_offset" field. It's identical to SourceByteOffsetEQ. +func SourceByteOffset(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldSourceByteOffset, v)) +} + +// SourceByteLength applies equality check predicate on the "source_byte_length" field. It's identical to SourceByteLengthEQ. +func SourceByteLength(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldSourceByteLength, v)) +} + +// MimeType applies equality check predicate on the "mime_type" field. It's identical to MimeTypeEQ. +func MimeType(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldMimeType, v)) +} + +// FileExtension applies equality check predicate on the "file_extension" field. It's identical to FileExtensionEQ. +func FileExtension(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldFileExtension, v)) +} + +// ImageCount applies equality check predicate on the "image_count" field. It's identical to ImageCountEQ. +func ImageCount(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldImageCount, v)) +} + +// ErrorCode applies equality check predicate on the "error_code" field. It's identical to ErrorCodeEQ. +func ErrorCode(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldErrorCode, v)) +} + +// ErrorMessage applies equality check predicate on the "error_message" field. It's identical to ErrorMessageEQ. +func ErrorMessage(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldErrorMessage, v)) +} + +// BilledAmount applies equality check predicate on the "billed_amount" field. It's identical to BilledAmountEQ. +func BilledAmount(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldBilledAmount, v)) +} + +// CreatedAt applies equality check predicate on the "created_at" field. It's identical to CreatedAtEQ. +func CreatedAt(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldCreatedAt, v)) +} + +// IndexedAt applies equality check predicate on the "indexed_at" field. It's identical to IndexedAtEQ. +func IndexedAt(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldIndexedAt, v)) +} + +// JobIDEQ applies the EQ predicate on the "job_id" field. +func JobIDEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldJobID, v)) +} + +// JobIDNEQ applies the NEQ predicate on the "job_id" field. +func JobIDNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldJobID, v)) +} + +// JobIDIn applies the In predicate on the "job_id" field. +func JobIDIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldJobID, vs...)) +} + +// JobIDNotIn applies the NotIn predicate on the "job_id" field. +func JobIDNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldJobID, vs...)) +} + +// JobIDGT applies the GT predicate on the "job_id" field. +func JobIDGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldJobID, v)) +} + +// JobIDGTE applies the GTE predicate on the "job_id" field. +func JobIDGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldJobID, v)) +} + +// JobIDLT applies the LT predicate on the "job_id" field. +func JobIDLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldJobID, v)) +} + +// JobIDLTE applies the LTE predicate on the "job_id" field. +func JobIDLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldJobID, v)) +} + +// JobIDContains applies the Contains predicate on the "job_id" field. +func JobIDContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldJobID, v)) +} + +// JobIDHasPrefix applies the HasPrefix predicate on the "job_id" field. +func JobIDHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldJobID, v)) +} + +// JobIDHasSuffix applies the HasSuffix predicate on the "job_id" field. +func JobIDHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldJobID, v)) +} + +// JobIDEqualFold applies the EqualFold predicate on the "job_id" field. +func JobIDEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldJobID, v)) +} + +// JobIDContainsFold applies the ContainsFold predicate on the "job_id" field. +func JobIDContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldJobID, v)) +} + +// CustomIDEQ applies the EQ predicate on the "custom_id" field. +func CustomIDEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldCustomID, v)) +} + +// CustomIDNEQ applies the NEQ predicate on the "custom_id" field. +func CustomIDNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldCustomID, v)) +} + +// CustomIDIn applies the In predicate on the "custom_id" field. +func CustomIDIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldCustomID, vs...)) +} + +// CustomIDNotIn applies the NotIn predicate on the "custom_id" field. +func CustomIDNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldCustomID, vs...)) +} + +// CustomIDGT applies the GT predicate on the "custom_id" field. +func CustomIDGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldCustomID, v)) +} + +// CustomIDGTE applies the GTE predicate on the "custom_id" field. +func CustomIDGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldCustomID, v)) +} + +// CustomIDLT applies the LT predicate on the "custom_id" field. +func CustomIDLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldCustomID, v)) +} + +// CustomIDLTE applies the LTE predicate on the "custom_id" field. +func CustomIDLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldCustomID, v)) +} + +// CustomIDContains applies the Contains predicate on the "custom_id" field. +func CustomIDContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldCustomID, v)) +} + +// CustomIDHasPrefix applies the HasPrefix predicate on the "custom_id" field. +func CustomIDHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldCustomID, v)) +} + +// CustomIDHasSuffix applies the HasSuffix predicate on the "custom_id" field. +func CustomIDHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldCustomID, v)) +} + +// CustomIDEqualFold applies the EqualFold predicate on the "custom_id" field. +func CustomIDEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldCustomID, v)) +} + +// CustomIDContainsFold applies the ContainsFold predicate on the "custom_id" field. +func CustomIDContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldCustomID, v)) +} + +// StatusEQ applies the EQ predicate on the "status" field. +func StatusEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldStatus, v)) +} + +// StatusNEQ applies the NEQ predicate on the "status" field. +func StatusNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldStatus, v)) +} + +// StatusIn applies the In predicate on the "status" field. +func StatusIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldStatus, vs...)) +} + +// StatusNotIn applies the NotIn predicate on the "status" field. +func StatusNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldStatus, vs...)) +} + +// StatusGT applies the GT predicate on the "status" field. +func StatusGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldStatus, v)) +} + +// StatusGTE applies the GTE predicate on the "status" field. +func StatusGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldStatus, v)) +} + +// StatusLT applies the LT predicate on the "status" field. +func StatusLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldStatus, v)) +} + +// StatusLTE applies the LTE predicate on the "status" field. +func StatusLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldStatus, v)) +} + +// StatusContains applies the Contains predicate on the "status" field. +func StatusContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldStatus, v)) +} + +// StatusHasPrefix applies the HasPrefix predicate on the "status" field. +func StatusHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldStatus, v)) +} + +// StatusHasSuffix applies the HasSuffix predicate on the "status" field. +func StatusHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldStatus, v)) +} + +// StatusEqualFold applies the EqualFold predicate on the "status" field. +func StatusEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldStatus, v)) +} + +// StatusContainsFold applies the ContainsFold predicate on the "status" field. +func StatusContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldStatus, v)) +} + +// RequestHashEQ applies the EQ predicate on the "request_hash" field. +func RequestHashEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldRequestHash, v)) +} + +// RequestHashNEQ applies the NEQ predicate on the "request_hash" field. +func RequestHashNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldRequestHash, v)) +} + +// RequestHashIn applies the In predicate on the "request_hash" field. +func RequestHashIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldRequestHash, vs...)) +} + +// RequestHashNotIn applies the NotIn predicate on the "request_hash" field. +func RequestHashNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldRequestHash, vs...)) +} + +// RequestHashGT applies the GT predicate on the "request_hash" field. +func RequestHashGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldRequestHash, v)) +} + +// RequestHashGTE applies the GTE predicate on the "request_hash" field. +func RequestHashGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldRequestHash, v)) +} + +// RequestHashLT applies the LT predicate on the "request_hash" field. +func RequestHashLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldRequestHash, v)) +} + +// RequestHashLTE applies the LTE predicate on the "request_hash" field. +func RequestHashLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldRequestHash, v)) +} + +// RequestHashContains applies the Contains predicate on the "request_hash" field. +func RequestHashContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldRequestHash, v)) +} + +// RequestHashHasPrefix applies the HasPrefix predicate on the "request_hash" field. +func RequestHashHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldRequestHash, v)) +} + +// RequestHashHasSuffix applies the HasSuffix predicate on the "request_hash" field. +func RequestHashHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldRequestHash, v)) +} + +// RequestHashIsNil applies the IsNil predicate on the "request_hash" field. +func RequestHashIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldRequestHash)) +} + +// RequestHashNotNil applies the NotNil predicate on the "request_hash" field. +func RequestHashNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldRequestHash)) +} + +// RequestHashEqualFold applies the EqualFold predicate on the "request_hash" field. +func RequestHashEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldRequestHash, v)) +} + +// RequestHashContainsFold applies the ContainsFold predicate on the "request_hash" field. +func RequestHashContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldRequestHash, v)) +} + +// PromptPreviewEQ applies the EQ predicate on the "prompt_preview" field. +func PromptPreviewEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldPromptPreview, v)) +} + +// PromptPreviewNEQ applies the NEQ predicate on the "prompt_preview" field. +func PromptPreviewNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldPromptPreview, v)) +} + +// PromptPreviewIn applies the In predicate on the "prompt_preview" field. +func PromptPreviewIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldPromptPreview, vs...)) +} + +// PromptPreviewNotIn applies the NotIn predicate on the "prompt_preview" field. +func PromptPreviewNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldPromptPreview, vs...)) +} + +// PromptPreviewGT applies the GT predicate on the "prompt_preview" field. +func PromptPreviewGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldPromptPreview, v)) +} + +// PromptPreviewGTE applies the GTE predicate on the "prompt_preview" field. +func PromptPreviewGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldPromptPreview, v)) +} + +// PromptPreviewLT applies the LT predicate on the "prompt_preview" field. +func PromptPreviewLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldPromptPreview, v)) +} + +// PromptPreviewLTE applies the LTE predicate on the "prompt_preview" field. +func PromptPreviewLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldPromptPreview, v)) +} + +// PromptPreviewContains applies the Contains predicate on the "prompt_preview" field. +func PromptPreviewContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldPromptPreview, v)) +} + +// PromptPreviewHasPrefix applies the HasPrefix predicate on the "prompt_preview" field. +func PromptPreviewHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldPromptPreview, v)) +} + +// PromptPreviewHasSuffix applies the HasSuffix predicate on the "prompt_preview" field. +func PromptPreviewHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldPromptPreview, v)) +} + +// PromptPreviewIsNil applies the IsNil predicate on the "prompt_preview" field. +func PromptPreviewIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldPromptPreview)) +} + +// PromptPreviewNotNil applies the NotNil predicate on the "prompt_preview" field. +func PromptPreviewNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldPromptPreview)) +} + +// PromptPreviewEqualFold applies the EqualFold predicate on the "prompt_preview" field. +func PromptPreviewEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldPromptPreview, v)) +} + +// PromptPreviewContainsFold applies the ContainsFold predicate on the "prompt_preview" field. +func PromptPreviewContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldPromptPreview, v)) +} + +// ProviderSourceObjectEQ applies the EQ predicate on the "provider_source_object" field. +func ProviderSourceObjectEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectNEQ applies the NEQ predicate on the "provider_source_object" field. +func ProviderSourceObjectNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectIn applies the In predicate on the "provider_source_object" field. +func ProviderSourceObjectIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldProviderSourceObject, vs...)) +} + +// ProviderSourceObjectNotIn applies the NotIn predicate on the "provider_source_object" field. +func ProviderSourceObjectNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldProviderSourceObject, vs...)) +} + +// ProviderSourceObjectGT applies the GT predicate on the "provider_source_object" field. +func ProviderSourceObjectGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectGTE applies the GTE predicate on the "provider_source_object" field. +func ProviderSourceObjectGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectLT applies the LT predicate on the "provider_source_object" field. +func ProviderSourceObjectLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectLTE applies the LTE predicate on the "provider_source_object" field. +func ProviderSourceObjectLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectContains applies the Contains predicate on the "provider_source_object" field. +func ProviderSourceObjectContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectHasPrefix applies the HasPrefix predicate on the "provider_source_object" field. +func ProviderSourceObjectHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectHasSuffix applies the HasSuffix predicate on the "provider_source_object" field. +func ProviderSourceObjectHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectIsNil applies the IsNil predicate on the "provider_source_object" field. +func ProviderSourceObjectIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldProviderSourceObject)) +} + +// ProviderSourceObjectNotNil applies the NotNil predicate on the "provider_source_object" field. +func ProviderSourceObjectNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldProviderSourceObject)) +} + +// ProviderSourceObjectEqualFold applies the EqualFold predicate on the "provider_source_object" field. +func ProviderSourceObjectEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldProviderSourceObject, v)) +} + +// ProviderSourceObjectContainsFold applies the ContainsFold predicate on the "provider_source_object" field. +func ProviderSourceObjectContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldProviderSourceObject, v)) +} + +// SourceLineNumberEQ applies the EQ predicate on the "source_line_number" field. +func SourceLineNumberEQ(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldSourceLineNumber, v)) +} + +// SourceLineNumberNEQ applies the NEQ predicate on the "source_line_number" field. +func SourceLineNumberNEQ(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldSourceLineNumber, v)) +} + +// SourceLineNumberIn applies the In predicate on the "source_line_number" field. +func SourceLineNumberIn(vs ...int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldSourceLineNumber, vs...)) +} + +// SourceLineNumberNotIn applies the NotIn predicate on the "source_line_number" field. +func SourceLineNumberNotIn(vs ...int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldSourceLineNumber, vs...)) +} + +// SourceLineNumberGT applies the GT predicate on the "source_line_number" field. +func SourceLineNumberGT(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldSourceLineNumber, v)) +} + +// SourceLineNumberGTE applies the GTE predicate on the "source_line_number" field. +func SourceLineNumberGTE(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldSourceLineNumber, v)) +} + +// SourceLineNumberLT applies the LT predicate on the "source_line_number" field. +func SourceLineNumberLT(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldSourceLineNumber, v)) +} + +// SourceLineNumberLTE applies the LTE predicate on the "source_line_number" field. +func SourceLineNumberLTE(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldSourceLineNumber, v)) +} + +// SourceLineNumberIsNil applies the IsNil predicate on the "source_line_number" field. +func SourceLineNumberIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldSourceLineNumber)) +} + +// SourceLineNumberNotNil applies the NotNil predicate on the "source_line_number" field. +func SourceLineNumberNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldSourceLineNumber)) +} + +// SourceByteOffsetEQ applies the EQ predicate on the "source_byte_offset" field. +func SourceByteOffsetEQ(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldSourceByteOffset, v)) +} + +// SourceByteOffsetNEQ applies the NEQ predicate on the "source_byte_offset" field. +func SourceByteOffsetNEQ(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldSourceByteOffset, v)) +} + +// SourceByteOffsetIn applies the In predicate on the "source_byte_offset" field. +func SourceByteOffsetIn(vs ...int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldSourceByteOffset, vs...)) +} + +// SourceByteOffsetNotIn applies the NotIn predicate on the "source_byte_offset" field. +func SourceByteOffsetNotIn(vs ...int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldSourceByteOffset, vs...)) +} + +// SourceByteOffsetGT applies the GT predicate on the "source_byte_offset" field. +func SourceByteOffsetGT(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldSourceByteOffset, v)) +} + +// SourceByteOffsetGTE applies the GTE predicate on the "source_byte_offset" field. +func SourceByteOffsetGTE(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldSourceByteOffset, v)) +} + +// SourceByteOffsetLT applies the LT predicate on the "source_byte_offset" field. +func SourceByteOffsetLT(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldSourceByteOffset, v)) +} + +// SourceByteOffsetLTE applies the LTE predicate on the "source_byte_offset" field. +func SourceByteOffsetLTE(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldSourceByteOffset, v)) +} + +// SourceByteOffsetIsNil applies the IsNil predicate on the "source_byte_offset" field. +func SourceByteOffsetIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldSourceByteOffset)) +} + +// SourceByteOffsetNotNil applies the NotNil predicate on the "source_byte_offset" field. +func SourceByteOffsetNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldSourceByteOffset)) +} + +// SourceByteLengthEQ applies the EQ predicate on the "source_byte_length" field. +func SourceByteLengthEQ(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldSourceByteLength, v)) +} + +// SourceByteLengthNEQ applies the NEQ predicate on the "source_byte_length" field. +func SourceByteLengthNEQ(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldSourceByteLength, v)) +} + +// SourceByteLengthIn applies the In predicate on the "source_byte_length" field. +func SourceByteLengthIn(vs ...int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldSourceByteLength, vs...)) +} + +// SourceByteLengthNotIn applies the NotIn predicate on the "source_byte_length" field. +func SourceByteLengthNotIn(vs ...int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldSourceByteLength, vs...)) +} + +// SourceByteLengthGT applies the GT predicate on the "source_byte_length" field. +func SourceByteLengthGT(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldSourceByteLength, v)) +} + +// SourceByteLengthGTE applies the GTE predicate on the "source_byte_length" field. +func SourceByteLengthGTE(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldSourceByteLength, v)) +} + +// SourceByteLengthLT applies the LT predicate on the "source_byte_length" field. +func SourceByteLengthLT(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldSourceByteLength, v)) +} + +// SourceByteLengthLTE applies the LTE predicate on the "source_byte_length" field. +func SourceByteLengthLTE(v int64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldSourceByteLength, v)) +} + +// SourceByteLengthIsNil applies the IsNil predicate on the "source_byte_length" field. +func SourceByteLengthIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldSourceByteLength)) +} + +// SourceByteLengthNotNil applies the NotNil predicate on the "source_byte_length" field. +func SourceByteLengthNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldSourceByteLength)) +} + +// MimeTypeEQ applies the EQ predicate on the "mime_type" field. +func MimeTypeEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldMimeType, v)) +} + +// MimeTypeNEQ applies the NEQ predicate on the "mime_type" field. +func MimeTypeNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldMimeType, v)) +} + +// MimeTypeIn applies the In predicate on the "mime_type" field. +func MimeTypeIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldMimeType, vs...)) +} + +// MimeTypeNotIn applies the NotIn predicate on the "mime_type" field. +func MimeTypeNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldMimeType, vs...)) +} + +// MimeTypeGT applies the GT predicate on the "mime_type" field. +func MimeTypeGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldMimeType, v)) +} + +// MimeTypeGTE applies the GTE predicate on the "mime_type" field. +func MimeTypeGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldMimeType, v)) +} + +// MimeTypeLT applies the LT predicate on the "mime_type" field. +func MimeTypeLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldMimeType, v)) +} + +// MimeTypeLTE applies the LTE predicate on the "mime_type" field. +func MimeTypeLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldMimeType, v)) +} + +// MimeTypeContains applies the Contains predicate on the "mime_type" field. +func MimeTypeContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldMimeType, v)) +} + +// MimeTypeHasPrefix applies the HasPrefix predicate on the "mime_type" field. +func MimeTypeHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldMimeType, v)) +} + +// MimeTypeHasSuffix applies the HasSuffix predicate on the "mime_type" field. +func MimeTypeHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldMimeType, v)) +} + +// MimeTypeIsNil applies the IsNil predicate on the "mime_type" field. +func MimeTypeIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldMimeType)) +} + +// MimeTypeNotNil applies the NotNil predicate on the "mime_type" field. +func MimeTypeNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldMimeType)) +} + +// MimeTypeEqualFold applies the EqualFold predicate on the "mime_type" field. +func MimeTypeEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldMimeType, v)) +} + +// MimeTypeContainsFold applies the ContainsFold predicate on the "mime_type" field. +func MimeTypeContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldMimeType, v)) +} + +// FileExtensionEQ applies the EQ predicate on the "file_extension" field. +func FileExtensionEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldFileExtension, v)) +} + +// FileExtensionNEQ applies the NEQ predicate on the "file_extension" field. +func FileExtensionNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldFileExtension, v)) +} + +// FileExtensionIn applies the In predicate on the "file_extension" field. +func FileExtensionIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldFileExtension, vs...)) +} + +// FileExtensionNotIn applies the NotIn predicate on the "file_extension" field. +func FileExtensionNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldFileExtension, vs...)) +} + +// FileExtensionGT applies the GT predicate on the "file_extension" field. +func FileExtensionGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldFileExtension, v)) +} + +// FileExtensionGTE applies the GTE predicate on the "file_extension" field. +func FileExtensionGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldFileExtension, v)) +} + +// FileExtensionLT applies the LT predicate on the "file_extension" field. +func FileExtensionLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldFileExtension, v)) +} + +// FileExtensionLTE applies the LTE predicate on the "file_extension" field. +func FileExtensionLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldFileExtension, v)) +} + +// FileExtensionContains applies the Contains predicate on the "file_extension" field. +func FileExtensionContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldFileExtension, v)) +} + +// FileExtensionHasPrefix applies the HasPrefix predicate on the "file_extension" field. +func FileExtensionHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldFileExtension, v)) +} + +// FileExtensionHasSuffix applies the HasSuffix predicate on the "file_extension" field. +func FileExtensionHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldFileExtension, v)) +} + +// FileExtensionIsNil applies the IsNil predicate on the "file_extension" field. +func FileExtensionIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldFileExtension)) +} + +// FileExtensionNotNil applies the NotNil predicate on the "file_extension" field. +func FileExtensionNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldFileExtension)) +} + +// FileExtensionEqualFold applies the EqualFold predicate on the "file_extension" field. +func FileExtensionEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldFileExtension, v)) +} + +// FileExtensionContainsFold applies the ContainsFold predicate on the "file_extension" field. +func FileExtensionContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldFileExtension, v)) +} + +// ImageCountEQ applies the EQ predicate on the "image_count" field. +func ImageCountEQ(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldImageCount, v)) +} + +// ImageCountNEQ applies the NEQ predicate on the "image_count" field. +func ImageCountNEQ(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldImageCount, v)) +} + +// ImageCountIn applies the In predicate on the "image_count" field. +func ImageCountIn(vs ...int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldImageCount, vs...)) +} + +// ImageCountNotIn applies the NotIn predicate on the "image_count" field. +func ImageCountNotIn(vs ...int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldImageCount, vs...)) +} + +// ImageCountGT applies the GT predicate on the "image_count" field. +func ImageCountGT(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldImageCount, v)) +} + +// ImageCountGTE applies the GTE predicate on the "image_count" field. +func ImageCountGTE(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldImageCount, v)) +} + +// ImageCountLT applies the LT predicate on the "image_count" field. +func ImageCountLT(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldImageCount, v)) +} + +// ImageCountLTE applies the LTE predicate on the "image_count" field. +func ImageCountLTE(v int) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldImageCount, v)) +} + +// ErrorCodeEQ applies the EQ predicate on the "error_code" field. +func ErrorCodeEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldErrorCode, v)) +} + +// ErrorCodeNEQ applies the NEQ predicate on the "error_code" field. +func ErrorCodeNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldErrorCode, v)) +} + +// ErrorCodeIn applies the In predicate on the "error_code" field. +func ErrorCodeIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldErrorCode, vs...)) +} + +// ErrorCodeNotIn applies the NotIn predicate on the "error_code" field. +func ErrorCodeNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldErrorCode, vs...)) +} + +// ErrorCodeGT applies the GT predicate on the "error_code" field. +func ErrorCodeGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldErrorCode, v)) +} + +// ErrorCodeGTE applies the GTE predicate on the "error_code" field. +func ErrorCodeGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldErrorCode, v)) +} + +// ErrorCodeLT applies the LT predicate on the "error_code" field. +func ErrorCodeLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldErrorCode, v)) +} + +// ErrorCodeLTE applies the LTE predicate on the "error_code" field. +func ErrorCodeLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldErrorCode, v)) +} + +// ErrorCodeContains applies the Contains predicate on the "error_code" field. +func ErrorCodeContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldErrorCode, v)) +} + +// ErrorCodeHasPrefix applies the HasPrefix predicate on the "error_code" field. +func ErrorCodeHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldErrorCode, v)) +} + +// ErrorCodeHasSuffix applies the HasSuffix predicate on the "error_code" field. +func ErrorCodeHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldErrorCode, v)) +} + +// ErrorCodeIsNil applies the IsNil predicate on the "error_code" field. +func ErrorCodeIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldErrorCode)) +} + +// ErrorCodeNotNil applies the NotNil predicate on the "error_code" field. +func ErrorCodeNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldErrorCode)) +} + +// ErrorCodeEqualFold applies the EqualFold predicate on the "error_code" field. +func ErrorCodeEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldErrorCode, v)) +} + +// ErrorCodeContainsFold applies the ContainsFold predicate on the "error_code" field. +func ErrorCodeContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldErrorCode, v)) +} + +// ErrorMessageEQ applies the EQ predicate on the "error_message" field. +func ErrorMessageEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldErrorMessage, v)) +} + +// ErrorMessageNEQ applies the NEQ predicate on the "error_message" field. +func ErrorMessageNEQ(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldErrorMessage, v)) +} + +// ErrorMessageIn applies the In predicate on the "error_message" field. +func ErrorMessageIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldErrorMessage, vs...)) +} + +// ErrorMessageNotIn applies the NotIn predicate on the "error_message" field. +func ErrorMessageNotIn(vs ...string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldErrorMessage, vs...)) +} + +// ErrorMessageGT applies the GT predicate on the "error_message" field. +func ErrorMessageGT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldErrorMessage, v)) +} + +// ErrorMessageGTE applies the GTE predicate on the "error_message" field. +func ErrorMessageGTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldErrorMessage, v)) +} + +// ErrorMessageLT applies the LT predicate on the "error_message" field. +func ErrorMessageLT(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldErrorMessage, v)) +} + +// ErrorMessageLTE applies the LTE predicate on the "error_message" field. +func ErrorMessageLTE(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldErrorMessage, v)) +} + +// ErrorMessageContains applies the Contains predicate on the "error_message" field. +func ErrorMessageContains(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContains(FieldErrorMessage, v)) +} + +// ErrorMessageHasPrefix applies the HasPrefix predicate on the "error_message" field. +func ErrorMessageHasPrefix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasPrefix(FieldErrorMessage, v)) +} + +// ErrorMessageHasSuffix applies the HasSuffix predicate on the "error_message" field. +func ErrorMessageHasSuffix(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldHasSuffix(FieldErrorMessage, v)) +} + +// ErrorMessageIsNil applies the IsNil predicate on the "error_message" field. +func ErrorMessageIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldErrorMessage)) +} + +// ErrorMessageNotNil applies the NotNil predicate on the "error_message" field. +func ErrorMessageNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldErrorMessage)) +} + +// ErrorMessageEqualFold applies the EqualFold predicate on the "error_message" field. +func ErrorMessageEqualFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEqualFold(FieldErrorMessage, v)) +} + +// ErrorMessageContainsFold applies the ContainsFold predicate on the "error_message" field. +func ErrorMessageContainsFold(v string) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldContainsFold(FieldErrorMessage, v)) +} + +// BilledAmountEQ applies the EQ predicate on the "billed_amount" field. +func BilledAmountEQ(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldBilledAmount, v)) +} + +// BilledAmountNEQ applies the NEQ predicate on the "billed_amount" field. +func BilledAmountNEQ(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldBilledAmount, v)) +} + +// BilledAmountIn applies the In predicate on the "billed_amount" field. +func BilledAmountIn(vs ...float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldBilledAmount, vs...)) +} + +// BilledAmountNotIn applies the NotIn predicate on the "billed_amount" field. +func BilledAmountNotIn(vs ...float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldBilledAmount, vs...)) +} + +// BilledAmountGT applies the GT predicate on the "billed_amount" field. +func BilledAmountGT(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldBilledAmount, v)) +} + +// BilledAmountGTE applies the GTE predicate on the "billed_amount" field. +func BilledAmountGTE(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldBilledAmount, v)) +} + +// BilledAmountLT applies the LT predicate on the "billed_amount" field. +func BilledAmountLT(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldBilledAmount, v)) +} + +// BilledAmountLTE applies the LTE predicate on the "billed_amount" field. +func BilledAmountLTE(v float64) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldBilledAmount, v)) +} + +// BilledAmountIsNil applies the IsNil predicate on the "billed_amount" field. +func BilledAmountIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldBilledAmount)) +} + +// BilledAmountNotNil applies the NotNil predicate on the "billed_amount" field. +func BilledAmountNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldBilledAmount)) +} + +// CreatedAtEQ applies the EQ predicate on the "created_at" field. +func CreatedAtEQ(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldCreatedAt, v)) +} + +// CreatedAtNEQ applies the NEQ predicate on the "created_at" field. +func CreatedAtNEQ(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldCreatedAt, v)) +} + +// CreatedAtIn applies the In predicate on the "created_at" field. +func CreatedAtIn(vs ...time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldCreatedAt, vs...)) +} + +// CreatedAtNotIn applies the NotIn predicate on the "created_at" field. +func CreatedAtNotIn(vs ...time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldCreatedAt, vs...)) +} + +// CreatedAtGT applies the GT predicate on the "created_at" field. +func CreatedAtGT(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldCreatedAt, v)) +} + +// CreatedAtGTE applies the GTE predicate on the "created_at" field. +func CreatedAtGTE(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldCreatedAt, v)) +} + +// CreatedAtLT applies the LT predicate on the "created_at" field. +func CreatedAtLT(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldCreatedAt, v)) +} + +// CreatedAtLTE applies the LTE predicate on the "created_at" field. +func CreatedAtLTE(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldCreatedAt, v)) +} + +// IndexedAtEQ applies the EQ predicate on the "indexed_at" field. +func IndexedAtEQ(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldEQ(FieldIndexedAt, v)) +} + +// IndexedAtNEQ applies the NEQ predicate on the "indexed_at" field. +func IndexedAtNEQ(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNEQ(FieldIndexedAt, v)) +} + +// IndexedAtIn applies the In predicate on the "indexed_at" field. +func IndexedAtIn(vs ...time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIn(FieldIndexedAt, vs...)) +} + +// IndexedAtNotIn applies the NotIn predicate on the "indexed_at" field. +func IndexedAtNotIn(vs ...time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotIn(FieldIndexedAt, vs...)) +} + +// IndexedAtGT applies the GT predicate on the "indexed_at" field. +func IndexedAtGT(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGT(FieldIndexedAt, v)) +} + +// IndexedAtGTE applies the GTE predicate on the "indexed_at" field. +func IndexedAtGTE(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldGTE(FieldIndexedAt, v)) +} + +// IndexedAtLT applies the LT predicate on the "indexed_at" field. +func IndexedAtLT(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLT(FieldIndexedAt, v)) +} + +// IndexedAtLTE applies the LTE predicate on the "indexed_at" field. +func IndexedAtLTE(v time.Time) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldLTE(FieldIndexedAt, v)) +} + +// IndexedAtIsNil applies the IsNil predicate on the "indexed_at" field. +func IndexedAtIsNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldIsNull(FieldIndexedAt)) +} + +// IndexedAtNotNil applies the NotNil predicate on the "indexed_at" field. +func IndexedAtNotNil() predicate.BatchImageItem { + return predicate.BatchImageItem(sql.FieldNotNull(FieldIndexedAt)) +} + +// And groups predicates with the AND operator between them. +func And(predicates ...predicate.BatchImageItem) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.AndPredicates(predicates...)) +} + +// Or groups predicates with the OR operator between them. +func Or(predicates ...predicate.BatchImageItem) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.OrPredicates(predicates...)) +} + +// Not applies the not operator on the given predicate. +func Not(p predicate.BatchImageItem) predicate.BatchImageItem { + return predicate.BatchImageItem(sql.NotPredicates(p)) +} diff --git a/backend/ent/batchimageitem_create.go b/backend/ent/batchimageitem_create.go new file mode 100644 index 0000000000..f9ee14998a --- /dev/null +++ b/backend/ent/batchimageitem_create.go @@ -0,0 +1,1745 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "errors" + "fmt" + "time" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" +) + +// BatchImageItemCreate is the builder for creating a BatchImageItem entity. +type BatchImageItemCreate struct { + config + mutation *BatchImageItemMutation + hooks []Hook + conflict []sql.ConflictOption +} + +// SetJobID sets the "job_id" field. +func (_c *BatchImageItemCreate) SetJobID(v string) *BatchImageItemCreate { + _c.mutation.SetJobID(v) + return _c +} + +// SetCustomID sets the "custom_id" field. +func (_c *BatchImageItemCreate) SetCustomID(v string) *BatchImageItemCreate { + _c.mutation.SetCustomID(v) + return _c +} + +// SetStatus sets the "status" field. +func (_c *BatchImageItemCreate) SetStatus(v string) *BatchImageItemCreate { + _c.mutation.SetStatus(v) + return _c +} + +// SetRequestHash sets the "request_hash" field. +func (_c *BatchImageItemCreate) SetRequestHash(v string) *BatchImageItemCreate { + _c.mutation.SetRequestHash(v) + return _c +} + +// SetNillableRequestHash sets the "request_hash" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableRequestHash(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetRequestHash(*v) + } + return _c +} + +// SetPromptPreview sets the "prompt_preview" field. +func (_c *BatchImageItemCreate) SetPromptPreview(v string) *BatchImageItemCreate { + _c.mutation.SetPromptPreview(v) + return _c +} + +// SetNillablePromptPreview sets the "prompt_preview" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillablePromptPreview(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetPromptPreview(*v) + } + return _c +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (_c *BatchImageItemCreate) SetProviderSourceObject(v string) *BatchImageItemCreate { + _c.mutation.SetProviderSourceObject(v) + return _c +} + +// SetNillableProviderSourceObject sets the "provider_source_object" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableProviderSourceObject(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetProviderSourceObject(*v) + } + return _c +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (_c *BatchImageItemCreate) SetSourceLineNumber(v int) *BatchImageItemCreate { + _c.mutation.SetSourceLineNumber(v) + return _c +} + +// SetNillableSourceLineNumber sets the "source_line_number" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableSourceLineNumber(v *int) *BatchImageItemCreate { + if v != nil { + _c.SetSourceLineNumber(*v) + } + return _c +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (_c *BatchImageItemCreate) SetSourceByteOffset(v int64) *BatchImageItemCreate { + _c.mutation.SetSourceByteOffset(v) + return _c +} + +// SetNillableSourceByteOffset sets the "source_byte_offset" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableSourceByteOffset(v *int64) *BatchImageItemCreate { + if v != nil { + _c.SetSourceByteOffset(*v) + } + return _c +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (_c *BatchImageItemCreate) SetSourceByteLength(v int64) *BatchImageItemCreate { + _c.mutation.SetSourceByteLength(v) + return _c +} + +// SetNillableSourceByteLength sets the "source_byte_length" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableSourceByteLength(v *int64) *BatchImageItemCreate { + if v != nil { + _c.SetSourceByteLength(*v) + } + return _c +} + +// SetMimeType sets the "mime_type" field. +func (_c *BatchImageItemCreate) SetMimeType(v string) *BatchImageItemCreate { + _c.mutation.SetMimeType(v) + return _c +} + +// SetNillableMimeType sets the "mime_type" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableMimeType(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetMimeType(*v) + } + return _c +} + +// SetFileExtension sets the "file_extension" field. +func (_c *BatchImageItemCreate) SetFileExtension(v string) *BatchImageItemCreate { + _c.mutation.SetFileExtension(v) + return _c +} + +// SetNillableFileExtension sets the "file_extension" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableFileExtension(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetFileExtension(*v) + } + return _c +} + +// SetImageCount sets the "image_count" field. +func (_c *BatchImageItemCreate) SetImageCount(v int) *BatchImageItemCreate { + _c.mutation.SetImageCount(v) + return _c +} + +// SetNillableImageCount sets the "image_count" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableImageCount(v *int) *BatchImageItemCreate { + if v != nil { + _c.SetImageCount(*v) + } + return _c +} + +// SetErrorCode sets the "error_code" field. +func (_c *BatchImageItemCreate) SetErrorCode(v string) *BatchImageItemCreate { + _c.mutation.SetErrorCode(v) + return _c +} + +// SetNillableErrorCode sets the "error_code" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableErrorCode(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetErrorCode(*v) + } + return _c +} + +// SetErrorMessage sets the "error_message" field. +func (_c *BatchImageItemCreate) SetErrorMessage(v string) *BatchImageItemCreate { + _c.mutation.SetErrorMessage(v) + return _c +} + +// SetNillableErrorMessage sets the "error_message" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableErrorMessage(v *string) *BatchImageItemCreate { + if v != nil { + _c.SetErrorMessage(*v) + } + return _c +} + +// SetBilledAmount sets the "billed_amount" field. +func (_c *BatchImageItemCreate) SetBilledAmount(v float64) *BatchImageItemCreate { + _c.mutation.SetBilledAmount(v) + return _c +} + +// SetNillableBilledAmount sets the "billed_amount" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableBilledAmount(v *float64) *BatchImageItemCreate { + if v != nil { + _c.SetBilledAmount(*v) + } + return _c +} + +// SetCreatedAt sets the "created_at" field. +func (_c *BatchImageItemCreate) SetCreatedAt(v time.Time) *BatchImageItemCreate { + _c.mutation.SetCreatedAt(v) + return _c +} + +// SetNillableCreatedAt sets the "created_at" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableCreatedAt(v *time.Time) *BatchImageItemCreate { + if v != nil { + _c.SetCreatedAt(*v) + } + return _c +} + +// SetIndexedAt sets the "indexed_at" field. +func (_c *BatchImageItemCreate) SetIndexedAt(v time.Time) *BatchImageItemCreate { + _c.mutation.SetIndexedAt(v) + return _c +} + +// SetNillableIndexedAt sets the "indexed_at" field if the given value is not nil. +func (_c *BatchImageItemCreate) SetNillableIndexedAt(v *time.Time) *BatchImageItemCreate { + if v != nil { + _c.SetIndexedAt(*v) + } + return _c +} + +// Mutation returns the BatchImageItemMutation object of the builder. +func (_c *BatchImageItemCreate) Mutation() *BatchImageItemMutation { + return _c.mutation +} + +// Save creates the BatchImageItem in the database. +func (_c *BatchImageItemCreate) Save(ctx context.Context) (*BatchImageItem, error) { + _c.defaults() + return withHooks(ctx, _c.sqlSave, _c.mutation, _c.hooks) +} + +// SaveX calls Save and panics if Save returns an error. +func (_c *BatchImageItemCreate) SaveX(ctx context.Context) *BatchImageItem { + v, err := _c.Save(ctx) + if err != nil { + panic(err) + } + return v +} + +// Exec executes the query. +func (_c *BatchImageItemCreate) Exec(ctx context.Context) error { + _, err := _c.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_c *BatchImageItemCreate) ExecX(ctx context.Context) { + if err := _c.Exec(ctx); err != nil { + panic(err) + } +} + +// defaults sets the default values of the builder before save. +func (_c *BatchImageItemCreate) defaults() { + if _, ok := _c.mutation.ImageCount(); !ok { + v := batchimageitem.DefaultImageCount + _c.mutation.SetImageCount(v) + } + if _, ok := _c.mutation.CreatedAt(); !ok { + v := batchimageitem.DefaultCreatedAt() + _c.mutation.SetCreatedAt(v) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_c *BatchImageItemCreate) check() error { + if _, ok := _c.mutation.JobID(); !ok { + return &ValidationError{Name: "job_id", err: errors.New(`ent: missing required field "BatchImageItem.job_id"`)} + } + if v, ok := _c.mutation.JobID(); ok { + if err := batchimageitem.JobIDValidator(v); err != nil { + return &ValidationError{Name: "job_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.job_id": %w`, err)} + } + } + if _, ok := _c.mutation.CustomID(); !ok { + return &ValidationError{Name: "custom_id", err: errors.New(`ent: missing required field "BatchImageItem.custom_id"`)} + } + if v, ok := _c.mutation.CustomID(); ok { + if err := batchimageitem.CustomIDValidator(v); err != nil { + return &ValidationError{Name: "custom_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.custom_id": %w`, err)} + } + } + if _, ok := _c.mutation.Status(); !ok { + return &ValidationError{Name: "status", err: errors.New(`ent: missing required field "BatchImageItem.status"`)} + } + if v, ok := _c.mutation.Status(); ok { + if err := batchimageitem.StatusValidator(v); err != nil { + return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.status": %w`, err)} + } + } + if v, ok := _c.mutation.RequestHash(); ok { + if err := batchimageitem.RequestHashValidator(v); err != nil { + return &ValidationError{Name: "request_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.request_hash": %w`, err)} + } + } + if v, ok := _c.mutation.ProviderSourceObject(); ok { + if err := batchimageitem.ProviderSourceObjectValidator(v); err != nil { + return &ValidationError{Name: "provider_source_object", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.provider_source_object": %w`, err)} + } + } + if v, ok := _c.mutation.MimeType(); ok { + if err := batchimageitem.MimeTypeValidator(v); err != nil { + return &ValidationError{Name: "mime_type", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.mime_type": %w`, err)} + } + } + if v, ok := _c.mutation.FileExtension(); ok { + if err := batchimageitem.FileExtensionValidator(v); err != nil { + return &ValidationError{Name: "file_extension", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.file_extension": %w`, err)} + } + } + if _, ok := _c.mutation.ImageCount(); !ok { + return &ValidationError{Name: "image_count", err: errors.New(`ent: missing required field "BatchImageItem.image_count"`)} + } + if v, ok := _c.mutation.ErrorCode(); ok { + if err := batchimageitem.ErrorCodeValidator(v); err != nil { + return &ValidationError{Name: "error_code", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.error_code": %w`, err)} + } + } + if _, ok := _c.mutation.CreatedAt(); !ok { + return &ValidationError{Name: "created_at", err: errors.New(`ent: missing required field "BatchImageItem.created_at"`)} + } + return nil +} + +func (_c *BatchImageItemCreate) sqlSave(ctx context.Context) (*BatchImageItem, error) { + if err := _c.check(); err != nil { + return nil, err + } + _node, _spec := _c.createSpec() + if err := sqlgraph.CreateNode(ctx, _c.driver, _spec); err != nil { + if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return nil, err + } + id := _spec.ID.Value.(int64) + _node.ID = int64(id) + _c.mutation.id = &_node.ID + _c.mutation.done = true + return _node, nil +} + +func (_c *BatchImageItemCreate) createSpec() (*BatchImageItem, *sqlgraph.CreateSpec) { + var ( + _node = &BatchImageItem{config: _c.config} + _spec = sqlgraph.NewCreateSpec(batchimageitem.Table, sqlgraph.NewFieldSpec(batchimageitem.FieldID, field.TypeInt64)) + ) + _spec.OnConflict = _c.conflict + if value, ok := _c.mutation.JobID(); ok { + _spec.SetField(batchimageitem.FieldJobID, field.TypeString, value) + _node.JobID = value + } + if value, ok := _c.mutation.CustomID(); ok { + _spec.SetField(batchimageitem.FieldCustomID, field.TypeString, value) + _node.CustomID = value + } + if value, ok := _c.mutation.Status(); ok { + _spec.SetField(batchimageitem.FieldStatus, field.TypeString, value) + _node.Status = value + } + if value, ok := _c.mutation.RequestHash(); ok { + _spec.SetField(batchimageitem.FieldRequestHash, field.TypeString, value) + _node.RequestHash = &value + } + if value, ok := _c.mutation.PromptPreview(); ok { + _spec.SetField(batchimageitem.FieldPromptPreview, field.TypeString, value) + _node.PromptPreview = &value + } + if value, ok := _c.mutation.ProviderSourceObject(); ok { + _spec.SetField(batchimageitem.FieldProviderSourceObject, field.TypeString, value) + _node.ProviderSourceObject = &value + } + if value, ok := _c.mutation.SourceLineNumber(); ok { + _spec.SetField(batchimageitem.FieldSourceLineNumber, field.TypeInt, value) + _node.SourceLineNumber = &value + } + if value, ok := _c.mutation.SourceByteOffset(); ok { + _spec.SetField(batchimageitem.FieldSourceByteOffset, field.TypeInt64, value) + _node.SourceByteOffset = &value + } + if value, ok := _c.mutation.SourceByteLength(); ok { + _spec.SetField(batchimageitem.FieldSourceByteLength, field.TypeInt64, value) + _node.SourceByteLength = &value + } + if value, ok := _c.mutation.MimeType(); ok { + _spec.SetField(batchimageitem.FieldMimeType, field.TypeString, value) + _node.MimeType = &value + } + if value, ok := _c.mutation.FileExtension(); ok { + _spec.SetField(batchimageitem.FieldFileExtension, field.TypeString, value) + _node.FileExtension = &value + } + if value, ok := _c.mutation.ImageCount(); ok { + _spec.SetField(batchimageitem.FieldImageCount, field.TypeInt, value) + _node.ImageCount = value + } + if value, ok := _c.mutation.ErrorCode(); ok { + _spec.SetField(batchimageitem.FieldErrorCode, field.TypeString, value) + _node.ErrorCode = &value + } + if value, ok := _c.mutation.ErrorMessage(); ok { + _spec.SetField(batchimageitem.FieldErrorMessage, field.TypeString, value) + _node.ErrorMessage = &value + } + if value, ok := _c.mutation.BilledAmount(); ok { + _spec.SetField(batchimageitem.FieldBilledAmount, field.TypeFloat64, value) + _node.BilledAmount = &value + } + if value, ok := _c.mutation.CreatedAt(); ok { + _spec.SetField(batchimageitem.FieldCreatedAt, field.TypeTime, value) + _node.CreatedAt = value + } + if value, ok := _c.mutation.IndexedAt(); ok { + _spec.SetField(batchimageitem.FieldIndexedAt, field.TypeTime, value) + _node.IndexedAt = &value + } + return _node, _spec +} + +// OnConflict allows configuring the `ON CONFLICT` / `ON DUPLICATE KEY` clause +// of the `INSERT` statement. For example: +// +// client.BatchImageItem.Create(). +// SetJobID(v). +// OnConflict( +// // Update the row with the new values +// // the was proposed for insertion. +// sql.ResolveWithNewValues(), +// ). +// // Override some of the fields with custom +// // update values. +// Update(func(u *ent.BatchImageItemUpsert) { +// SetJobID(v+v). +// }). +// Exec(ctx) +func (_c *BatchImageItemCreate) OnConflict(opts ...sql.ConflictOption) *BatchImageItemUpsertOne { + _c.conflict = opts + return &BatchImageItemUpsertOne{ + create: _c, + } +} + +// OnConflictColumns calls `OnConflict` and configures the columns +// as conflict target. Using this option is equivalent to using: +// +// client.BatchImageItem.Create(). +// OnConflict(sql.ConflictColumns(columns...)). +// Exec(ctx) +func (_c *BatchImageItemCreate) OnConflictColumns(columns ...string) *BatchImageItemUpsertOne { + _c.conflict = append(_c.conflict, sql.ConflictColumns(columns...)) + return &BatchImageItemUpsertOne{ + create: _c, + } +} + +type ( + // BatchImageItemUpsertOne is the builder for "upsert"-ing + // one BatchImageItem node. + BatchImageItemUpsertOne struct { + create *BatchImageItemCreate + } + + // BatchImageItemUpsert is the "OnConflict" setter. + BatchImageItemUpsert struct { + *sql.UpdateSet + } +) + +// SetJobID sets the "job_id" field. +func (u *BatchImageItemUpsert) SetJobID(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldJobID, v) + return u +} + +// UpdateJobID sets the "job_id" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateJobID() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldJobID) + return u +} + +// SetCustomID sets the "custom_id" field. +func (u *BatchImageItemUpsert) SetCustomID(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldCustomID, v) + return u +} + +// UpdateCustomID sets the "custom_id" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateCustomID() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldCustomID) + return u +} + +// SetStatus sets the "status" field. +func (u *BatchImageItemUpsert) SetStatus(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldStatus, v) + return u +} + +// UpdateStatus sets the "status" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateStatus() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldStatus) + return u +} + +// SetRequestHash sets the "request_hash" field. +func (u *BatchImageItemUpsert) SetRequestHash(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldRequestHash, v) + return u +} + +// UpdateRequestHash sets the "request_hash" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateRequestHash() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldRequestHash) + return u +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (u *BatchImageItemUpsert) ClearRequestHash() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldRequestHash) + return u +} + +// SetPromptPreview sets the "prompt_preview" field. +func (u *BatchImageItemUpsert) SetPromptPreview(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldPromptPreview, v) + return u +} + +// UpdatePromptPreview sets the "prompt_preview" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdatePromptPreview() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldPromptPreview) + return u +} + +// ClearPromptPreview clears the value of the "prompt_preview" field. +func (u *BatchImageItemUpsert) ClearPromptPreview() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldPromptPreview) + return u +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (u *BatchImageItemUpsert) SetProviderSourceObject(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldProviderSourceObject, v) + return u +} + +// UpdateProviderSourceObject sets the "provider_source_object" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateProviderSourceObject() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldProviderSourceObject) + return u +} + +// ClearProviderSourceObject clears the value of the "provider_source_object" field. +func (u *BatchImageItemUpsert) ClearProviderSourceObject() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldProviderSourceObject) + return u +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (u *BatchImageItemUpsert) SetSourceLineNumber(v int) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldSourceLineNumber, v) + return u +} + +// UpdateSourceLineNumber sets the "source_line_number" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateSourceLineNumber() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldSourceLineNumber) + return u +} + +// AddSourceLineNumber adds v to the "source_line_number" field. +func (u *BatchImageItemUpsert) AddSourceLineNumber(v int) *BatchImageItemUpsert { + u.Add(batchimageitem.FieldSourceLineNumber, v) + return u +} + +// ClearSourceLineNumber clears the value of the "source_line_number" field. +func (u *BatchImageItemUpsert) ClearSourceLineNumber() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldSourceLineNumber) + return u +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (u *BatchImageItemUpsert) SetSourceByteOffset(v int64) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldSourceByteOffset, v) + return u +} + +// UpdateSourceByteOffset sets the "source_byte_offset" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateSourceByteOffset() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldSourceByteOffset) + return u +} + +// AddSourceByteOffset adds v to the "source_byte_offset" field. +func (u *BatchImageItemUpsert) AddSourceByteOffset(v int64) *BatchImageItemUpsert { + u.Add(batchimageitem.FieldSourceByteOffset, v) + return u +} + +// ClearSourceByteOffset clears the value of the "source_byte_offset" field. +func (u *BatchImageItemUpsert) ClearSourceByteOffset() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldSourceByteOffset) + return u +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (u *BatchImageItemUpsert) SetSourceByteLength(v int64) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldSourceByteLength, v) + return u +} + +// UpdateSourceByteLength sets the "source_byte_length" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateSourceByteLength() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldSourceByteLength) + return u +} + +// AddSourceByteLength adds v to the "source_byte_length" field. +func (u *BatchImageItemUpsert) AddSourceByteLength(v int64) *BatchImageItemUpsert { + u.Add(batchimageitem.FieldSourceByteLength, v) + return u +} + +// ClearSourceByteLength clears the value of the "source_byte_length" field. +func (u *BatchImageItemUpsert) ClearSourceByteLength() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldSourceByteLength) + return u +} + +// SetMimeType sets the "mime_type" field. +func (u *BatchImageItemUpsert) SetMimeType(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldMimeType, v) + return u +} + +// UpdateMimeType sets the "mime_type" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateMimeType() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldMimeType) + return u +} + +// ClearMimeType clears the value of the "mime_type" field. +func (u *BatchImageItemUpsert) ClearMimeType() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldMimeType) + return u +} + +// SetFileExtension sets the "file_extension" field. +func (u *BatchImageItemUpsert) SetFileExtension(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldFileExtension, v) + return u +} + +// UpdateFileExtension sets the "file_extension" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateFileExtension() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldFileExtension) + return u +} + +// ClearFileExtension clears the value of the "file_extension" field. +func (u *BatchImageItemUpsert) ClearFileExtension() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldFileExtension) + return u +} + +// SetImageCount sets the "image_count" field. +func (u *BatchImageItemUpsert) SetImageCount(v int) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldImageCount, v) + return u +} + +// UpdateImageCount sets the "image_count" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateImageCount() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldImageCount) + return u +} + +// AddImageCount adds v to the "image_count" field. +func (u *BatchImageItemUpsert) AddImageCount(v int) *BatchImageItemUpsert { + u.Add(batchimageitem.FieldImageCount, v) + return u +} + +// SetErrorCode sets the "error_code" field. +func (u *BatchImageItemUpsert) SetErrorCode(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldErrorCode, v) + return u +} + +// UpdateErrorCode sets the "error_code" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateErrorCode() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldErrorCode) + return u +} + +// ClearErrorCode clears the value of the "error_code" field. +func (u *BatchImageItemUpsert) ClearErrorCode() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldErrorCode) + return u +} + +// SetErrorMessage sets the "error_message" field. +func (u *BatchImageItemUpsert) SetErrorMessage(v string) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldErrorMessage, v) + return u +} + +// UpdateErrorMessage sets the "error_message" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateErrorMessage() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldErrorMessage) + return u +} + +// ClearErrorMessage clears the value of the "error_message" field. +func (u *BatchImageItemUpsert) ClearErrorMessage() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldErrorMessage) + return u +} + +// SetBilledAmount sets the "billed_amount" field. +func (u *BatchImageItemUpsert) SetBilledAmount(v float64) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldBilledAmount, v) + return u +} + +// UpdateBilledAmount sets the "billed_amount" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateBilledAmount() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldBilledAmount) + return u +} + +// AddBilledAmount adds v to the "billed_amount" field. +func (u *BatchImageItemUpsert) AddBilledAmount(v float64) *BatchImageItemUpsert { + u.Add(batchimageitem.FieldBilledAmount, v) + return u +} + +// ClearBilledAmount clears the value of the "billed_amount" field. +func (u *BatchImageItemUpsert) ClearBilledAmount() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldBilledAmount) + return u +} + +// SetIndexedAt sets the "indexed_at" field. +func (u *BatchImageItemUpsert) SetIndexedAt(v time.Time) *BatchImageItemUpsert { + u.Set(batchimageitem.FieldIndexedAt, v) + return u +} + +// UpdateIndexedAt sets the "indexed_at" field to the value that was provided on create. +func (u *BatchImageItemUpsert) UpdateIndexedAt() *BatchImageItemUpsert { + u.SetExcluded(batchimageitem.FieldIndexedAt) + return u +} + +// ClearIndexedAt clears the value of the "indexed_at" field. +func (u *BatchImageItemUpsert) ClearIndexedAt() *BatchImageItemUpsert { + u.SetNull(batchimageitem.FieldIndexedAt) + return u +} + +// UpdateNewValues updates the mutable fields using the new values that were set on create. +// Using this option is equivalent to using: +// +// client.BatchImageItem.Create(). +// OnConflict( +// sql.ResolveWithNewValues(), +// ). +// Exec(ctx) +func (u *BatchImageItemUpsertOne) UpdateNewValues() *BatchImageItemUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWithNewValues()) + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(s *sql.UpdateSet) { + if _, exists := u.create.mutation.CreatedAt(); exists { + s.SetIgnore(batchimageitem.FieldCreatedAt) + } + })) + return u +} + +// Ignore sets each column to itself in case of conflict. +// Using this option is equivalent to using: +// +// client.BatchImageItem.Create(). +// OnConflict(sql.ResolveWithIgnore()). +// Exec(ctx) +func (u *BatchImageItemUpsertOne) Ignore() *BatchImageItemUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWithIgnore()) + return u +} + +// DoNothing configures the conflict_action to `DO NOTHING`. +// Supported only by SQLite and PostgreSQL. +func (u *BatchImageItemUpsertOne) DoNothing() *BatchImageItemUpsertOne { + u.create.conflict = append(u.create.conflict, sql.DoNothing()) + return u +} + +// Update allows overriding fields `UPDATE` values. See the BatchImageItemCreate.OnConflict +// documentation for more info. +func (u *BatchImageItemUpsertOne) Update(set func(*BatchImageItemUpsert)) *BatchImageItemUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(update *sql.UpdateSet) { + set(&BatchImageItemUpsert{UpdateSet: update}) + })) + return u +} + +// SetJobID sets the "job_id" field. +func (u *BatchImageItemUpsertOne) SetJobID(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetJobID(v) + }) +} + +// UpdateJobID sets the "job_id" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateJobID() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateJobID() + }) +} + +// SetCustomID sets the "custom_id" field. +func (u *BatchImageItemUpsertOne) SetCustomID(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetCustomID(v) + }) +} + +// UpdateCustomID sets the "custom_id" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateCustomID() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateCustomID() + }) +} + +// SetStatus sets the "status" field. +func (u *BatchImageItemUpsertOne) SetStatus(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetStatus(v) + }) +} + +// UpdateStatus sets the "status" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateStatus() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateStatus() + }) +} + +// SetRequestHash sets the "request_hash" field. +func (u *BatchImageItemUpsertOne) SetRequestHash(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetRequestHash(v) + }) +} + +// UpdateRequestHash sets the "request_hash" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateRequestHash() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateRequestHash() + }) +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (u *BatchImageItemUpsertOne) ClearRequestHash() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearRequestHash() + }) +} + +// SetPromptPreview sets the "prompt_preview" field. +func (u *BatchImageItemUpsertOne) SetPromptPreview(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetPromptPreview(v) + }) +} + +// UpdatePromptPreview sets the "prompt_preview" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdatePromptPreview() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdatePromptPreview() + }) +} + +// ClearPromptPreview clears the value of the "prompt_preview" field. +func (u *BatchImageItemUpsertOne) ClearPromptPreview() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearPromptPreview() + }) +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (u *BatchImageItemUpsertOne) SetProviderSourceObject(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetProviderSourceObject(v) + }) +} + +// UpdateProviderSourceObject sets the "provider_source_object" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateProviderSourceObject() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateProviderSourceObject() + }) +} + +// ClearProviderSourceObject clears the value of the "provider_source_object" field. +func (u *BatchImageItemUpsertOne) ClearProviderSourceObject() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearProviderSourceObject() + }) +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (u *BatchImageItemUpsertOne) SetSourceLineNumber(v int) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetSourceLineNumber(v) + }) +} + +// AddSourceLineNumber adds v to the "source_line_number" field. +func (u *BatchImageItemUpsertOne) AddSourceLineNumber(v int) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddSourceLineNumber(v) + }) +} + +// UpdateSourceLineNumber sets the "source_line_number" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateSourceLineNumber() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateSourceLineNumber() + }) +} + +// ClearSourceLineNumber clears the value of the "source_line_number" field. +func (u *BatchImageItemUpsertOne) ClearSourceLineNumber() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearSourceLineNumber() + }) +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (u *BatchImageItemUpsertOne) SetSourceByteOffset(v int64) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetSourceByteOffset(v) + }) +} + +// AddSourceByteOffset adds v to the "source_byte_offset" field. +func (u *BatchImageItemUpsertOne) AddSourceByteOffset(v int64) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddSourceByteOffset(v) + }) +} + +// UpdateSourceByteOffset sets the "source_byte_offset" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateSourceByteOffset() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateSourceByteOffset() + }) +} + +// ClearSourceByteOffset clears the value of the "source_byte_offset" field. +func (u *BatchImageItemUpsertOne) ClearSourceByteOffset() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearSourceByteOffset() + }) +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (u *BatchImageItemUpsertOne) SetSourceByteLength(v int64) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetSourceByteLength(v) + }) +} + +// AddSourceByteLength adds v to the "source_byte_length" field. +func (u *BatchImageItemUpsertOne) AddSourceByteLength(v int64) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddSourceByteLength(v) + }) +} + +// UpdateSourceByteLength sets the "source_byte_length" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateSourceByteLength() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateSourceByteLength() + }) +} + +// ClearSourceByteLength clears the value of the "source_byte_length" field. +func (u *BatchImageItemUpsertOne) ClearSourceByteLength() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearSourceByteLength() + }) +} + +// SetMimeType sets the "mime_type" field. +func (u *BatchImageItemUpsertOne) SetMimeType(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetMimeType(v) + }) +} + +// UpdateMimeType sets the "mime_type" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateMimeType() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateMimeType() + }) +} + +// ClearMimeType clears the value of the "mime_type" field. +func (u *BatchImageItemUpsertOne) ClearMimeType() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearMimeType() + }) +} + +// SetFileExtension sets the "file_extension" field. +func (u *BatchImageItemUpsertOne) SetFileExtension(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetFileExtension(v) + }) +} + +// UpdateFileExtension sets the "file_extension" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateFileExtension() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateFileExtension() + }) +} + +// ClearFileExtension clears the value of the "file_extension" field. +func (u *BatchImageItemUpsertOne) ClearFileExtension() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearFileExtension() + }) +} + +// SetImageCount sets the "image_count" field. +func (u *BatchImageItemUpsertOne) SetImageCount(v int) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetImageCount(v) + }) +} + +// AddImageCount adds v to the "image_count" field. +func (u *BatchImageItemUpsertOne) AddImageCount(v int) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddImageCount(v) + }) +} + +// UpdateImageCount sets the "image_count" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateImageCount() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateImageCount() + }) +} + +// SetErrorCode sets the "error_code" field. +func (u *BatchImageItemUpsertOne) SetErrorCode(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetErrorCode(v) + }) +} + +// UpdateErrorCode sets the "error_code" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateErrorCode() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateErrorCode() + }) +} + +// ClearErrorCode clears the value of the "error_code" field. +func (u *BatchImageItemUpsertOne) ClearErrorCode() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearErrorCode() + }) +} + +// SetErrorMessage sets the "error_message" field. +func (u *BatchImageItemUpsertOne) SetErrorMessage(v string) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetErrorMessage(v) + }) +} + +// UpdateErrorMessage sets the "error_message" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateErrorMessage() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateErrorMessage() + }) +} + +// ClearErrorMessage clears the value of the "error_message" field. +func (u *BatchImageItemUpsertOne) ClearErrorMessage() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearErrorMessage() + }) +} + +// SetBilledAmount sets the "billed_amount" field. +func (u *BatchImageItemUpsertOne) SetBilledAmount(v float64) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetBilledAmount(v) + }) +} + +// AddBilledAmount adds v to the "billed_amount" field. +func (u *BatchImageItemUpsertOne) AddBilledAmount(v float64) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddBilledAmount(v) + }) +} + +// UpdateBilledAmount sets the "billed_amount" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateBilledAmount() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateBilledAmount() + }) +} + +// ClearBilledAmount clears the value of the "billed_amount" field. +func (u *BatchImageItemUpsertOne) ClearBilledAmount() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearBilledAmount() + }) +} + +// SetIndexedAt sets the "indexed_at" field. +func (u *BatchImageItemUpsertOne) SetIndexedAt(v time.Time) *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetIndexedAt(v) + }) +} + +// UpdateIndexedAt sets the "indexed_at" field to the value that was provided on create. +func (u *BatchImageItemUpsertOne) UpdateIndexedAt() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateIndexedAt() + }) +} + +// ClearIndexedAt clears the value of the "indexed_at" field. +func (u *BatchImageItemUpsertOne) ClearIndexedAt() *BatchImageItemUpsertOne { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearIndexedAt() + }) +} + +// Exec executes the query. +func (u *BatchImageItemUpsertOne) Exec(ctx context.Context) error { + if len(u.create.conflict) == 0 { + return errors.New("ent: missing options for BatchImageItemCreate.OnConflict") + } + return u.create.Exec(ctx) +} + +// ExecX is like Exec, but panics if an error occurs. +func (u *BatchImageItemUpsertOne) ExecX(ctx context.Context) { + if err := u.create.Exec(ctx); err != nil { + panic(err) + } +} + +// Exec executes the UPSERT query and returns the inserted/updated ID. +func (u *BatchImageItemUpsertOne) ID(ctx context.Context) (id int64, err error) { + node, err := u.create.Save(ctx) + if err != nil { + return id, err + } + return node.ID, nil +} + +// IDX is like ID, but panics if an error occurs. +func (u *BatchImageItemUpsertOne) IDX(ctx context.Context) int64 { + id, err := u.ID(ctx) + if err != nil { + panic(err) + } + return id +} + +// BatchImageItemCreateBulk is the builder for creating many BatchImageItem entities in bulk. +type BatchImageItemCreateBulk struct { + config + err error + builders []*BatchImageItemCreate + conflict []sql.ConflictOption +} + +// Save creates the BatchImageItem entities in the database. +func (_c *BatchImageItemCreateBulk) Save(ctx context.Context) ([]*BatchImageItem, error) { + if _c.err != nil { + return nil, _c.err + } + specs := make([]*sqlgraph.CreateSpec, len(_c.builders)) + nodes := make([]*BatchImageItem, len(_c.builders)) + mutators := make([]Mutator, len(_c.builders)) + for i := range _c.builders { + func(i int, root context.Context) { + builder := _c.builders[i] + builder.defaults() + var mut Mutator = MutateFunc(func(ctx context.Context, m Mutation) (Value, error) { + mutation, ok := m.(*BatchImageItemMutation) + if !ok { + return nil, fmt.Errorf("unexpected mutation type %T", m) + } + if err := builder.check(); err != nil { + return nil, err + } + builder.mutation = mutation + var err error + nodes[i], specs[i] = builder.createSpec() + if i < len(mutators)-1 { + _, err = mutators[i+1].Mutate(root, _c.builders[i+1].mutation) + } else { + spec := &sqlgraph.BatchCreateSpec{Nodes: specs} + spec.OnConflict = _c.conflict + // Invoke the actual operation on the latest mutation in the chain. + if err = sqlgraph.BatchCreate(ctx, _c.driver, spec); err != nil { + if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + } + } + if err != nil { + return nil, err + } + mutation.id = &nodes[i].ID + if specs[i].ID.Value != nil { + id := specs[i].ID.Value.(int64) + nodes[i].ID = int64(id) + } + mutation.done = true + return nodes[i], nil + }) + for i := len(builder.hooks) - 1; i >= 0; i-- { + mut = builder.hooks[i](mut) + } + mutators[i] = mut + }(i, ctx) + } + if len(mutators) > 0 { + if _, err := mutators[0].Mutate(ctx, _c.builders[0].mutation); err != nil { + return nil, err + } + } + return nodes, nil +} + +// SaveX is like Save, but panics if an error occurs. +func (_c *BatchImageItemCreateBulk) SaveX(ctx context.Context) []*BatchImageItem { + v, err := _c.Save(ctx) + if err != nil { + panic(err) + } + return v +} + +// Exec executes the query. +func (_c *BatchImageItemCreateBulk) Exec(ctx context.Context) error { + _, err := _c.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_c *BatchImageItemCreateBulk) ExecX(ctx context.Context) { + if err := _c.Exec(ctx); err != nil { + panic(err) + } +} + +// OnConflict allows configuring the `ON CONFLICT` / `ON DUPLICATE KEY` clause +// of the `INSERT` statement. For example: +// +// client.BatchImageItem.CreateBulk(builders...). +// OnConflict( +// // Update the row with the new values +// // the was proposed for insertion. +// sql.ResolveWithNewValues(), +// ). +// // Override some of the fields with custom +// // update values. +// Update(func(u *ent.BatchImageItemUpsert) { +// SetJobID(v+v). +// }). +// Exec(ctx) +func (_c *BatchImageItemCreateBulk) OnConflict(opts ...sql.ConflictOption) *BatchImageItemUpsertBulk { + _c.conflict = opts + return &BatchImageItemUpsertBulk{ + create: _c, + } +} + +// OnConflictColumns calls `OnConflict` and configures the columns +// as conflict target. Using this option is equivalent to using: +// +// client.BatchImageItem.Create(). +// OnConflict(sql.ConflictColumns(columns...)). +// Exec(ctx) +func (_c *BatchImageItemCreateBulk) OnConflictColumns(columns ...string) *BatchImageItemUpsertBulk { + _c.conflict = append(_c.conflict, sql.ConflictColumns(columns...)) + return &BatchImageItemUpsertBulk{ + create: _c, + } +} + +// BatchImageItemUpsertBulk is the builder for "upsert"-ing +// a bulk of BatchImageItem nodes. +type BatchImageItemUpsertBulk struct { + create *BatchImageItemCreateBulk +} + +// UpdateNewValues updates the mutable fields using the new values that +// were set on create. Using this option is equivalent to using: +// +// client.BatchImageItem.Create(). +// OnConflict( +// sql.ResolveWithNewValues(), +// ). +// Exec(ctx) +func (u *BatchImageItemUpsertBulk) UpdateNewValues() *BatchImageItemUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWithNewValues()) + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(s *sql.UpdateSet) { + for _, b := range u.create.builders { + if _, exists := b.mutation.CreatedAt(); exists { + s.SetIgnore(batchimageitem.FieldCreatedAt) + } + } + })) + return u +} + +// Ignore sets each column to itself in case of conflict. +// Using this option is equivalent to using: +// +// client.BatchImageItem.Create(). +// OnConflict(sql.ResolveWithIgnore()). +// Exec(ctx) +func (u *BatchImageItemUpsertBulk) Ignore() *BatchImageItemUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWithIgnore()) + return u +} + +// DoNothing configures the conflict_action to `DO NOTHING`. +// Supported only by SQLite and PostgreSQL. +func (u *BatchImageItemUpsertBulk) DoNothing() *BatchImageItemUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.DoNothing()) + return u +} + +// Update allows overriding fields `UPDATE` values. See the BatchImageItemCreateBulk.OnConflict +// documentation for more info. +func (u *BatchImageItemUpsertBulk) Update(set func(*BatchImageItemUpsert)) *BatchImageItemUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(update *sql.UpdateSet) { + set(&BatchImageItemUpsert{UpdateSet: update}) + })) + return u +} + +// SetJobID sets the "job_id" field. +func (u *BatchImageItemUpsertBulk) SetJobID(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetJobID(v) + }) +} + +// UpdateJobID sets the "job_id" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateJobID() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateJobID() + }) +} + +// SetCustomID sets the "custom_id" field. +func (u *BatchImageItemUpsertBulk) SetCustomID(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetCustomID(v) + }) +} + +// UpdateCustomID sets the "custom_id" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateCustomID() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateCustomID() + }) +} + +// SetStatus sets the "status" field. +func (u *BatchImageItemUpsertBulk) SetStatus(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetStatus(v) + }) +} + +// UpdateStatus sets the "status" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateStatus() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateStatus() + }) +} + +// SetRequestHash sets the "request_hash" field. +func (u *BatchImageItemUpsertBulk) SetRequestHash(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetRequestHash(v) + }) +} + +// UpdateRequestHash sets the "request_hash" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateRequestHash() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateRequestHash() + }) +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (u *BatchImageItemUpsertBulk) ClearRequestHash() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearRequestHash() + }) +} + +// SetPromptPreview sets the "prompt_preview" field. +func (u *BatchImageItemUpsertBulk) SetPromptPreview(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetPromptPreview(v) + }) +} + +// UpdatePromptPreview sets the "prompt_preview" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdatePromptPreview() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdatePromptPreview() + }) +} + +// ClearPromptPreview clears the value of the "prompt_preview" field. +func (u *BatchImageItemUpsertBulk) ClearPromptPreview() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearPromptPreview() + }) +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (u *BatchImageItemUpsertBulk) SetProviderSourceObject(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetProviderSourceObject(v) + }) +} + +// UpdateProviderSourceObject sets the "provider_source_object" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateProviderSourceObject() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateProviderSourceObject() + }) +} + +// ClearProviderSourceObject clears the value of the "provider_source_object" field. +func (u *BatchImageItemUpsertBulk) ClearProviderSourceObject() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearProviderSourceObject() + }) +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (u *BatchImageItemUpsertBulk) SetSourceLineNumber(v int) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetSourceLineNumber(v) + }) +} + +// AddSourceLineNumber adds v to the "source_line_number" field. +func (u *BatchImageItemUpsertBulk) AddSourceLineNumber(v int) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddSourceLineNumber(v) + }) +} + +// UpdateSourceLineNumber sets the "source_line_number" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateSourceLineNumber() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateSourceLineNumber() + }) +} + +// ClearSourceLineNumber clears the value of the "source_line_number" field. +func (u *BatchImageItemUpsertBulk) ClearSourceLineNumber() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearSourceLineNumber() + }) +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (u *BatchImageItemUpsertBulk) SetSourceByteOffset(v int64) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetSourceByteOffset(v) + }) +} + +// AddSourceByteOffset adds v to the "source_byte_offset" field. +func (u *BatchImageItemUpsertBulk) AddSourceByteOffset(v int64) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddSourceByteOffset(v) + }) +} + +// UpdateSourceByteOffset sets the "source_byte_offset" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateSourceByteOffset() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateSourceByteOffset() + }) +} + +// ClearSourceByteOffset clears the value of the "source_byte_offset" field. +func (u *BatchImageItemUpsertBulk) ClearSourceByteOffset() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearSourceByteOffset() + }) +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (u *BatchImageItemUpsertBulk) SetSourceByteLength(v int64) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetSourceByteLength(v) + }) +} + +// AddSourceByteLength adds v to the "source_byte_length" field. +func (u *BatchImageItemUpsertBulk) AddSourceByteLength(v int64) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddSourceByteLength(v) + }) +} + +// UpdateSourceByteLength sets the "source_byte_length" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateSourceByteLength() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateSourceByteLength() + }) +} + +// ClearSourceByteLength clears the value of the "source_byte_length" field. +func (u *BatchImageItemUpsertBulk) ClearSourceByteLength() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearSourceByteLength() + }) +} + +// SetMimeType sets the "mime_type" field. +func (u *BatchImageItemUpsertBulk) SetMimeType(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetMimeType(v) + }) +} + +// UpdateMimeType sets the "mime_type" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateMimeType() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateMimeType() + }) +} + +// ClearMimeType clears the value of the "mime_type" field. +func (u *BatchImageItemUpsertBulk) ClearMimeType() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearMimeType() + }) +} + +// SetFileExtension sets the "file_extension" field. +func (u *BatchImageItemUpsertBulk) SetFileExtension(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetFileExtension(v) + }) +} + +// UpdateFileExtension sets the "file_extension" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateFileExtension() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateFileExtension() + }) +} + +// ClearFileExtension clears the value of the "file_extension" field. +func (u *BatchImageItemUpsertBulk) ClearFileExtension() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearFileExtension() + }) +} + +// SetImageCount sets the "image_count" field. +func (u *BatchImageItemUpsertBulk) SetImageCount(v int) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetImageCount(v) + }) +} + +// AddImageCount adds v to the "image_count" field. +func (u *BatchImageItemUpsertBulk) AddImageCount(v int) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddImageCount(v) + }) +} + +// UpdateImageCount sets the "image_count" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateImageCount() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateImageCount() + }) +} + +// SetErrorCode sets the "error_code" field. +func (u *BatchImageItemUpsertBulk) SetErrorCode(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetErrorCode(v) + }) +} + +// UpdateErrorCode sets the "error_code" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateErrorCode() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateErrorCode() + }) +} + +// ClearErrorCode clears the value of the "error_code" field. +func (u *BatchImageItemUpsertBulk) ClearErrorCode() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearErrorCode() + }) +} + +// SetErrorMessage sets the "error_message" field. +func (u *BatchImageItemUpsertBulk) SetErrorMessage(v string) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetErrorMessage(v) + }) +} + +// UpdateErrorMessage sets the "error_message" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateErrorMessage() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateErrorMessage() + }) +} + +// ClearErrorMessage clears the value of the "error_message" field. +func (u *BatchImageItemUpsertBulk) ClearErrorMessage() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearErrorMessage() + }) +} + +// SetBilledAmount sets the "billed_amount" field. +func (u *BatchImageItemUpsertBulk) SetBilledAmount(v float64) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetBilledAmount(v) + }) +} + +// AddBilledAmount adds v to the "billed_amount" field. +func (u *BatchImageItemUpsertBulk) AddBilledAmount(v float64) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.AddBilledAmount(v) + }) +} + +// UpdateBilledAmount sets the "billed_amount" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateBilledAmount() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateBilledAmount() + }) +} + +// ClearBilledAmount clears the value of the "billed_amount" field. +func (u *BatchImageItemUpsertBulk) ClearBilledAmount() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearBilledAmount() + }) +} + +// SetIndexedAt sets the "indexed_at" field. +func (u *BatchImageItemUpsertBulk) SetIndexedAt(v time.Time) *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.SetIndexedAt(v) + }) +} + +// UpdateIndexedAt sets the "indexed_at" field to the value that was provided on create. +func (u *BatchImageItemUpsertBulk) UpdateIndexedAt() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.UpdateIndexedAt() + }) +} + +// ClearIndexedAt clears the value of the "indexed_at" field. +func (u *BatchImageItemUpsertBulk) ClearIndexedAt() *BatchImageItemUpsertBulk { + return u.Update(func(s *BatchImageItemUpsert) { + s.ClearIndexedAt() + }) +} + +// Exec executes the query. +func (u *BatchImageItemUpsertBulk) Exec(ctx context.Context) error { + if u.create.err != nil { + return u.create.err + } + for i, b := range u.create.builders { + if len(b.conflict) != 0 { + return fmt.Errorf("ent: OnConflict was set for builder %d. Set it on the BatchImageItemCreateBulk instead", i) + } + } + if len(u.create.conflict) == 0 { + return errors.New("ent: missing options for BatchImageItemCreateBulk.OnConflict") + } + return u.create.Exec(ctx) +} + +// ExecX is like Exec, but panics if an error occurs. +func (u *BatchImageItemUpsertBulk) ExecX(ctx context.Context) { + if err := u.create.Exec(ctx); err != nil { + panic(err) + } +} diff --git a/backend/ent/batchimageitem_delete.go b/backend/ent/batchimageitem_delete.go new file mode 100644 index 0000000000..7aa3bf32e8 --- /dev/null +++ b/backend/ent/batchimageitem_delete.go @@ -0,0 +1,88 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageItemDelete is the builder for deleting a BatchImageItem entity. +type BatchImageItemDelete struct { + config + hooks []Hook + mutation *BatchImageItemMutation +} + +// Where appends a list predicates to the BatchImageItemDelete builder. +func (_d *BatchImageItemDelete) Where(ps ...predicate.BatchImageItem) *BatchImageItemDelete { + _d.mutation.Where(ps...) + return _d +} + +// Exec executes the deletion query and returns how many vertices were deleted. +func (_d *BatchImageItemDelete) Exec(ctx context.Context) (int, error) { + return withHooks(ctx, _d.sqlExec, _d.mutation, _d.hooks) +} + +// ExecX is like Exec, but panics if an error occurs. +func (_d *BatchImageItemDelete) ExecX(ctx context.Context) int { + n, err := _d.Exec(ctx) + if err != nil { + panic(err) + } + return n +} + +func (_d *BatchImageItemDelete) sqlExec(ctx context.Context) (int, error) { + _spec := sqlgraph.NewDeleteSpec(batchimageitem.Table, sqlgraph.NewFieldSpec(batchimageitem.FieldID, field.TypeInt64)) + if ps := _d.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + affected, err := sqlgraph.DeleteNodes(ctx, _d.driver, _spec) + if err != nil && sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + _d.mutation.done = true + return affected, err +} + +// BatchImageItemDeleteOne is the builder for deleting a single BatchImageItem entity. +type BatchImageItemDeleteOne struct { + _d *BatchImageItemDelete +} + +// Where appends a list predicates to the BatchImageItemDelete builder. +func (_d *BatchImageItemDeleteOne) Where(ps ...predicate.BatchImageItem) *BatchImageItemDeleteOne { + _d._d.mutation.Where(ps...) + return _d +} + +// Exec executes the deletion query. +func (_d *BatchImageItemDeleteOne) Exec(ctx context.Context) error { + n, err := _d._d.Exec(ctx) + switch { + case err != nil: + return err + case n == 0: + return &NotFoundError{batchimageitem.Label} + default: + return nil + } +} + +// ExecX is like Exec, but panics if an error occurs. +func (_d *BatchImageItemDeleteOne) ExecX(ctx context.Context) { + if err := _d.Exec(ctx); err != nil { + panic(err) + } +} diff --git a/backend/ent/batchimageitem_query.go b/backend/ent/batchimageitem_query.go new file mode 100644 index 0000000000..7e1d08f7be --- /dev/null +++ b/backend/ent/batchimageitem_query.go @@ -0,0 +1,564 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "fmt" + "math" + + "entgo.io/ent" + "entgo.io/ent/dialect" + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageItemQuery is the builder for querying BatchImageItem entities. +type BatchImageItemQuery struct { + config + ctx *QueryContext + order []batchimageitem.OrderOption + inters []Interceptor + predicates []predicate.BatchImageItem + modifiers []func(*sql.Selector) + // intermediate query (i.e. traversal path). + sql *sql.Selector + path func(context.Context) (*sql.Selector, error) +} + +// Where adds a new predicate for the BatchImageItemQuery builder. +func (_q *BatchImageItemQuery) Where(ps ...predicate.BatchImageItem) *BatchImageItemQuery { + _q.predicates = append(_q.predicates, ps...) + return _q +} + +// Limit the number of records to be returned by this query. +func (_q *BatchImageItemQuery) Limit(limit int) *BatchImageItemQuery { + _q.ctx.Limit = &limit + return _q +} + +// Offset to start from. +func (_q *BatchImageItemQuery) Offset(offset int) *BatchImageItemQuery { + _q.ctx.Offset = &offset + return _q +} + +// Unique configures the query builder to filter duplicate records on query. +// By default, unique is set to true, and can be disabled using this method. +func (_q *BatchImageItemQuery) Unique(unique bool) *BatchImageItemQuery { + _q.ctx.Unique = &unique + return _q +} + +// Order specifies how the records should be ordered. +func (_q *BatchImageItemQuery) Order(o ...batchimageitem.OrderOption) *BatchImageItemQuery { + _q.order = append(_q.order, o...) + return _q +} + +// First returns the first BatchImageItem entity from the query. +// Returns a *NotFoundError when no BatchImageItem was found. +func (_q *BatchImageItemQuery) First(ctx context.Context) (*BatchImageItem, error) { + nodes, err := _q.Limit(1).All(setContextOp(ctx, _q.ctx, ent.OpQueryFirst)) + if err != nil { + return nil, err + } + if len(nodes) == 0 { + return nil, &NotFoundError{batchimageitem.Label} + } + return nodes[0], nil +} + +// FirstX is like First, but panics if an error occurs. +func (_q *BatchImageItemQuery) FirstX(ctx context.Context) *BatchImageItem { + node, err := _q.First(ctx) + if err != nil && !IsNotFound(err) { + panic(err) + } + return node +} + +// FirstID returns the first BatchImageItem ID from the query. +// Returns a *NotFoundError when no BatchImageItem ID was found. +func (_q *BatchImageItemQuery) FirstID(ctx context.Context) (id int64, err error) { + var ids []int64 + if ids, err = _q.Limit(1).IDs(setContextOp(ctx, _q.ctx, ent.OpQueryFirstID)); err != nil { + return + } + if len(ids) == 0 { + err = &NotFoundError{batchimageitem.Label} + return + } + return ids[0], nil +} + +// FirstIDX is like FirstID, but panics if an error occurs. +func (_q *BatchImageItemQuery) FirstIDX(ctx context.Context) int64 { + id, err := _q.FirstID(ctx) + if err != nil && !IsNotFound(err) { + panic(err) + } + return id +} + +// Only returns a single BatchImageItem entity found by the query, ensuring it only returns one. +// Returns a *NotSingularError when more than one BatchImageItem entity is found. +// Returns a *NotFoundError when no BatchImageItem entities are found. +func (_q *BatchImageItemQuery) Only(ctx context.Context) (*BatchImageItem, error) { + nodes, err := _q.Limit(2).All(setContextOp(ctx, _q.ctx, ent.OpQueryOnly)) + if err != nil { + return nil, err + } + switch len(nodes) { + case 1: + return nodes[0], nil + case 0: + return nil, &NotFoundError{batchimageitem.Label} + default: + return nil, &NotSingularError{batchimageitem.Label} + } +} + +// OnlyX is like Only, but panics if an error occurs. +func (_q *BatchImageItemQuery) OnlyX(ctx context.Context) *BatchImageItem { + node, err := _q.Only(ctx) + if err != nil { + panic(err) + } + return node +} + +// OnlyID is like Only, but returns the only BatchImageItem ID in the query. +// Returns a *NotSingularError when more than one BatchImageItem ID is found. +// Returns a *NotFoundError when no entities are found. +func (_q *BatchImageItemQuery) OnlyID(ctx context.Context) (id int64, err error) { + var ids []int64 + if ids, err = _q.Limit(2).IDs(setContextOp(ctx, _q.ctx, ent.OpQueryOnlyID)); err != nil { + return + } + switch len(ids) { + case 1: + id = ids[0] + case 0: + err = &NotFoundError{batchimageitem.Label} + default: + err = &NotSingularError{batchimageitem.Label} + } + return +} + +// OnlyIDX is like OnlyID, but panics if an error occurs. +func (_q *BatchImageItemQuery) OnlyIDX(ctx context.Context) int64 { + id, err := _q.OnlyID(ctx) + if err != nil { + panic(err) + } + return id +} + +// All executes the query and returns a list of BatchImageItems. +func (_q *BatchImageItemQuery) All(ctx context.Context) ([]*BatchImageItem, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryAll) + if err := _q.prepareQuery(ctx); err != nil { + return nil, err + } + qr := querierAll[[]*BatchImageItem, *BatchImageItemQuery]() + return withInterceptors[[]*BatchImageItem](ctx, _q, qr, _q.inters) +} + +// AllX is like All, but panics if an error occurs. +func (_q *BatchImageItemQuery) AllX(ctx context.Context) []*BatchImageItem { + nodes, err := _q.All(ctx) + if err != nil { + panic(err) + } + return nodes +} + +// IDs executes the query and returns a list of BatchImageItem IDs. +func (_q *BatchImageItemQuery) IDs(ctx context.Context) (ids []int64, err error) { + if _q.ctx.Unique == nil && _q.path != nil { + _q.Unique(true) + } + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryIDs) + if err = _q.Select(batchimageitem.FieldID).Scan(ctx, &ids); err != nil { + return nil, err + } + return ids, nil +} + +// IDsX is like IDs, but panics if an error occurs. +func (_q *BatchImageItemQuery) IDsX(ctx context.Context) []int64 { + ids, err := _q.IDs(ctx) + if err != nil { + panic(err) + } + return ids +} + +// Count returns the count of the given query. +func (_q *BatchImageItemQuery) Count(ctx context.Context) (int, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryCount) + if err := _q.prepareQuery(ctx); err != nil { + return 0, err + } + return withInterceptors[int](ctx, _q, querierCount[*BatchImageItemQuery](), _q.inters) +} + +// CountX is like Count, but panics if an error occurs. +func (_q *BatchImageItemQuery) CountX(ctx context.Context) int { + count, err := _q.Count(ctx) + if err != nil { + panic(err) + } + return count +} + +// Exist returns true if the query has elements in the graph. +func (_q *BatchImageItemQuery) Exist(ctx context.Context) (bool, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryExist) + switch _, err := _q.FirstID(ctx); { + case IsNotFound(err): + return false, nil + case err != nil: + return false, fmt.Errorf("ent: check existence: %w", err) + default: + return true, nil + } +} + +// ExistX is like Exist, but panics if an error occurs. +func (_q *BatchImageItemQuery) ExistX(ctx context.Context) bool { + exist, err := _q.Exist(ctx) + if err != nil { + panic(err) + } + return exist +} + +// Clone returns a duplicate of the BatchImageItemQuery builder, including all associated steps. It can be +// used to prepare common query builders and use them differently after the clone is made. +func (_q *BatchImageItemQuery) Clone() *BatchImageItemQuery { + if _q == nil { + return nil + } + return &BatchImageItemQuery{ + config: _q.config, + ctx: _q.ctx.Clone(), + order: append([]batchimageitem.OrderOption{}, _q.order...), + inters: append([]Interceptor{}, _q.inters...), + predicates: append([]predicate.BatchImageItem{}, _q.predicates...), + // clone intermediate query. + sql: _q.sql.Clone(), + path: _q.path, + } +} + +// GroupBy is used to group vertices by one or more fields/columns. +// It is often used with aggregate functions, like: count, max, mean, min, sum. +// +// Example: +// +// var v []struct { +// JobID string `json:"job_id,omitempty"` +// Count int `json:"count,omitempty"` +// } +// +// client.BatchImageItem.Query(). +// GroupBy(batchimageitem.FieldJobID). +// Aggregate(ent.Count()). +// Scan(ctx, &v) +func (_q *BatchImageItemQuery) GroupBy(field string, fields ...string) *BatchImageItemGroupBy { + _q.ctx.Fields = append([]string{field}, fields...) + grbuild := &BatchImageItemGroupBy{build: _q} + grbuild.flds = &_q.ctx.Fields + grbuild.label = batchimageitem.Label + grbuild.scan = grbuild.Scan + return grbuild +} + +// Select allows the selection one or more fields/columns for the given query, +// instead of selecting all fields in the entity. +// +// Example: +// +// var v []struct { +// JobID string `json:"job_id,omitempty"` +// } +// +// client.BatchImageItem.Query(). +// Select(batchimageitem.FieldJobID). +// Scan(ctx, &v) +func (_q *BatchImageItemQuery) Select(fields ...string) *BatchImageItemSelect { + _q.ctx.Fields = append(_q.ctx.Fields, fields...) + sbuild := &BatchImageItemSelect{BatchImageItemQuery: _q} + sbuild.label = batchimageitem.Label + sbuild.flds, sbuild.scan = &_q.ctx.Fields, sbuild.Scan + return sbuild +} + +// Aggregate returns a BatchImageItemSelect configured with the given aggregations. +func (_q *BatchImageItemQuery) Aggregate(fns ...AggregateFunc) *BatchImageItemSelect { + return _q.Select().Aggregate(fns...) +} + +func (_q *BatchImageItemQuery) prepareQuery(ctx context.Context) error { + for _, inter := range _q.inters { + if inter == nil { + return fmt.Errorf("ent: uninitialized interceptor (forgotten import ent/runtime?)") + } + if trv, ok := inter.(Traverser); ok { + if err := trv.Traverse(ctx, _q); err != nil { + return err + } + } + } + for _, f := range _q.ctx.Fields { + if !batchimageitem.ValidColumn(f) { + return &ValidationError{Name: f, err: fmt.Errorf("ent: invalid field %q for query", f)} + } + } + if _q.path != nil { + prev, err := _q.path(ctx) + if err != nil { + return err + } + _q.sql = prev + } + return nil +} + +func (_q *BatchImageItemQuery) sqlAll(ctx context.Context, hooks ...queryHook) ([]*BatchImageItem, error) { + var ( + nodes = []*BatchImageItem{} + _spec = _q.querySpec() + ) + _spec.ScanValues = func(columns []string) ([]any, error) { + return (*BatchImageItem).scanValues(nil, columns) + } + _spec.Assign = func(columns []string, values []any) error { + node := &BatchImageItem{config: _q.config} + nodes = append(nodes, node) + return node.assignValues(columns, values) + } + if len(_q.modifiers) > 0 { + _spec.Modifiers = _q.modifiers + } + for i := range hooks { + hooks[i](ctx, _spec) + } + if err := sqlgraph.QueryNodes(ctx, _q.driver, _spec); err != nil { + return nil, err + } + if len(nodes) == 0 { + return nodes, nil + } + return nodes, nil +} + +func (_q *BatchImageItemQuery) sqlCount(ctx context.Context) (int, error) { + _spec := _q.querySpec() + if len(_q.modifiers) > 0 { + _spec.Modifiers = _q.modifiers + } + _spec.Node.Columns = _q.ctx.Fields + if len(_q.ctx.Fields) > 0 { + _spec.Unique = _q.ctx.Unique != nil && *_q.ctx.Unique + } + return sqlgraph.CountNodes(ctx, _q.driver, _spec) +} + +func (_q *BatchImageItemQuery) querySpec() *sqlgraph.QuerySpec { + _spec := sqlgraph.NewQuerySpec(batchimageitem.Table, batchimageitem.Columns, sqlgraph.NewFieldSpec(batchimageitem.FieldID, field.TypeInt64)) + _spec.From = _q.sql + if unique := _q.ctx.Unique; unique != nil { + _spec.Unique = *unique + } else if _q.path != nil { + _spec.Unique = true + } + if fields := _q.ctx.Fields; len(fields) > 0 { + _spec.Node.Columns = make([]string, 0, len(fields)) + _spec.Node.Columns = append(_spec.Node.Columns, batchimageitem.FieldID) + for i := range fields { + if fields[i] != batchimageitem.FieldID { + _spec.Node.Columns = append(_spec.Node.Columns, fields[i]) + } + } + } + if ps := _q.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if limit := _q.ctx.Limit; limit != nil { + _spec.Limit = *limit + } + if offset := _q.ctx.Offset; offset != nil { + _spec.Offset = *offset + } + if ps := _q.order; len(ps) > 0 { + _spec.Order = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + return _spec +} + +func (_q *BatchImageItemQuery) sqlQuery(ctx context.Context) *sql.Selector { + builder := sql.Dialect(_q.driver.Dialect()) + t1 := builder.Table(batchimageitem.Table) + columns := _q.ctx.Fields + if len(columns) == 0 { + columns = batchimageitem.Columns + } + selector := builder.Select(t1.Columns(columns...)...).From(t1) + if _q.sql != nil { + selector = _q.sql + selector.Select(selector.Columns(columns...)...) + } + if _q.ctx.Unique != nil && *_q.ctx.Unique { + selector.Distinct() + } + for _, m := range _q.modifiers { + m(selector) + } + for _, p := range _q.predicates { + p(selector) + } + for _, p := range _q.order { + p(selector) + } + if offset := _q.ctx.Offset; offset != nil { + // limit is mandatory for offset clause. We start + // with default value, and override it below if needed. + selector.Offset(*offset).Limit(math.MaxInt32) + } + if limit := _q.ctx.Limit; limit != nil { + selector.Limit(*limit) + } + return selector +} + +// ForUpdate locks the selected rows against concurrent updates, and prevent them from being +// updated, deleted or "selected ... for update" by other sessions, until the transaction is +// either committed or rolled-back. +func (_q *BatchImageItemQuery) ForUpdate(opts ...sql.LockOption) *BatchImageItemQuery { + if _q.driver.Dialect() == dialect.Postgres { + _q.Unique(false) + } + _q.modifiers = append(_q.modifiers, func(s *sql.Selector) { + s.ForUpdate(opts...) + }) + return _q +} + +// ForShare behaves similarly to ForUpdate, except that it acquires a shared mode lock +// on any rows that are read. Other sessions can read the rows, but cannot modify them +// until your transaction commits. +func (_q *BatchImageItemQuery) ForShare(opts ...sql.LockOption) *BatchImageItemQuery { + if _q.driver.Dialect() == dialect.Postgres { + _q.Unique(false) + } + _q.modifiers = append(_q.modifiers, func(s *sql.Selector) { + s.ForShare(opts...) + }) + return _q +} + +// BatchImageItemGroupBy is the group-by builder for BatchImageItem entities. +type BatchImageItemGroupBy struct { + selector + build *BatchImageItemQuery +} + +// Aggregate adds the given aggregation functions to the group-by query. +func (_g *BatchImageItemGroupBy) Aggregate(fns ...AggregateFunc) *BatchImageItemGroupBy { + _g.fns = append(_g.fns, fns...) + return _g +} + +// Scan applies the selector query and scans the result into the given value. +func (_g *BatchImageItemGroupBy) Scan(ctx context.Context, v any) error { + ctx = setContextOp(ctx, _g.build.ctx, ent.OpQueryGroupBy) + if err := _g.build.prepareQuery(ctx); err != nil { + return err + } + return scanWithInterceptors[*BatchImageItemQuery, *BatchImageItemGroupBy](ctx, _g.build, _g, _g.build.inters, v) +} + +func (_g *BatchImageItemGroupBy) sqlScan(ctx context.Context, root *BatchImageItemQuery, v any) error { + selector := root.sqlQuery(ctx).Select() + aggregation := make([]string, 0, len(_g.fns)) + for _, fn := range _g.fns { + aggregation = append(aggregation, fn(selector)) + } + if len(selector.SelectedColumns()) == 0 { + columns := make([]string, 0, len(*_g.flds)+len(_g.fns)) + for _, f := range *_g.flds { + columns = append(columns, selector.C(f)) + } + columns = append(columns, aggregation...) + selector.Select(columns...) + } + selector.GroupBy(selector.Columns(*_g.flds...)...) + if err := selector.Err(); err != nil { + return err + } + rows := &sql.Rows{} + query, args := selector.Query() + if err := _g.build.driver.Query(ctx, query, args, rows); err != nil { + return err + } + defer rows.Close() + return sql.ScanSlice(rows, v) +} + +// BatchImageItemSelect is the builder for selecting fields of BatchImageItem entities. +type BatchImageItemSelect struct { + *BatchImageItemQuery + selector +} + +// Aggregate adds the given aggregation functions to the selector query. +func (_s *BatchImageItemSelect) Aggregate(fns ...AggregateFunc) *BatchImageItemSelect { + _s.fns = append(_s.fns, fns...) + return _s +} + +// Scan applies the selector query and scans the result into the given value. +func (_s *BatchImageItemSelect) Scan(ctx context.Context, v any) error { + ctx = setContextOp(ctx, _s.ctx, ent.OpQuerySelect) + if err := _s.prepareQuery(ctx); err != nil { + return err + } + return scanWithInterceptors[*BatchImageItemQuery, *BatchImageItemSelect](ctx, _s.BatchImageItemQuery, _s, _s.inters, v) +} + +func (_s *BatchImageItemSelect) sqlScan(ctx context.Context, root *BatchImageItemQuery, v any) error { + selector := root.sqlQuery(ctx) + aggregation := make([]string, 0, len(_s.fns)) + for _, fn := range _s.fns { + aggregation = append(aggregation, fn(selector)) + } + switch n := len(*_s.selector.flds); { + case n == 0 && len(aggregation) > 0: + selector.Select(aggregation...) + case n != 0 && len(aggregation) > 0: + selector.AppendSelect(aggregation...) + } + rows := &sql.Rows{} + query, args := selector.Query() + if err := _s.driver.Query(ctx, query, args, rows); err != nil { + return err + } + defer rows.Close() + return sql.ScanSlice(rows, v) +} diff --git a/backend/ent/batchimageitem_update.go b/backend/ent/batchimageitem_update.go new file mode 100644 index 0000000000..edca025b7e --- /dev/null +++ b/backend/ent/batchimageitem_update.go @@ -0,0 +1,1132 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "errors" + "fmt" + "time" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageItemUpdate is the builder for updating BatchImageItem entities. +type BatchImageItemUpdate struct { + config + hooks []Hook + mutation *BatchImageItemMutation +} + +// Where appends a list predicates to the BatchImageItemUpdate builder. +func (_u *BatchImageItemUpdate) Where(ps ...predicate.BatchImageItem) *BatchImageItemUpdate { + _u.mutation.Where(ps...) + return _u +} + +// SetJobID sets the "job_id" field. +func (_u *BatchImageItemUpdate) SetJobID(v string) *BatchImageItemUpdate { + _u.mutation.SetJobID(v) + return _u +} + +// SetNillableJobID sets the "job_id" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableJobID(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetJobID(*v) + } + return _u +} + +// SetCustomID sets the "custom_id" field. +func (_u *BatchImageItemUpdate) SetCustomID(v string) *BatchImageItemUpdate { + _u.mutation.SetCustomID(v) + return _u +} + +// SetNillableCustomID sets the "custom_id" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableCustomID(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetCustomID(*v) + } + return _u +} + +// SetStatus sets the "status" field. +func (_u *BatchImageItemUpdate) SetStatus(v string) *BatchImageItemUpdate { + _u.mutation.SetStatus(v) + return _u +} + +// SetNillableStatus sets the "status" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableStatus(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetStatus(*v) + } + return _u +} + +// SetRequestHash sets the "request_hash" field. +func (_u *BatchImageItemUpdate) SetRequestHash(v string) *BatchImageItemUpdate { + _u.mutation.SetRequestHash(v) + return _u +} + +// SetNillableRequestHash sets the "request_hash" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableRequestHash(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetRequestHash(*v) + } + return _u +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (_u *BatchImageItemUpdate) ClearRequestHash() *BatchImageItemUpdate { + _u.mutation.ClearRequestHash() + return _u +} + +// SetPromptPreview sets the "prompt_preview" field. +func (_u *BatchImageItemUpdate) SetPromptPreview(v string) *BatchImageItemUpdate { + _u.mutation.SetPromptPreview(v) + return _u +} + +// SetNillablePromptPreview sets the "prompt_preview" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillablePromptPreview(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetPromptPreview(*v) + } + return _u +} + +// ClearPromptPreview clears the value of the "prompt_preview" field. +func (_u *BatchImageItemUpdate) ClearPromptPreview() *BatchImageItemUpdate { + _u.mutation.ClearPromptPreview() + return _u +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (_u *BatchImageItemUpdate) SetProviderSourceObject(v string) *BatchImageItemUpdate { + _u.mutation.SetProviderSourceObject(v) + return _u +} + +// SetNillableProviderSourceObject sets the "provider_source_object" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableProviderSourceObject(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetProviderSourceObject(*v) + } + return _u +} + +// ClearProviderSourceObject clears the value of the "provider_source_object" field. +func (_u *BatchImageItemUpdate) ClearProviderSourceObject() *BatchImageItemUpdate { + _u.mutation.ClearProviderSourceObject() + return _u +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (_u *BatchImageItemUpdate) SetSourceLineNumber(v int) *BatchImageItemUpdate { + _u.mutation.ResetSourceLineNumber() + _u.mutation.SetSourceLineNumber(v) + return _u +} + +// SetNillableSourceLineNumber sets the "source_line_number" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableSourceLineNumber(v *int) *BatchImageItemUpdate { + if v != nil { + _u.SetSourceLineNumber(*v) + } + return _u +} + +// AddSourceLineNumber adds value to the "source_line_number" field. +func (_u *BatchImageItemUpdate) AddSourceLineNumber(v int) *BatchImageItemUpdate { + _u.mutation.AddSourceLineNumber(v) + return _u +} + +// ClearSourceLineNumber clears the value of the "source_line_number" field. +func (_u *BatchImageItemUpdate) ClearSourceLineNumber() *BatchImageItemUpdate { + _u.mutation.ClearSourceLineNumber() + return _u +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (_u *BatchImageItemUpdate) SetSourceByteOffset(v int64) *BatchImageItemUpdate { + _u.mutation.ResetSourceByteOffset() + _u.mutation.SetSourceByteOffset(v) + return _u +} + +// SetNillableSourceByteOffset sets the "source_byte_offset" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableSourceByteOffset(v *int64) *BatchImageItemUpdate { + if v != nil { + _u.SetSourceByteOffset(*v) + } + return _u +} + +// AddSourceByteOffset adds value to the "source_byte_offset" field. +func (_u *BatchImageItemUpdate) AddSourceByteOffset(v int64) *BatchImageItemUpdate { + _u.mutation.AddSourceByteOffset(v) + return _u +} + +// ClearSourceByteOffset clears the value of the "source_byte_offset" field. +func (_u *BatchImageItemUpdate) ClearSourceByteOffset() *BatchImageItemUpdate { + _u.mutation.ClearSourceByteOffset() + return _u +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (_u *BatchImageItemUpdate) SetSourceByteLength(v int64) *BatchImageItemUpdate { + _u.mutation.ResetSourceByteLength() + _u.mutation.SetSourceByteLength(v) + return _u +} + +// SetNillableSourceByteLength sets the "source_byte_length" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableSourceByteLength(v *int64) *BatchImageItemUpdate { + if v != nil { + _u.SetSourceByteLength(*v) + } + return _u +} + +// AddSourceByteLength adds value to the "source_byte_length" field. +func (_u *BatchImageItemUpdate) AddSourceByteLength(v int64) *BatchImageItemUpdate { + _u.mutation.AddSourceByteLength(v) + return _u +} + +// ClearSourceByteLength clears the value of the "source_byte_length" field. +func (_u *BatchImageItemUpdate) ClearSourceByteLength() *BatchImageItemUpdate { + _u.mutation.ClearSourceByteLength() + return _u +} + +// SetMimeType sets the "mime_type" field. +func (_u *BatchImageItemUpdate) SetMimeType(v string) *BatchImageItemUpdate { + _u.mutation.SetMimeType(v) + return _u +} + +// SetNillableMimeType sets the "mime_type" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableMimeType(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetMimeType(*v) + } + return _u +} + +// ClearMimeType clears the value of the "mime_type" field. +func (_u *BatchImageItemUpdate) ClearMimeType() *BatchImageItemUpdate { + _u.mutation.ClearMimeType() + return _u +} + +// SetFileExtension sets the "file_extension" field. +func (_u *BatchImageItemUpdate) SetFileExtension(v string) *BatchImageItemUpdate { + _u.mutation.SetFileExtension(v) + return _u +} + +// SetNillableFileExtension sets the "file_extension" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableFileExtension(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetFileExtension(*v) + } + return _u +} + +// ClearFileExtension clears the value of the "file_extension" field. +func (_u *BatchImageItemUpdate) ClearFileExtension() *BatchImageItemUpdate { + _u.mutation.ClearFileExtension() + return _u +} + +// SetImageCount sets the "image_count" field. +func (_u *BatchImageItemUpdate) SetImageCount(v int) *BatchImageItemUpdate { + _u.mutation.ResetImageCount() + _u.mutation.SetImageCount(v) + return _u +} + +// SetNillableImageCount sets the "image_count" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableImageCount(v *int) *BatchImageItemUpdate { + if v != nil { + _u.SetImageCount(*v) + } + return _u +} + +// AddImageCount adds value to the "image_count" field. +func (_u *BatchImageItemUpdate) AddImageCount(v int) *BatchImageItemUpdate { + _u.mutation.AddImageCount(v) + return _u +} + +// SetErrorCode sets the "error_code" field. +func (_u *BatchImageItemUpdate) SetErrorCode(v string) *BatchImageItemUpdate { + _u.mutation.SetErrorCode(v) + return _u +} + +// SetNillableErrorCode sets the "error_code" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableErrorCode(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetErrorCode(*v) + } + return _u +} + +// ClearErrorCode clears the value of the "error_code" field. +func (_u *BatchImageItemUpdate) ClearErrorCode() *BatchImageItemUpdate { + _u.mutation.ClearErrorCode() + return _u +} + +// SetErrorMessage sets the "error_message" field. +func (_u *BatchImageItemUpdate) SetErrorMessage(v string) *BatchImageItemUpdate { + _u.mutation.SetErrorMessage(v) + return _u +} + +// SetNillableErrorMessage sets the "error_message" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableErrorMessage(v *string) *BatchImageItemUpdate { + if v != nil { + _u.SetErrorMessage(*v) + } + return _u +} + +// ClearErrorMessage clears the value of the "error_message" field. +func (_u *BatchImageItemUpdate) ClearErrorMessage() *BatchImageItemUpdate { + _u.mutation.ClearErrorMessage() + return _u +} + +// SetBilledAmount sets the "billed_amount" field. +func (_u *BatchImageItemUpdate) SetBilledAmount(v float64) *BatchImageItemUpdate { + _u.mutation.ResetBilledAmount() + _u.mutation.SetBilledAmount(v) + return _u +} + +// SetNillableBilledAmount sets the "billed_amount" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableBilledAmount(v *float64) *BatchImageItemUpdate { + if v != nil { + _u.SetBilledAmount(*v) + } + return _u +} + +// AddBilledAmount adds value to the "billed_amount" field. +func (_u *BatchImageItemUpdate) AddBilledAmount(v float64) *BatchImageItemUpdate { + _u.mutation.AddBilledAmount(v) + return _u +} + +// ClearBilledAmount clears the value of the "billed_amount" field. +func (_u *BatchImageItemUpdate) ClearBilledAmount() *BatchImageItemUpdate { + _u.mutation.ClearBilledAmount() + return _u +} + +// SetIndexedAt sets the "indexed_at" field. +func (_u *BatchImageItemUpdate) SetIndexedAt(v time.Time) *BatchImageItemUpdate { + _u.mutation.SetIndexedAt(v) + return _u +} + +// SetNillableIndexedAt sets the "indexed_at" field if the given value is not nil. +func (_u *BatchImageItemUpdate) SetNillableIndexedAt(v *time.Time) *BatchImageItemUpdate { + if v != nil { + _u.SetIndexedAt(*v) + } + return _u +} + +// ClearIndexedAt clears the value of the "indexed_at" field. +func (_u *BatchImageItemUpdate) ClearIndexedAt() *BatchImageItemUpdate { + _u.mutation.ClearIndexedAt() + return _u +} + +// Mutation returns the BatchImageItemMutation object of the builder. +func (_u *BatchImageItemUpdate) Mutation() *BatchImageItemMutation { + return _u.mutation +} + +// Save executes the query and returns the number of nodes affected by the update operation. +func (_u *BatchImageItemUpdate) Save(ctx context.Context) (int, error) { + return withHooks(ctx, _u.sqlSave, _u.mutation, _u.hooks) +} + +// SaveX is like Save, but panics if an error occurs. +func (_u *BatchImageItemUpdate) SaveX(ctx context.Context) int { + affected, err := _u.Save(ctx) + if err != nil { + panic(err) + } + return affected +} + +// Exec executes the query. +func (_u *BatchImageItemUpdate) Exec(ctx context.Context) error { + _, err := _u.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_u *BatchImageItemUpdate) ExecX(ctx context.Context) { + if err := _u.Exec(ctx); err != nil { + panic(err) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_u *BatchImageItemUpdate) check() error { + if v, ok := _u.mutation.JobID(); ok { + if err := batchimageitem.JobIDValidator(v); err != nil { + return &ValidationError{Name: "job_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.job_id": %w`, err)} + } + } + if v, ok := _u.mutation.CustomID(); ok { + if err := batchimageitem.CustomIDValidator(v); err != nil { + return &ValidationError{Name: "custom_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.custom_id": %w`, err)} + } + } + if v, ok := _u.mutation.Status(); ok { + if err := batchimageitem.StatusValidator(v); err != nil { + return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.status": %w`, err)} + } + } + if v, ok := _u.mutation.RequestHash(); ok { + if err := batchimageitem.RequestHashValidator(v); err != nil { + return &ValidationError{Name: "request_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.request_hash": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderSourceObject(); ok { + if err := batchimageitem.ProviderSourceObjectValidator(v); err != nil { + return &ValidationError{Name: "provider_source_object", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.provider_source_object": %w`, err)} + } + } + if v, ok := _u.mutation.MimeType(); ok { + if err := batchimageitem.MimeTypeValidator(v); err != nil { + return &ValidationError{Name: "mime_type", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.mime_type": %w`, err)} + } + } + if v, ok := _u.mutation.FileExtension(); ok { + if err := batchimageitem.FileExtensionValidator(v); err != nil { + return &ValidationError{Name: "file_extension", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.file_extension": %w`, err)} + } + } + if v, ok := _u.mutation.ErrorCode(); ok { + if err := batchimageitem.ErrorCodeValidator(v); err != nil { + return &ValidationError{Name: "error_code", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.error_code": %w`, err)} + } + } + return nil +} + +func (_u *BatchImageItemUpdate) sqlSave(ctx context.Context) (_node int, err error) { + if err := _u.check(); err != nil { + return _node, err + } + _spec := sqlgraph.NewUpdateSpec(batchimageitem.Table, batchimageitem.Columns, sqlgraph.NewFieldSpec(batchimageitem.FieldID, field.TypeInt64)) + if ps := _u.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if value, ok := _u.mutation.JobID(); ok { + _spec.SetField(batchimageitem.FieldJobID, field.TypeString, value) + } + if value, ok := _u.mutation.CustomID(); ok { + _spec.SetField(batchimageitem.FieldCustomID, field.TypeString, value) + } + if value, ok := _u.mutation.Status(); ok { + _spec.SetField(batchimageitem.FieldStatus, field.TypeString, value) + } + if value, ok := _u.mutation.RequestHash(); ok { + _spec.SetField(batchimageitem.FieldRequestHash, field.TypeString, value) + } + if _u.mutation.RequestHashCleared() { + _spec.ClearField(batchimageitem.FieldRequestHash, field.TypeString) + } + if value, ok := _u.mutation.PromptPreview(); ok { + _spec.SetField(batchimageitem.FieldPromptPreview, field.TypeString, value) + } + if _u.mutation.PromptPreviewCleared() { + _spec.ClearField(batchimageitem.FieldPromptPreview, field.TypeString) + } + if value, ok := _u.mutation.ProviderSourceObject(); ok { + _spec.SetField(batchimageitem.FieldProviderSourceObject, field.TypeString, value) + } + if _u.mutation.ProviderSourceObjectCleared() { + _spec.ClearField(batchimageitem.FieldProviderSourceObject, field.TypeString) + } + if value, ok := _u.mutation.SourceLineNumber(); ok { + _spec.SetField(batchimageitem.FieldSourceLineNumber, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedSourceLineNumber(); ok { + _spec.AddField(batchimageitem.FieldSourceLineNumber, field.TypeInt, value) + } + if _u.mutation.SourceLineNumberCleared() { + _spec.ClearField(batchimageitem.FieldSourceLineNumber, field.TypeInt) + } + if value, ok := _u.mutation.SourceByteOffset(); ok { + _spec.SetField(batchimageitem.FieldSourceByteOffset, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedSourceByteOffset(); ok { + _spec.AddField(batchimageitem.FieldSourceByteOffset, field.TypeInt64, value) + } + if _u.mutation.SourceByteOffsetCleared() { + _spec.ClearField(batchimageitem.FieldSourceByteOffset, field.TypeInt64) + } + if value, ok := _u.mutation.SourceByteLength(); ok { + _spec.SetField(batchimageitem.FieldSourceByteLength, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedSourceByteLength(); ok { + _spec.AddField(batchimageitem.FieldSourceByteLength, field.TypeInt64, value) + } + if _u.mutation.SourceByteLengthCleared() { + _spec.ClearField(batchimageitem.FieldSourceByteLength, field.TypeInt64) + } + if value, ok := _u.mutation.MimeType(); ok { + _spec.SetField(batchimageitem.FieldMimeType, field.TypeString, value) + } + if _u.mutation.MimeTypeCleared() { + _spec.ClearField(batchimageitem.FieldMimeType, field.TypeString) + } + if value, ok := _u.mutation.FileExtension(); ok { + _spec.SetField(batchimageitem.FieldFileExtension, field.TypeString, value) + } + if _u.mutation.FileExtensionCleared() { + _spec.ClearField(batchimageitem.FieldFileExtension, field.TypeString) + } + if value, ok := _u.mutation.ImageCount(); ok { + _spec.SetField(batchimageitem.FieldImageCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedImageCount(); ok { + _spec.AddField(batchimageitem.FieldImageCount, field.TypeInt, value) + } + if value, ok := _u.mutation.ErrorCode(); ok { + _spec.SetField(batchimageitem.FieldErrorCode, field.TypeString, value) + } + if _u.mutation.ErrorCodeCleared() { + _spec.ClearField(batchimageitem.FieldErrorCode, field.TypeString) + } + if value, ok := _u.mutation.ErrorMessage(); ok { + _spec.SetField(batchimageitem.FieldErrorMessage, field.TypeString, value) + } + if _u.mutation.ErrorMessageCleared() { + _spec.ClearField(batchimageitem.FieldErrorMessage, field.TypeString) + } + if value, ok := _u.mutation.BilledAmount(); ok { + _spec.SetField(batchimageitem.FieldBilledAmount, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedBilledAmount(); ok { + _spec.AddField(batchimageitem.FieldBilledAmount, field.TypeFloat64, value) + } + if _u.mutation.BilledAmountCleared() { + _spec.ClearField(batchimageitem.FieldBilledAmount, field.TypeFloat64) + } + if value, ok := _u.mutation.IndexedAt(); ok { + _spec.SetField(batchimageitem.FieldIndexedAt, field.TypeTime, value) + } + if _u.mutation.IndexedAtCleared() { + _spec.ClearField(batchimageitem.FieldIndexedAt, field.TypeTime) + } + if _node, err = sqlgraph.UpdateNodes(ctx, _u.driver, _spec); err != nil { + if _, ok := err.(*sqlgraph.NotFoundError); ok { + err = &NotFoundError{batchimageitem.Label} + } else if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return 0, err + } + _u.mutation.done = true + return _node, nil +} + +// BatchImageItemUpdateOne is the builder for updating a single BatchImageItem entity. +type BatchImageItemUpdateOne struct { + config + fields []string + hooks []Hook + mutation *BatchImageItemMutation +} + +// SetJobID sets the "job_id" field. +func (_u *BatchImageItemUpdateOne) SetJobID(v string) *BatchImageItemUpdateOne { + _u.mutation.SetJobID(v) + return _u +} + +// SetNillableJobID sets the "job_id" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableJobID(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetJobID(*v) + } + return _u +} + +// SetCustomID sets the "custom_id" field. +func (_u *BatchImageItemUpdateOne) SetCustomID(v string) *BatchImageItemUpdateOne { + _u.mutation.SetCustomID(v) + return _u +} + +// SetNillableCustomID sets the "custom_id" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableCustomID(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetCustomID(*v) + } + return _u +} + +// SetStatus sets the "status" field. +func (_u *BatchImageItemUpdateOne) SetStatus(v string) *BatchImageItemUpdateOne { + _u.mutation.SetStatus(v) + return _u +} + +// SetNillableStatus sets the "status" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableStatus(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetStatus(*v) + } + return _u +} + +// SetRequestHash sets the "request_hash" field. +func (_u *BatchImageItemUpdateOne) SetRequestHash(v string) *BatchImageItemUpdateOne { + _u.mutation.SetRequestHash(v) + return _u +} + +// SetNillableRequestHash sets the "request_hash" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableRequestHash(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetRequestHash(*v) + } + return _u +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (_u *BatchImageItemUpdateOne) ClearRequestHash() *BatchImageItemUpdateOne { + _u.mutation.ClearRequestHash() + return _u +} + +// SetPromptPreview sets the "prompt_preview" field. +func (_u *BatchImageItemUpdateOne) SetPromptPreview(v string) *BatchImageItemUpdateOne { + _u.mutation.SetPromptPreview(v) + return _u +} + +// SetNillablePromptPreview sets the "prompt_preview" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillablePromptPreview(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetPromptPreview(*v) + } + return _u +} + +// ClearPromptPreview clears the value of the "prompt_preview" field. +func (_u *BatchImageItemUpdateOne) ClearPromptPreview() *BatchImageItemUpdateOne { + _u.mutation.ClearPromptPreview() + return _u +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (_u *BatchImageItemUpdateOne) SetProviderSourceObject(v string) *BatchImageItemUpdateOne { + _u.mutation.SetProviderSourceObject(v) + return _u +} + +// SetNillableProviderSourceObject sets the "provider_source_object" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableProviderSourceObject(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetProviderSourceObject(*v) + } + return _u +} + +// ClearProviderSourceObject clears the value of the "provider_source_object" field. +func (_u *BatchImageItemUpdateOne) ClearProviderSourceObject() *BatchImageItemUpdateOne { + _u.mutation.ClearProviderSourceObject() + return _u +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (_u *BatchImageItemUpdateOne) SetSourceLineNumber(v int) *BatchImageItemUpdateOne { + _u.mutation.ResetSourceLineNumber() + _u.mutation.SetSourceLineNumber(v) + return _u +} + +// SetNillableSourceLineNumber sets the "source_line_number" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableSourceLineNumber(v *int) *BatchImageItemUpdateOne { + if v != nil { + _u.SetSourceLineNumber(*v) + } + return _u +} + +// AddSourceLineNumber adds value to the "source_line_number" field. +func (_u *BatchImageItemUpdateOne) AddSourceLineNumber(v int) *BatchImageItemUpdateOne { + _u.mutation.AddSourceLineNumber(v) + return _u +} + +// ClearSourceLineNumber clears the value of the "source_line_number" field. +func (_u *BatchImageItemUpdateOne) ClearSourceLineNumber() *BatchImageItemUpdateOne { + _u.mutation.ClearSourceLineNumber() + return _u +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (_u *BatchImageItemUpdateOne) SetSourceByteOffset(v int64) *BatchImageItemUpdateOne { + _u.mutation.ResetSourceByteOffset() + _u.mutation.SetSourceByteOffset(v) + return _u +} + +// SetNillableSourceByteOffset sets the "source_byte_offset" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableSourceByteOffset(v *int64) *BatchImageItemUpdateOne { + if v != nil { + _u.SetSourceByteOffset(*v) + } + return _u +} + +// AddSourceByteOffset adds value to the "source_byte_offset" field. +func (_u *BatchImageItemUpdateOne) AddSourceByteOffset(v int64) *BatchImageItemUpdateOne { + _u.mutation.AddSourceByteOffset(v) + return _u +} + +// ClearSourceByteOffset clears the value of the "source_byte_offset" field. +func (_u *BatchImageItemUpdateOne) ClearSourceByteOffset() *BatchImageItemUpdateOne { + _u.mutation.ClearSourceByteOffset() + return _u +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (_u *BatchImageItemUpdateOne) SetSourceByteLength(v int64) *BatchImageItemUpdateOne { + _u.mutation.ResetSourceByteLength() + _u.mutation.SetSourceByteLength(v) + return _u +} + +// SetNillableSourceByteLength sets the "source_byte_length" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableSourceByteLength(v *int64) *BatchImageItemUpdateOne { + if v != nil { + _u.SetSourceByteLength(*v) + } + return _u +} + +// AddSourceByteLength adds value to the "source_byte_length" field. +func (_u *BatchImageItemUpdateOne) AddSourceByteLength(v int64) *BatchImageItemUpdateOne { + _u.mutation.AddSourceByteLength(v) + return _u +} + +// ClearSourceByteLength clears the value of the "source_byte_length" field. +func (_u *BatchImageItemUpdateOne) ClearSourceByteLength() *BatchImageItemUpdateOne { + _u.mutation.ClearSourceByteLength() + return _u +} + +// SetMimeType sets the "mime_type" field. +func (_u *BatchImageItemUpdateOne) SetMimeType(v string) *BatchImageItemUpdateOne { + _u.mutation.SetMimeType(v) + return _u +} + +// SetNillableMimeType sets the "mime_type" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableMimeType(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetMimeType(*v) + } + return _u +} + +// ClearMimeType clears the value of the "mime_type" field. +func (_u *BatchImageItemUpdateOne) ClearMimeType() *BatchImageItemUpdateOne { + _u.mutation.ClearMimeType() + return _u +} + +// SetFileExtension sets the "file_extension" field. +func (_u *BatchImageItemUpdateOne) SetFileExtension(v string) *BatchImageItemUpdateOne { + _u.mutation.SetFileExtension(v) + return _u +} + +// SetNillableFileExtension sets the "file_extension" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableFileExtension(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetFileExtension(*v) + } + return _u +} + +// ClearFileExtension clears the value of the "file_extension" field. +func (_u *BatchImageItemUpdateOne) ClearFileExtension() *BatchImageItemUpdateOne { + _u.mutation.ClearFileExtension() + return _u +} + +// SetImageCount sets the "image_count" field. +func (_u *BatchImageItemUpdateOne) SetImageCount(v int) *BatchImageItemUpdateOne { + _u.mutation.ResetImageCount() + _u.mutation.SetImageCount(v) + return _u +} + +// SetNillableImageCount sets the "image_count" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableImageCount(v *int) *BatchImageItemUpdateOne { + if v != nil { + _u.SetImageCount(*v) + } + return _u +} + +// AddImageCount adds value to the "image_count" field. +func (_u *BatchImageItemUpdateOne) AddImageCount(v int) *BatchImageItemUpdateOne { + _u.mutation.AddImageCount(v) + return _u +} + +// SetErrorCode sets the "error_code" field. +func (_u *BatchImageItemUpdateOne) SetErrorCode(v string) *BatchImageItemUpdateOne { + _u.mutation.SetErrorCode(v) + return _u +} + +// SetNillableErrorCode sets the "error_code" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableErrorCode(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetErrorCode(*v) + } + return _u +} + +// ClearErrorCode clears the value of the "error_code" field. +func (_u *BatchImageItemUpdateOne) ClearErrorCode() *BatchImageItemUpdateOne { + _u.mutation.ClearErrorCode() + return _u +} + +// SetErrorMessage sets the "error_message" field. +func (_u *BatchImageItemUpdateOne) SetErrorMessage(v string) *BatchImageItemUpdateOne { + _u.mutation.SetErrorMessage(v) + return _u +} + +// SetNillableErrorMessage sets the "error_message" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableErrorMessage(v *string) *BatchImageItemUpdateOne { + if v != nil { + _u.SetErrorMessage(*v) + } + return _u +} + +// ClearErrorMessage clears the value of the "error_message" field. +func (_u *BatchImageItemUpdateOne) ClearErrorMessage() *BatchImageItemUpdateOne { + _u.mutation.ClearErrorMessage() + return _u +} + +// SetBilledAmount sets the "billed_amount" field. +func (_u *BatchImageItemUpdateOne) SetBilledAmount(v float64) *BatchImageItemUpdateOne { + _u.mutation.ResetBilledAmount() + _u.mutation.SetBilledAmount(v) + return _u +} + +// SetNillableBilledAmount sets the "billed_amount" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableBilledAmount(v *float64) *BatchImageItemUpdateOne { + if v != nil { + _u.SetBilledAmount(*v) + } + return _u +} + +// AddBilledAmount adds value to the "billed_amount" field. +func (_u *BatchImageItemUpdateOne) AddBilledAmount(v float64) *BatchImageItemUpdateOne { + _u.mutation.AddBilledAmount(v) + return _u +} + +// ClearBilledAmount clears the value of the "billed_amount" field. +func (_u *BatchImageItemUpdateOne) ClearBilledAmount() *BatchImageItemUpdateOne { + _u.mutation.ClearBilledAmount() + return _u +} + +// SetIndexedAt sets the "indexed_at" field. +func (_u *BatchImageItemUpdateOne) SetIndexedAt(v time.Time) *BatchImageItemUpdateOne { + _u.mutation.SetIndexedAt(v) + return _u +} + +// SetNillableIndexedAt sets the "indexed_at" field if the given value is not nil. +func (_u *BatchImageItemUpdateOne) SetNillableIndexedAt(v *time.Time) *BatchImageItemUpdateOne { + if v != nil { + _u.SetIndexedAt(*v) + } + return _u +} + +// ClearIndexedAt clears the value of the "indexed_at" field. +func (_u *BatchImageItemUpdateOne) ClearIndexedAt() *BatchImageItemUpdateOne { + _u.mutation.ClearIndexedAt() + return _u +} + +// Mutation returns the BatchImageItemMutation object of the builder. +func (_u *BatchImageItemUpdateOne) Mutation() *BatchImageItemMutation { + return _u.mutation +} + +// Where appends a list predicates to the BatchImageItemUpdate builder. +func (_u *BatchImageItemUpdateOne) Where(ps ...predicate.BatchImageItem) *BatchImageItemUpdateOne { + _u.mutation.Where(ps...) + return _u +} + +// Select allows selecting one or more fields (columns) of the returned entity. +// The default is selecting all fields defined in the entity schema. +func (_u *BatchImageItemUpdateOne) Select(field string, fields ...string) *BatchImageItemUpdateOne { + _u.fields = append([]string{field}, fields...) + return _u +} + +// Save executes the query and returns the updated BatchImageItem entity. +func (_u *BatchImageItemUpdateOne) Save(ctx context.Context) (*BatchImageItem, error) { + return withHooks(ctx, _u.sqlSave, _u.mutation, _u.hooks) +} + +// SaveX is like Save, but panics if an error occurs. +func (_u *BatchImageItemUpdateOne) SaveX(ctx context.Context) *BatchImageItem { + node, err := _u.Save(ctx) + if err != nil { + panic(err) + } + return node +} + +// Exec executes the query on the entity. +func (_u *BatchImageItemUpdateOne) Exec(ctx context.Context) error { + _, err := _u.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_u *BatchImageItemUpdateOne) ExecX(ctx context.Context) { + if err := _u.Exec(ctx); err != nil { + panic(err) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_u *BatchImageItemUpdateOne) check() error { + if v, ok := _u.mutation.JobID(); ok { + if err := batchimageitem.JobIDValidator(v); err != nil { + return &ValidationError{Name: "job_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.job_id": %w`, err)} + } + } + if v, ok := _u.mutation.CustomID(); ok { + if err := batchimageitem.CustomIDValidator(v); err != nil { + return &ValidationError{Name: "custom_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.custom_id": %w`, err)} + } + } + if v, ok := _u.mutation.Status(); ok { + if err := batchimageitem.StatusValidator(v); err != nil { + return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.status": %w`, err)} + } + } + if v, ok := _u.mutation.RequestHash(); ok { + if err := batchimageitem.RequestHashValidator(v); err != nil { + return &ValidationError{Name: "request_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.request_hash": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderSourceObject(); ok { + if err := batchimageitem.ProviderSourceObjectValidator(v); err != nil { + return &ValidationError{Name: "provider_source_object", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.provider_source_object": %w`, err)} + } + } + if v, ok := _u.mutation.MimeType(); ok { + if err := batchimageitem.MimeTypeValidator(v); err != nil { + return &ValidationError{Name: "mime_type", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.mime_type": %w`, err)} + } + } + if v, ok := _u.mutation.FileExtension(); ok { + if err := batchimageitem.FileExtensionValidator(v); err != nil { + return &ValidationError{Name: "file_extension", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.file_extension": %w`, err)} + } + } + if v, ok := _u.mutation.ErrorCode(); ok { + if err := batchimageitem.ErrorCodeValidator(v); err != nil { + return &ValidationError{Name: "error_code", err: fmt.Errorf(`ent: validator failed for field "BatchImageItem.error_code": %w`, err)} + } + } + return nil +} + +func (_u *BatchImageItemUpdateOne) sqlSave(ctx context.Context) (_node *BatchImageItem, err error) { + if err := _u.check(); err != nil { + return _node, err + } + _spec := sqlgraph.NewUpdateSpec(batchimageitem.Table, batchimageitem.Columns, sqlgraph.NewFieldSpec(batchimageitem.FieldID, field.TypeInt64)) + id, ok := _u.mutation.ID() + if !ok { + return nil, &ValidationError{Name: "id", err: errors.New(`ent: missing "BatchImageItem.id" for update`)} + } + _spec.Node.ID.Value = id + if fields := _u.fields; len(fields) > 0 { + _spec.Node.Columns = make([]string, 0, len(fields)) + _spec.Node.Columns = append(_spec.Node.Columns, batchimageitem.FieldID) + for _, f := range fields { + if !batchimageitem.ValidColumn(f) { + return nil, &ValidationError{Name: f, err: fmt.Errorf("ent: invalid field %q for query", f)} + } + if f != batchimageitem.FieldID { + _spec.Node.Columns = append(_spec.Node.Columns, f) + } + } + } + if ps := _u.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if value, ok := _u.mutation.JobID(); ok { + _spec.SetField(batchimageitem.FieldJobID, field.TypeString, value) + } + if value, ok := _u.mutation.CustomID(); ok { + _spec.SetField(batchimageitem.FieldCustomID, field.TypeString, value) + } + if value, ok := _u.mutation.Status(); ok { + _spec.SetField(batchimageitem.FieldStatus, field.TypeString, value) + } + if value, ok := _u.mutation.RequestHash(); ok { + _spec.SetField(batchimageitem.FieldRequestHash, field.TypeString, value) + } + if _u.mutation.RequestHashCleared() { + _spec.ClearField(batchimageitem.FieldRequestHash, field.TypeString) + } + if value, ok := _u.mutation.PromptPreview(); ok { + _spec.SetField(batchimageitem.FieldPromptPreview, field.TypeString, value) + } + if _u.mutation.PromptPreviewCleared() { + _spec.ClearField(batchimageitem.FieldPromptPreview, field.TypeString) + } + if value, ok := _u.mutation.ProviderSourceObject(); ok { + _spec.SetField(batchimageitem.FieldProviderSourceObject, field.TypeString, value) + } + if _u.mutation.ProviderSourceObjectCleared() { + _spec.ClearField(batchimageitem.FieldProviderSourceObject, field.TypeString) + } + if value, ok := _u.mutation.SourceLineNumber(); ok { + _spec.SetField(batchimageitem.FieldSourceLineNumber, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedSourceLineNumber(); ok { + _spec.AddField(batchimageitem.FieldSourceLineNumber, field.TypeInt, value) + } + if _u.mutation.SourceLineNumberCleared() { + _spec.ClearField(batchimageitem.FieldSourceLineNumber, field.TypeInt) + } + if value, ok := _u.mutation.SourceByteOffset(); ok { + _spec.SetField(batchimageitem.FieldSourceByteOffset, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedSourceByteOffset(); ok { + _spec.AddField(batchimageitem.FieldSourceByteOffset, field.TypeInt64, value) + } + if _u.mutation.SourceByteOffsetCleared() { + _spec.ClearField(batchimageitem.FieldSourceByteOffset, field.TypeInt64) + } + if value, ok := _u.mutation.SourceByteLength(); ok { + _spec.SetField(batchimageitem.FieldSourceByteLength, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedSourceByteLength(); ok { + _spec.AddField(batchimageitem.FieldSourceByteLength, field.TypeInt64, value) + } + if _u.mutation.SourceByteLengthCleared() { + _spec.ClearField(batchimageitem.FieldSourceByteLength, field.TypeInt64) + } + if value, ok := _u.mutation.MimeType(); ok { + _spec.SetField(batchimageitem.FieldMimeType, field.TypeString, value) + } + if _u.mutation.MimeTypeCleared() { + _spec.ClearField(batchimageitem.FieldMimeType, field.TypeString) + } + if value, ok := _u.mutation.FileExtension(); ok { + _spec.SetField(batchimageitem.FieldFileExtension, field.TypeString, value) + } + if _u.mutation.FileExtensionCleared() { + _spec.ClearField(batchimageitem.FieldFileExtension, field.TypeString) + } + if value, ok := _u.mutation.ImageCount(); ok { + _spec.SetField(batchimageitem.FieldImageCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedImageCount(); ok { + _spec.AddField(batchimageitem.FieldImageCount, field.TypeInt, value) + } + if value, ok := _u.mutation.ErrorCode(); ok { + _spec.SetField(batchimageitem.FieldErrorCode, field.TypeString, value) + } + if _u.mutation.ErrorCodeCleared() { + _spec.ClearField(batchimageitem.FieldErrorCode, field.TypeString) + } + if value, ok := _u.mutation.ErrorMessage(); ok { + _spec.SetField(batchimageitem.FieldErrorMessage, field.TypeString, value) + } + if _u.mutation.ErrorMessageCleared() { + _spec.ClearField(batchimageitem.FieldErrorMessage, field.TypeString) + } + if value, ok := _u.mutation.BilledAmount(); ok { + _spec.SetField(batchimageitem.FieldBilledAmount, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedBilledAmount(); ok { + _spec.AddField(batchimageitem.FieldBilledAmount, field.TypeFloat64, value) + } + if _u.mutation.BilledAmountCleared() { + _spec.ClearField(batchimageitem.FieldBilledAmount, field.TypeFloat64) + } + if value, ok := _u.mutation.IndexedAt(); ok { + _spec.SetField(batchimageitem.FieldIndexedAt, field.TypeTime, value) + } + if _u.mutation.IndexedAtCleared() { + _spec.ClearField(batchimageitem.FieldIndexedAt, field.TypeTime) + } + _node = &BatchImageItem{config: _u.config} + _spec.Assign = _node.assignValues + _spec.ScanValues = _node.scanValues + if err = sqlgraph.UpdateNode(ctx, _u.driver, _spec); err != nil { + if _, ok := err.(*sqlgraph.NotFoundError); ok { + err = &NotFoundError{batchimageitem.Label} + } else if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return nil, err + } + _u.mutation.done = true + return _node, nil +} diff --git a/backend/ent/batchimagejob.go b/backend/ent/batchimagejob.go new file mode 100644 index 0000000000..29f09bb289 --- /dev/null +++ b/backend/ent/batchimagejob.go @@ -0,0 +1,609 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "fmt" + "strings" + "time" + + "entgo.io/ent" + "entgo.io/ent/dialect/sql" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" +) + +// BatchImageJob is the model entity for the BatchImageJob schema. +type BatchImageJob struct { + config `json:"-"` + // ID of the ent. + ID int64 `json:"id,omitempty"` + // BatchID holds the value of the "batch_id" field. + BatchID string `json:"batch_id,omitempty"` + // UserID holds the value of the "user_id" field. + UserID int64 `json:"user_id,omitempty"` + // APIKeyID holds the value of the "api_key_id" field. + APIKeyID *int64 `json:"api_key_id,omitempty"` + // AccountID holds the value of the "account_id" field. + AccountID *int64 `json:"account_id,omitempty"` + // Provider holds the value of the "provider" field. + Provider string `json:"provider,omitempty"` + // Model holds the value of the "model" field. + Model string `json:"model,omitempty"` + // TaskName holds the value of the "task_name" field. + TaskName string `json:"task_name,omitempty"` + // Status holds the value of the "status" field. + Status string `json:"status,omitempty"` + // ProviderJobName holds the value of the "provider_job_name" field. + ProviderJobName *string `json:"provider_job_name,omitempty"` + // ProviderInputRef holds the value of the "provider_input_ref" field. + ProviderInputRef *string `json:"provider_input_ref,omitempty"` + // ProviderOutputRef holds the value of the "provider_output_ref" field. + ProviderOutputRef *string `json:"provider_output_ref,omitempty"` + // GcsInputURI holds the value of the "gcs_input_uri" field. + GcsInputURI *string `json:"gcs_input_uri,omitempty"` + // GcsOutputURI holds the value of the "gcs_output_uri" field. + GcsOutputURI *string `json:"gcs_output_uri,omitempty"` + // ItemCount holds the value of the "item_count" field. + ItemCount int `json:"item_count,omitempty"` + // SuccessCount holds the value of the "success_count" field. + SuccessCount int `json:"success_count,omitempty"` + // FailCount holds the value of the "fail_count" field. + FailCount int `json:"fail_count,omitempty"` + // CancelledCount holds the value of the "cancelled_count" field. + CancelledCount int `json:"cancelled_count,omitempty"` + // EstimatedCost holds the value of the "estimated_cost" field. + EstimatedCost float64 `json:"estimated_cost,omitempty"` + // HoldAmount holds the value of the "hold_amount" field. + HoldAmount *float64 `json:"hold_amount,omitempty"` + // ActualCost holds the value of the "actual_cost" field. + ActualCost *float64 `json:"actual_cost,omitempty"` + // Currency holds the value of the "currency" field. + Currency string `json:"currency,omitempty"` + // HoldID holds the value of the "hold_id" field. + HoldID *string `json:"hold_id,omitempty"` + // IdempotencyKey holds the value of the "idempotency_key" field. + IdempotencyKey *string `json:"idempotency_key,omitempty"` + // RequestHash holds the value of the "request_hash" field. + RequestHash *string `json:"request_hash,omitempty"` + // ManifestHash holds the value of the "manifest_hash" field. + ManifestHash *string `json:"manifest_hash,omitempty"` + // RetryCount holds the value of the "retry_count" field. + RetryCount int `json:"retry_count,omitempty"` + // Version holds the value of the "version" field. + Version int `json:"version,omitempty"` + // OutputExpiresAt holds the value of the "output_expires_at" field. + OutputExpiresAt *time.Time `json:"output_expires_at,omitempty"` + // InputDeletedAt holds the value of the "input_deleted_at" field. + InputDeletedAt *time.Time `json:"input_deleted_at,omitempty"` + // OutputDeletedAt holds the value of the "output_deleted_at" field. + OutputDeletedAt *time.Time `json:"output_deleted_at,omitempty"` + // DownloadedAt holds the value of the "downloaded_at" field. + DownloadedAt *time.Time `json:"downloaded_at,omitempty"` + // UserDeletedAt holds the value of the "user_deleted_at" field. + UserDeletedAt *time.Time `json:"user_deleted_at,omitempty"` + // LastErrorCode holds the value of the "last_error_code" field. + LastErrorCode *string `json:"last_error_code,omitempty"` + // LastErrorMessage holds the value of the "last_error_message" field. + LastErrorMessage *string `json:"last_error_message,omitempty"` + // CreatedAt holds the value of the "created_at" field. + CreatedAt time.Time `json:"created_at,omitempty"` + // UpdatedAt holds the value of the "updated_at" field. + UpdatedAt time.Time `json:"updated_at,omitempty"` + // SubmittedAt holds the value of the "submitted_at" field. + SubmittedAt *time.Time `json:"submitted_at,omitempty"` + // StartedAt holds the value of the "started_at" field. + StartedAt *time.Time `json:"started_at,omitempty"` + // FinishedAt holds the value of the "finished_at" field. + FinishedAt *time.Time `json:"finished_at,omitempty"` + // SettledAt holds the value of the "settled_at" field. + SettledAt *time.Time `json:"settled_at,omitempty"` + selectValues sql.SelectValues +} + +// scanValues returns the types for scanning values from sql.Rows. +func (*BatchImageJob) scanValues(columns []string) ([]any, error) { + values := make([]any, len(columns)) + for i := range columns { + switch columns[i] { + case batchimagejob.FieldEstimatedCost, batchimagejob.FieldHoldAmount, batchimagejob.FieldActualCost: + values[i] = new(sql.NullFloat64) + case batchimagejob.FieldID, batchimagejob.FieldUserID, batchimagejob.FieldAPIKeyID, batchimagejob.FieldAccountID, batchimagejob.FieldItemCount, batchimagejob.FieldSuccessCount, batchimagejob.FieldFailCount, batchimagejob.FieldCancelledCount, batchimagejob.FieldRetryCount, batchimagejob.FieldVersion: + values[i] = new(sql.NullInt64) + case batchimagejob.FieldBatchID, batchimagejob.FieldProvider, batchimagejob.FieldModel, batchimagejob.FieldTaskName, batchimagejob.FieldStatus, batchimagejob.FieldProviderJobName, batchimagejob.FieldProviderInputRef, batchimagejob.FieldProviderOutputRef, batchimagejob.FieldGcsInputURI, batchimagejob.FieldGcsOutputURI, batchimagejob.FieldCurrency, batchimagejob.FieldHoldID, batchimagejob.FieldIdempotencyKey, batchimagejob.FieldRequestHash, batchimagejob.FieldManifestHash, batchimagejob.FieldLastErrorCode, batchimagejob.FieldLastErrorMessage: + values[i] = new(sql.NullString) + case batchimagejob.FieldOutputExpiresAt, batchimagejob.FieldInputDeletedAt, batchimagejob.FieldOutputDeletedAt, batchimagejob.FieldDownloadedAt, batchimagejob.FieldUserDeletedAt, batchimagejob.FieldCreatedAt, batchimagejob.FieldUpdatedAt, batchimagejob.FieldSubmittedAt, batchimagejob.FieldStartedAt, batchimagejob.FieldFinishedAt, batchimagejob.FieldSettledAt: + values[i] = new(sql.NullTime) + default: + values[i] = new(sql.UnknownType) + } + } + return values, nil +} + +// assignValues assigns the values that were returned from sql.Rows (after scanning) +// to the BatchImageJob fields. +func (_m *BatchImageJob) assignValues(columns []string, values []any) error { + if m, n := len(values), len(columns); m < n { + return fmt.Errorf("mismatch number of scan values: %d != %d", m, n) + } + for i := range columns { + switch columns[i] { + case batchimagejob.FieldID: + value, ok := values[i].(*sql.NullInt64) + if !ok { + return fmt.Errorf("unexpected type %T for field id", value) + } + _m.ID = int64(value.Int64) + case batchimagejob.FieldBatchID: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field batch_id", values[i]) + } else if value.Valid { + _m.BatchID = value.String + } + case batchimagejob.FieldUserID: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field user_id", values[i]) + } else if value.Valid { + _m.UserID = value.Int64 + } + case batchimagejob.FieldAPIKeyID: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field api_key_id", values[i]) + } else if value.Valid { + _m.APIKeyID = new(int64) + *_m.APIKeyID = value.Int64 + } + case batchimagejob.FieldAccountID: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field account_id", values[i]) + } else if value.Valid { + _m.AccountID = new(int64) + *_m.AccountID = value.Int64 + } + case batchimagejob.FieldProvider: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field provider", values[i]) + } else if value.Valid { + _m.Provider = value.String + } + case batchimagejob.FieldModel: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field model", values[i]) + } else if value.Valid { + _m.Model = value.String + } + case batchimagejob.FieldTaskName: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field task_name", values[i]) + } else if value.Valid { + _m.TaskName = value.String + } + case batchimagejob.FieldStatus: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field status", values[i]) + } else if value.Valid { + _m.Status = value.String + } + case batchimagejob.FieldProviderJobName: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field provider_job_name", values[i]) + } else if value.Valid { + _m.ProviderJobName = new(string) + *_m.ProviderJobName = value.String + } + case batchimagejob.FieldProviderInputRef: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field provider_input_ref", values[i]) + } else if value.Valid { + _m.ProviderInputRef = new(string) + *_m.ProviderInputRef = value.String + } + case batchimagejob.FieldProviderOutputRef: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field provider_output_ref", values[i]) + } else if value.Valid { + _m.ProviderOutputRef = new(string) + *_m.ProviderOutputRef = value.String + } + case batchimagejob.FieldGcsInputURI: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field gcs_input_uri", values[i]) + } else if value.Valid { + _m.GcsInputURI = new(string) + *_m.GcsInputURI = value.String + } + case batchimagejob.FieldGcsOutputURI: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field gcs_output_uri", values[i]) + } else if value.Valid { + _m.GcsOutputURI = new(string) + *_m.GcsOutputURI = value.String + } + case batchimagejob.FieldItemCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field item_count", values[i]) + } else if value.Valid { + _m.ItemCount = int(value.Int64) + } + case batchimagejob.FieldSuccessCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field success_count", values[i]) + } else if value.Valid { + _m.SuccessCount = int(value.Int64) + } + case batchimagejob.FieldFailCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field fail_count", values[i]) + } else if value.Valid { + _m.FailCount = int(value.Int64) + } + case batchimagejob.FieldCancelledCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field cancelled_count", values[i]) + } else if value.Valid { + _m.CancelledCount = int(value.Int64) + } + case batchimagejob.FieldEstimatedCost: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field estimated_cost", values[i]) + } else if value.Valid { + _m.EstimatedCost = value.Float64 + } + case batchimagejob.FieldHoldAmount: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field hold_amount", values[i]) + } else if value.Valid { + _m.HoldAmount = new(float64) + *_m.HoldAmount = value.Float64 + } + case batchimagejob.FieldActualCost: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field actual_cost", values[i]) + } else if value.Valid { + _m.ActualCost = new(float64) + *_m.ActualCost = value.Float64 + } + case batchimagejob.FieldCurrency: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field currency", values[i]) + } else if value.Valid { + _m.Currency = value.String + } + case batchimagejob.FieldHoldID: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field hold_id", values[i]) + } else if value.Valid { + _m.HoldID = new(string) + *_m.HoldID = value.String + } + case batchimagejob.FieldIdempotencyKey: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field idempotency_key", values[i]) + } else if value.Valid { + _m.IdempotencyKey = new(string) + *_m.IdempotencyKey = value.String + } + case batchimagejob.FieldRequestHash: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field request_hash", values[i]) + } else if value.Valid { + _m.RequestHash = new(string) + *_m.RequestHash = value.String + } + case batchimagejob.FieldManifestHash: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field manifest_hash", values[i]) + } else if value.Valid { + _m.ManifestHash = new(string) + *_m.ManifestHash = value.String + } + case batchimagejob.FieldRetryCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field retry_count", values[i]) + } else if value.Valid { + _m.RetryCount = int(value.Int64) + } + case batchimagejob.FieldVersion: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field version", values[i]) + } else if value.Valid { + _m.Version = int(value.Int64) + } + case batchimagejob.FieldOutputExpiresAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field output_expires_at", values[i]) + } else if value.Valid { + _m.OutputExpiresAt = new(time.Time) + *_m.OutputExpiresAt = value.Time + } + case batchimagejob.FieldInputDeletedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field input_deleted_at", values[i]) + } else if value.Valid { + _m.InputDeletedAt = new(time.Time) + *_m.InputDeletedAt = value.Time + } + case batchimagejob.FieldOutputDeletedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field output_deleted_at", values[i]) + } else if value.Valid { + _m.OutputDeletedAt = new(time.Time) + *_m.OutputDeletedAt = value.Time + } + case batchimagejob.FieldDownloadedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field downloaded_at", values[i]) + } else if value.Valid { + _m.DownloadedAt = new(time.Time) + *_m.DownloadedAt = value.Time + } + case batchimagejob.FieldUserDeletedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field user_deleted_at", values[i]) + } else if value.Valid { + _m.UserDeletedAt = new(time.Time) + *_m.UserDeletedAt = value.Time + } + case batchimagejob.FieldLastErrorCode: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field last_error_code", values[i]) + } else if value.Valid { + _m.LastErrorCode = new(string) + *_m.LastErrorCode = value.String + } + case batchimagejob.FieldLastErrorMessage: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field last_error_message", values[i]) + } else if value.Valid { + _m.LastErrorMessage = new(string) + *_m.LastErrorMessage = value.String + } + case batchimagejob.FieldCreatedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field created_at", values[i]) + } else if value.Valid { + _m.CreatedAt = value.Time + } + case batchimagejob.FieldUpdatedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field updated_at", values[i]) + } else if value.Valid { + _m.UpdatedAt = value.Time + } + case batchimagejob.FieldSubmittedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field submitted_at", values[i]) + } else if value.Valid { + _m.SubmittedAt = new(time.Time) + *_m.SubmittedAt = value.Time + } + case batchimagejob.FieldStartedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field started_at", values[i]) + } else if value.Valid { + _m.StartedAt = new(time.Time) + *_m.StartedAt = value.Time + } + case batchimagejob.FieldFinishedAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field finished_at", values[i]) + } else if value.Valid { + _m.FinishedAt = new(time.Time) + *_m.FinishedAt = value.Time + } + case batchimagejob.FieldSettledAt: + if value, ok := values[i].(*sql.NullTime); !ok { + return fmt.Errorf("unexpected type %T for field settled_at", values[i]) + } else if value.Valid { + _m.SettledAt = new(time.Time) + *_m.SettledAt = value.Time + } + default: + _m.selectValues.Set(columns[i], values[i]) + } + } + return nil +} + +// Value returns the ent.Value that was dynamically selected and assigned to the BatchImageJob. +// This includes values selected through modifiers, order, etc. +func (_m *BatchImageJob) Value(name string) (ent.Value, error) { + return _m.selectValues.Get(name) +} + +// Update returns a builder for updating this BatchImageJob. +// Note that you need to call BatchImageJob.Unwrap() before calling this method if this BatchImageJob +// was returned from a transaction, and the transaction was committed or rolled back. +func (_m *BatchImageJob) Update() *BatchImageJobUpdateOne { + return NewBatchImageJobClient(_m.config).UpdateOne(_m) +} + +// Unwrap unwraps the BatchImageJob entity that was returned from a transaction after it was closed, +// so that all future queries will be executed through the driver which created the transaction. +func (_m *BatchImageJob) Unwrap() *BatchImageJob { + _tx, ok := _m.config.driver.(*txDriver) + if !ok { + panic("ent: BatchImageJob is not a transactional entity") + } + _m.config.driver = _tx.drv + return _m +} + +// String implements the fmt.Stringer. +func (_m *BatchImageJob) String() string { + var builder strings.Builder + builder.WriteString("BatchImageJob(") + builder.WriteString(fmt.Sprintf("id=%v, ", _m.ID)) + builder.WriteString("batch_id=") + builder.WriteString(_m.BatchID) + builder.WriteString(", ") + builder.WriteString("user_id=") + builder.WriteString(fmt.Sprintf("%v", _m.UserID)) + builder.WriteString(", ") + if v := _m.APIKeyID; v != nil { + builder.WriteString("api_key_id=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.AccountID; v != nil { + builder.WriteString("account_id=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + builder.WriteString("provider=") + builder.WriteString(_m.Provider) + builder.WriteString(", ") + builder.WriteString("model=") + builder.WriteString(_m.Model) + builder.WriteString(", ") + builder.WriteString("task_name=") + builder.WriteString(_m.TaskName) + builder.WriteString(", ") + builder.WriteString("status=") + builder.WriteString(_m.Status) + builder.WriteString(", ") + if v := _m.ProviderJobName; v != nil { + builder.WriteString("provider_job_name=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.ProviderInputRef; v != nil { + builder.WriteString("provider_input_ref=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.ProviderOutputRef; v != nil { + builder.WriteString("provider_output_ref=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.GcsInputURI; v != nil { + builder.WriteString("gcs_input_uri=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.GcsOutputURI; v != nil { + builder.WriteString("gcs_output_uri=") + builder.WriteString(*v) + } + builder.WriteString(", ") + builder.WriteString("item_count=") + builder.WriteString(fmt.Sprintf("%v", _m.ItemCount)) + builder.WriteString(", ") + builder.WriteString("success_count=") + builder.WriteString(fmt.Sprintf("%v", _m.SuccessCount)) + builder.WriteString(", ") + builder.WriteString("fail_count=") + builder.WriteString(fmt.Sprintf("%v", _m.FailCount)) + builder.WriteString(", ") + builder.WriteString("cancelled_count=") + builder.WriteString(fmt.Sprintf("%v", _m.CancelledCount)) + builder.WriteString(", ") + builder.WriteString("estimated_cost=") + builder.WriteString(fmt.Sprintf("%v", _m.EstimatedCost)) + builder.WriteString(", ") + if v := _m.HoldAmount; v != nil { + builder.WriteString("hold_amount=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.ActualCost; v != nil { + builder.WriteString("actual_cost=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + builder.WriteString("currency=") + builder.WriteString(_m.Currency) + builder.WriteString(", ") + if v := _m.HoldID; v != nil { + builder.WriteString("hold_id=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.IdempotencyKey; v != nil { + builder.WriteString("idempotency_key=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.RequestHash; v != nil { + builder.WriteString("request_hash=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.ManifestHash; v != nil { + builder.WriteString("manifest_hash=") + builder.WriteString(*v) + } + builder.WriteString(", ") + builder.WriteString("retry_count=") + builder.WriteString(fmt.Sprintf("%v", _m.RetryCount)) + builder.WriteString(", ") + builder.WriteString("version=") + builder.WriteString(fmt.Sprintf("%v", _m.Version)) + builder.WriteString(", ") + if v := _m.OutputExpiresAt; v != nil { + builder.WriteString("output_expires_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.InputDeletedAt; v != nil { + builder.WriteString("input_deleted_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.OutputDeletedAt; v != nil { + builder.WriteString("output_deleted_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.DownloadedAt; v != nil { + builder.WriteString("downloaded_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.UserDeletedAt; v != nil { + builder.WriteString("user_deleted_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.LastErrorCode; v != nil { + builder.WriteString("last_error_code=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.LastErrorMessage; v != nil { + builder.WriteString("last_error_message=") + builder.WriteString(*v) + } + builder.WriteString(", ") + builder.WriteString("created_at=") + builder.WriteString(_m.CreatedAt.Format(time.ANSIC)) + builder.WriteString(", ") + builder.WriteString("updated_at=") + builder.WriteString(_m.UpdatedAt.Format(time.ANSIC)) + builder.WriteString(", ") + if v := _m.SubmittedAt; v != nil { + builder.WriteString("submitted_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.StartedAt; v != nil { + builder.WriteString("started_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.FinishedAt; v != nil { + builder.WriteString("finished_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteString(", ") + if v := _m.SettledAt; v != nil { + builder.WriteString("settled_at=") + builder.WriteString(v.Format(time.ANSIC)) + } + builder.WriteByte(')') + return builder.String() +} + +// BatchImageJobs is a parsable slice of BatchImageJob. +type BatchImageJobs []*BatchImageJob diff --git a/backend/ent/batchimagejob/batchimagejob.go b/backend/ent/batchimagejob/batchimagejob.go new file mode 100644 index 0000000000..be819183da --- /dev/null +++ b/backend/ent/batchimagejob/batchimagejob.go @@ -0,0 +1,420 @@ +// Code generated by ent, DO NOT EDIT. + +package batchimagejob + +import ( + "time" + + "entgo.io/ent/dialect/sql" +) + +const ( + // Label holds the string label denoting the batchimagejob type in the database. + Label = "batch_image_job" + // FieldID holds the string denoting the id field in the database. + FieldID = "id" + // FieldBatchID holds the string denoting the batch_id field in the database. + FieldBatchID = "batch_id" + // FieldUserID holds the string denoting the user_id field in the database. + FieldUserID = "user_id" + // FieldAPIKeyID holds the string denoting the api_key_id field in the database. + FieldAPIKeyID = "api_key_id" + // FieldAccountID holds the string denoting the account_id field in the database. + FieldAccountID = "account_id" + // FieldProvider holds the string denoting the provider field in the database. + FieldProvider = "provider" + // FieldModel holds the string denoting the model field in the database. + FieldModel = "model" + // FieldTaskName holds the string denoting the task_name field in the database. + FieldTaskName = "task_name" + // FieldStatus holds the string denoting the status field in the database. + FieldStatus = "status" + // FieldProviderJobName holds the string denoting the provider_job_name field in the database. + FieldProviderJobName = "provider_job_name" + // FieldProviderInputRef holds the string denoting the provider_input_ref field in the database. + FieldProviderInputRef = "provider_input_ref" + // FieldProviderOutputRef holds the string denoting the provider_output_ref field in the database. + FieldProviderOutputRef = "provider_output_ref" + // FieldGcsInputURI holds the string denoting the gcs_input_uri field in the database. + FieldGcsInputURI = "gcs_input_uri" + // FieldGcsOutputURI holds the string denoting the gcs_output_uri field in the database. + FieldGcsOutputURI = "gcs_output_uri" + // FieldItemCount holds the string denoting the item_count field in the database. + FieldItemCount = "item_count" + // FieldSuccessCount holds the string denoting the success_count field in the database. + FieldSuccessCount = "success_count" + // FieldFailCount holds the string denoting the fail_count field in the database. + FieldFailCount = "fail_count" + // FieldCancelledCount holds the string denoting the cancelled_count field in the database. + FieldCancelledCount = "cancelled_count" + // FieldEstimatedCost holds the string denoting the estimated_cost field in the database. + FieldEstimatedCost = "estimated_cost" + // FieldHoldAmount holds the string denoting the hold_amount field in the database. + FieldHoldAmount = "hold_amount" + // FieldActualCost holds the string denoting the actual_cost field in the database. + FieldActualCost = "actual_cost" + // FieldCurrency holds the string denoting the currency field in the database. + FieldCurrency = "currency" + // FieldHoldID holds the string denoting the hold_id field in the database. + FieldHoldID = "hold_id" + // FieldIdempotencyKey holds the string denoting the idempotency_key field in the database. + FieldIdempotencyKey = "idempotency_key" + // FieldRequestHash holds the string denoting the request_hash field in the database. + FieldRequestHash = "request_hash" + // FieldManifestHash holds the string denoting the manifest_hash field in the database. + FieldManifestHash = "manifest_hash" + // FieldRetryCount holds the string denoting the retry_count field in the database. + FieldRetryCount = "retry_count" + // FieldVersion holds the string denoting the version field in the database. + FieldVersion = "version" + // FieldOutputExpiresAt holds the string denoting the output_expires_at field in the database. + FieldOutputExpiresAt = "output_expires_at" + // FieldInputDeletedAt holds the string denoting the input_deleted_at field in the database. + FieldInputDeletedAt = "input_deleted_at" + // FieldOutputDeletedAt holds the string denoting the output_deleted_at field in the database. + FieldOutputDeletedAt = "output_deleted_at" + // FieldDownloadedAt holds the string denoting the downloaded_at field in the database. + FieldDownloadedAt = "downloaded_at" + // FieldUserDeletedAt holds the string denoting the user_deleted_at field in the database. + FieldUserDeletedAt = "user_deleted_at" + // FieldLastErrorCode holds the string denoting the last_error_code field in the database. + FieldLastErrorCode = "last_error_code" + // FieldLastErrorMessage holds the string denoting the last_error_message field in the database. + FieldLastErrorMessage = "last_error_message" + // FieldCreatedAt holds the string denoting the created_at field in the database. + FieldCreatedAt = "created_at" + // FieldUpdatedAt holds the string denoting the updated_at field in the database. + FieldUpdatedAt = "updated_at" + // FieldSubmittedAt holds the string denoting the submitted_at field in the database. + FieldSubmittedAt = "submitted_at" + // FieldStartedAt holds the string denoting the started_at field in the database. + FieldStartedAt = "started_at" + // FieldFinishedAt holds the string denoting the finished_at field in the database. + FieldFinishedAt = "finished_at" + // FieldSettledAt holds the string denoting the settled_at field in the database. + FieldSettledAt = "settled_at" + // Table holds the table name of the batchimagejob in the database. + Table = "batch_image_jobs" +) + +// Columns holds all SQL columns for batchimagejob fields. +var Columns = []string{ + FieldID, + FieldBatchID, + FieldUserID, + FieldAPIKeyID, + FieldAccountID, + FieldProvider, + FieldModel, + FieldTaskName, + FieldStatus, + FieldProviderJobName, + FieldProviderInputRef, + FieldProviderOutputRef, + FieldGcsInputURI, + FieldGcsOutputURI, + FieldItemCount, + FieldSuccessCount, + FieldFailCount, + FieldCancelledCount, + FieldEstimatedCost, + FieldHoldAmount, + FieldActualCost, + FieldCurrency, + FieldHoldID, + FieldIdempotencyKey, + FieldRequestHash, + FieldManifestHash, + FieldRetryCount, + FieldVersion, + FieldOutputExpiresAt, + FieldInputDeletedAt, + FieldOutputDeletedAt, + FieldDownloadedAt, + FieldUserDeletedAt, + FieldLastErrorCode, + FieldLastErrorMessage, + FieldCreatedAt, + FieldUpdatedAt, + FieldSubmittedAt, + FieldStartedAt, + FieldFinishedAt, + FieldSettledAt, +} + +// ValidColumn reports if the column name is valid (part of the table columns). +func ValidColumn(column string) bool { + for i := range Columns { + if column == Columns[i] { + return true + } + } + return false +} + +var ( + // BatchIDValidator is a validator for the "batch_id" field. It is called by the builders before save. + BatchIDValidator func(string) error + // ProviderValidator is a validator for the "provider" field. It is called by the builders before save. + ProviderValidator func(string) error + // ModelValidator is a validator for the "model" field. It is called by the builders before save. + ModelValidator func(string) error + // DefaultTaskName holds the default value on creation for the "task_name" field. + DefaultTaskName string + // TaskNameValidator is a validator for the "task_name" field. It is called by the builders before save. + TaskNameValidator func(string) error + // DefaultStatus holds the default value on creation for the "status" field. + DefaultStatus string + // StatusValidator is a validator for the "status" field. It is called by the builders before save. + StatusValidator func(string) error + // ProviderJobNameValidator is a validator for the "provider_job_name" field. It is called by the builders before save. + ProviderJobNameValidator func(string) error + // ProviderInputRefValidator is a validator for the "provider_input_ref" field. It is called by the builders before save. + ProviderInputRefValidator func(string) error + // ProviderOutputRefValidator is a validator for the "provider_output_ref" field. It is called by the builders before save. + ProviderOutputRefValidator func(string) error + // GcsInputURIValidator is a validator for the "gcs_input_uri" field. It is called by the builders before save. + GcsInputURIValidator func(string) error + // GcsOutputURIValidator is a validator for the "gcs_output_uri" field. It is called by the builders before save. + GcsOutputURIValidator func(string) error + // DefaultSuccessCount holds the default value on creation for the "success_count" field. + DefaultSuccessCount int + // DefaultFailCount holds the default value on creation for the "fail_count" field. + DefaultFailCount int + // DefaultCancelledCount holds the default value on creation for the "cancelled_count" field. + DefaultCancelledCount int + // DefaultEstimatedCost holds the default value on creation for the "estimated_cost" field. + DefaultEstimatedCost float64 + // DefaultCurrency holds the default value on creation for the "currency" field. + DefaultCurrency string + // CurrencyValidator is a validator for the "currency" field. It is called by the builders before save. + CurrencyValidator func(string) error + // HoldIDValidator is a validator for the "hold_id" field. It is called by the builders before save. + HoldIDValidator func(string) error + // IdempotencyKeyValidator is a validator for the "idempotency_key" field. It is called by the builders before save. + IdempotencyKeyValidator func(string) error + // RequestHashValidator is a validator for the "request_hash" field. It is called by the builders before save. + RequestHashValidator func(string) error + // ManifestHashValidator is a validator for the "manifest_hash" field. It is called by the builders before save. + ManifestHashValidator func(string) error + // DefaultRetryCount holds the default value on creation for the "retry_count" field. + DefaultRetryCount int + // DefaultVersion holds the default value on creation for the "version" field. + DefaultVersion int + // LastErrorCodeValidator is a validator for the "last_error_code" field. It is called by the builders before save. + LastErrorCodeValidator func(string) error + // DefaultCreatedAt holds the default value on creation for the "created_at" field. + DefaultCreatedAt func() time.Time + // DefaultUpdatedAt holds the default value on creation for the "updated_at" field. + DefaultUpdatedAt func() time.Time + // UpdateDefaultUpdatedAt holds the default value on update for the "updated_at" field. + UpdateDefaultUpdatedAt func() time.Time +) + +// OrderOption defines the ordering options for the BatchImageJob queries. +type OrderOption func(*sql.Selector) + +// ByID orders the results by the id field. +func ByID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldID, opts...).ToFunc() +} + +// ByBatchID orders the results by the batch_id field. +func ByBatchID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldBatchID, opts...).ToFunc() +} + +// ByUserID orders the results by the user_id field. +func ByUserID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldUserID, opts...).ToFunc() +} + +// ByAPIKeyID orders the results by the api_key_id field. +func ByAPIKeyID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldAPIKeyID, opts...).ToFunc() +} + +// ByAccountID orders the results by the account_id field. +func ByAccountID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldAccountID, opts...).ToFunc() +} + +// ByProvider orders the results by the provider field. +func ByProvider(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldProvider, opts...).ToFunc() +} + +// ByModel orders the results by the model field. +func ByModel(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldModel, opts...).ToFunc() +} + +// ByTaskName orders the results by the task_name field. +func ByTaskName(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldTaskName, opts...).ToFunc() +} + +// ByStatus orders the results by the status field. +func ByStatus(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldStatus, opts...).ToFunc() +} + +// ByProviderJobName orders the results by the provider_job_name field. +func ByProviderJobName(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldProviderJobName, opts...).ToFunc() +} + +// ByProviderInputRef orders the results by the provider_input_ref field. +func ByProviderInputRef(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldProviderInputRef, opts...).ToFunc() +} + +// ByProviderOutputRef orders the results by the provider_output_ref field. +func ByProviderOutputRef(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldProviderOutputRef, opts...).ToFunc() +} + +// ByGcsInputURI orders the results by the gcs_input_uri field. +func ByGcsInputURI(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldGcsInputURI, opts...).ToFunc() +} + +// ByGcsOutputURI orders the results by the gcs_output_uri field. +func ByGcsOutputURI(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldGcsOutputURI, opts...).ToFunc() +} + +// ByItemCount orders the results by the item_count field. +func ByItemCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldItemCount, opts...).ToFunc() +} + +// BySuccessCount orders the results by the success_count field. +func BySuccessCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldSuccessCount, opts...).ToFunc() +} + +// ByFailCount orders the results by the fail_count field. +func ByFailCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldFailCount, opts...).ToFunc() +} + +// ByCancelledCount orders the results by the cancelled_count field. +func ByCancelledCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldCancelledCount, opts...).ToFunc() +} + +// ByEstimatedCost orders the results by the estimated_cost field. +func ByEstimatedCost(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldEstimatedCost, opts...).ToFunc() +} + +// ByHoldAmount orders the results by the hold_amount field. +func ByHoldAmount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldHoldAmount, opts...).ToFunc() +} + +// ByActualCost orders the results by the actual_cost field. +func ByActualCost(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldActualCost, opts...).ToFunc() +} + +// ByCurrency orders the results by the currency field. +func ByCurrency(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldCurrency, opts...).ToFunc() +} + +// ByHoldID orders the results by the hold_id field. +func ByHoldID(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldHoldID, opts...).ToFunc() +} + +// ByIdempotencyKey orders the results by the idempotency_key field. +func ByIdempotencyKey(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldIdempotencyKey, opts...).ToFunc() +} + +// ByRequestHash orders the results by the request_hash field. +func ByRequestHash(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldRequestHash, opts...).ToFunc() +} + +// ByManifestHash orders the results by the manifest_hash field. +func ByManifestHash(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldManifestHash, opts...).ToFunc() +} + +// ByRetryCount orders the results by the retry_count field. +func ByRetryCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldRetryCount, opts...).ToFunc() +} + +// ByVersion orders the results by the version field. +func ByVersion(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVersion, opts...).ToFunc() +} + +// ByOutputExpiresAt orders the results by the output_expires_at field. +func ByOutputExpiresAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldOutputExpiresAt, opts...).ToFunc() +} + +// ByInputDeletedAt orders the results by the input_deleted_at field. +func ByInputDeletedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldInputDeletedAt, opts...).ToFunc() +} + +// ByOutputDeletedAt orders the results by the output_deleted_at field. +func ByOutputDeletedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldOutputDeletedAt, opts...).ToFunc() +} + +// ByDownloadedAt orders the results by the downloaded_at field. +func ByDownloadedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldDownloadedAt, opts...).ToFunc() +} + +// ByUserDeletedAt orders the results by the user_deleted_at field. +func ByUserDeletedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldUserDeletedAt, opts...).ToFunc() +} + +// ByLastErrorCode orders the results by the last_error_code field. +func ByLastErrorCode(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldLastErrorCode, opts...).ToFunc() +} + +// ByLastErrorMessage orders the results by the last_error_message field. +func ByLastErrorMessage(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldLastErrorMessage, opts...).ToFunc() +} + +// ByCreatedAt orders the results by the created_at field. +func ByCreatedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldCreatedAt, opts...).ToFunc() +} + +// ByUpdatedAt orders the results by the updated_at field. +func ByUpdatedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldUpdatedAt, opts...).ToFunc() +} + +// BySubmittedAt orders the results by the submitted_at field. +func BySubmittedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldSubmittedAt, opts...).ToFunc() +} + +// ByStartedAt orders the results by the started_at field. +func ByStartedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldStartedAt, opts...).ToFunc() +} + +// ByFinishedAt orders the results by the finished_at field. +func ByFinishedAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldFinishedAt, opts...).ToFunc() +} + +// BySettledAt orders the results by the settled_at field. +func BySettledAt(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldSettledAt, opts...).ToFunc() +} diff --git a/backend/ent/batchimagejob/where.go b/backend/ent/batchimagejob/where.go new file mode 100644 index 0000000000..b94722e41d --- /dev/null +++ b/backend/ent/batchimagejob/where.go @@ -0,0 +1,2535 @@ +// Code generated by ent, DO NOT EDIT. + +package batchimagejob + +import ( + "time" + + "entgo.io/ent/dialect/sql" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// ID filters vertices based on their ID field. +func ID(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldID, id)) +} + +// IDEQ applies the EQ predicate on the ID field. +func IDEQ(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldID, id)) +} + +// IDNEQ applies the NEQ predicate on the ID field. +func IDNEQ(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldID, id)) +} + +// IDIn applies the In predicate on the ID field. +func IDIn(ids ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldID, ids...)) +} + +// IDNotIn applies the NotIn predicate on the ID field. +func IDNotIn(ids ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldID, ids...)) +} + +// IDGT applies the GT predicate on the ID field. +func IDGT(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldID, id)) +} + +// IDGTE applies the GTE predicate on the ID field. +func IDGTE(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldID, id)) +} + +// IDLT applies the LT predicate on the ID field. +func IDLT(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldID, id)) +} + +// IDLTE applies the LTE predicate on the ID field. +func IDLTE(id int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldID, id)) +} + +// BatchID applies equality check predicate on the "batch_id" field. It's identical to BatchIDEQ. +func BatchID(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldBatchID, v)) +} + +// UserID applies equality check predicate on the "user_id" field. It's identical to UserIDEQ. +func UserID(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldUserID, v)) +} + +// APIKeyID applies equality check predicate on the "api_key_id" field. It's identical to APIKeyIDEQ. +func APIKeyID(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldAPIKeyID, v)) +} + +// AccountID applies equality check predicate on the "account_id" field. It's identical to AccountIDEQ. +func AccountID(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldAccountID, v)) +} + +// Provider applies equality check predicate on the "provider" field. It's identical to ProviderEQ. +func Provider(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProvider, v)) +} + +// Model applies equality check predicate on the "model" field. It's identical to ModelEQ. +func Model(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldModel, v)) +} + +// TaskName applies equality check predicate on the "task_name" field. It's identical to TaskNameEQ. +func TaskName(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldTaskName, v)) +} + +// Status applies equality check predicate on the "status" field. It's identical to StatusEQ. +func Status(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldStatus, v)) +} + +// ProviderJobName applies equality check predicate on the "provider_job_name" field. It's identical to ProviderJobNameEQ. +func ProviderJobName(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProviderJobName, v)) +} + +// ProviderInputRef applies equality check predicate on the "provider_input_ref" field. It's identical to ProviderInputRefEQ. +func ProviderInputRef(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProviderInputRef, v)) +} + +// ProviderOutputRef applies equality check predicate on the "provider_output_ref" field. It's identical to ProviderOutputRefEQ. +func ProviderOutputRef(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProviderOutputRef, v)) +} + +// GcsInputURI applies equality check predicate on the "gcs_input_uri" field. It's identical to GcsInputURIEQ. +func GcsInputURI(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldGcsInputURI, v)) +} + +// GcsOutputURI applies equality check predicate on the "gcs_output_uri" field. It's identical to GcsOutputURIEQ. +func GcsOutputURI(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldGcsOutputURI, v)) +} + +// ItemCount applies equality check predicate on the "item_count" field. It's identical to ItemCountEQ. +func ItemCount(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldItemCount, v)) +} + +// SuccessCount applies equality check predicate on the "success_count" field. It's identical to SuccessCountEQ. +func SuccessCount(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldSuccessCount, v)) +} + +// FailCount applies equality check predicate on the "fail_count" field. It's identical to FailCountEQ. +func FailCount(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldFailCount, v)) +} + +// CancelledCount applies equality check predicate on the "cancelled_count" field. It's identical to CancelledCountEQ. +func CancelledCount(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldCancelledCount, v)) +} + +// EstimatedCost applies equality check predicate on the "estimated_cost" field. It's identical to EstimatedCostEQ. +func EstimatedCost(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldEstimatedCost, v)) +} + +// HoldAmount applies equality check predicate on the "hold_amount" field. It's identical to HoldAmountEQ. +func HoldAmount(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldHoldAmount, v)) +} + +// ActualCost applies equality check predicate on the "actual_cost" field. It's identical to ActualCostEQ. +func ActualCost(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldActualCost, v)) +} + +// Currency applies equality check predicate on the "currency" field. It's identical to CurrencyEQ. +func Currency(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldCurrency, v)) +} + +// HoldID applies equality check predicate on the "hold_id" field. It's identical to HoldIDEQ. +func HoldID(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldHoldID, v)) +} + +// IdempotencyKey applies equality check predicate on the "idempotency_key" field. It's identical to IdempotencyKeyEQ. +func IdempotencyKey(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldIdempotencyKey, v)) +} + +// RequestHash applies equality check predicate on the "request_hash" field. It's identical to RequestHashEQ. +func RequestHash(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldRequestHash, v)) +} + +// ManifestHash applies equality check predicate on the "manifest_hash" field. It's identical to ManifestHashEQ. +func ManifestHash(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldManifestHash, v)) +} + +// RetryCount applies equality check predicate on the "retry_count" field. It's identical to RetryCountEQ. +func RetryCount(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldRetryCount, v)) +} + +// Version applies equality check predicate on the "version" field. It's identical to VersionEQ. +func Version(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldVersion, v)) +} + +// OutputExpiresAt applies equality check predicate on the "output_expires_at" field. It's identical to OutputExpiresAtEQ. +func OutputExpiresAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldOutputExpiresAt, v)) +} + +// InputDeletedAt applies equality check predicate on the "input_deleted_at" field. It's identical to InputDeletedAtEQ. +func InputDeletedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldInputDeletedAt, v)) +} + +// OutputDeletedAt applies equality check predicate on the "output_deleted_at" field. It's identical to OutputDeletedAtEQ. +func OutputDeletedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldOutputDeletedAt, v)) +} + +// DownloadedAt applies equality check predicate on the "downloaded_at" field. It's identical to DownloadedAtEQ. +func DownloadedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldDownloadedAt, v)) +} + +// UserDeletedAt applies equality check predicate on the "user_deleted_at" field. It's identical to UserDeletedAtEQ. +func UserDeletedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldUserDeletedAt, v)) +} + +// LastErrorCode applies equality check predicate on the "last_error_code" field. It's identical to LastErrorCodeEQ. +func LastErrorCode(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldLastErrorCode, v)) +} + +// LastErrorMessage applies equality check predicate on the "last_error_message" field. It's identical to LastErrorMessageEQ. +func LastErrorMessage(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldLastErrorMessage, v)) +} + +// CreatedAt applies equality check predicate on the "created_at" field. It's identical to CreatedAtEQ. +func CreatedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldCreatedAt, v)) +} + +// UpdatedAt applies equality check predicate on the "updated_at" field. It's identical to UpdatedAtEQ. +func UpdatedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldUpdatedAt, v)) +} + +// SubmittedAt applies equality check predicate on the "submitted_at" field. It's identical to SubmittedAtEQ. +func SubmittedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldSubmittedAt, v)) +} + +// StartedAt applies equality check predicate on the "started_at" field. It's identical to StartedAtEQ. +func StartedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldStartedAt, v)) +} + +// FinishedAt applies equality check predicate on the "finished_at" field. It's identical to FinishedAtEQ. +func FinishedAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldFinishedAt, v)) +} + +// SettledAt applies equality check predicate on the "settled_at" field. It's identical to SettledAtEQ. +func SettledAt(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldSettledAt, v)) +} + +// BatchIDEQ applies the EQ predicate on the "batch_id" field. +func BatchIDEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldBatchID, v)) +} + +// BatchIDNEQ applies the NEQ predicate on the "batch_id" field. +func BatchIDNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldBatchID, v)) +} + +// BatchIDIn applies the In predicate on the "batch_id" field. +func BatchIDIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldBatchID, vs...)) +} + +// BatchIDNotIn applies the NotIn predicate on the "batch_id" field. +func BatchIDNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldBatchID, vs...)) +} + +// BatchIDGT applies the GT predicate on the "batch_id" field. +func BatchIDGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldBatchID, v)) +} + +// BatchIDGTE applies the GTE predicate on the "batch_id" field. +func BatchIDGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldBatchID, v)) +} + +// BatchIDLT applies the LT predicate on the "batch_id" field. +func BatchIDLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldBatchID, v)) +} + +// BatchIDLTE applies the LTE predicate on the "batch_id" field. +func BatchIDLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldBatchID, v)) +} + +// BatchIDContains applies the Contains predicate on the "batch_id" field. +func BatchIDContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldBatchID, v)) +} + +// BatchIDHasPrefix applies the HasPrefix predicate on the "batch_id" field. +func BatchIDHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldBatchID, v)) +} + +// BatchIDHasSuffix applies the HasSuffix predicate on the "batch_id" field. +func BatchIDHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldBatchID, v)) +} + +// BatchIDEqualFold applies the EqualFold predicate on the "batch_id" field. +func BatchIDEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldBatchID, v)) +} + +// BatchIDContainsFold applies the ContainsFold predicate on the "batch_id" field. +func BatchIDContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldBatchID, v)) +} + +// UserIDEQ applies the EQ predicate on the "user_id" field. +func UserIDEQ(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldUserID, v)) +} + +// UserIDNEQ applies the NEQ predicate on the "user_id" field. +func UserIDNEQ(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldUserID, v)) +} + +// UserIDIn applies the In predicate on the "user_id" field. +func UserIDIn(vs ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldUserID, vs...)) +} + +// UserIDNotIn applies the NotIn predicate on the "user_id" field. +func UserIDNotIn(vs ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldUserID, vs...)) +} + +// UserIDGT applies the GT predicate on the "user_id" field. +func UserIDGT(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldUserID, v)) +} + +// UserIDGTE applies the GTE predicate on the "user_id" field. +func UserIDGTE(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldUserID, v)) +} + +// UserIDLT applies the LT predicate on the "user_id" field. +func UserIDLT(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldUserID, v)) +} + +// UserIDLTE applies the LTE predicate on the "user_id" field. +func UserIDLTE(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldUserID, v)) +} + +// APIKeyIDEQ applies the EQ predicate on the "api_key_id" field. +func APIKeyIDEQ(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldAPIKeyID, v)) +} + +// APIKeyIDNEQ applies the NEQ predicate on the "api_key_id" field. +func APIKeyIDNEQ(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldAPIKeyID, v)) +} + +// APIKeyIDIn applies the In predicate on the "api_key_id" field. +func APIKeyIDIn(vs ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldAPIKeyID, vs...)) +} + +// APIKeyIDNotIn applies the NotIn predicate on the "api_key_id" field. +func APIKeyIDNotIn(vs ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldAPIKeyID, vs...)) +} + +// APIKeyIDGT applies the GT predicate on the "api_key_id" field. +func APIKeyIDGT(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldAPIKeyID, v)) +} + +// APIKeyIDGTE applies the GTE predicate on the "api_key_id" field. +func APIKeyIDGTE(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldAPIKeyID, v)) +} + +// APIKeyIDLT applies the LT predicate on the "api_key_id" field. +func APIKeyIDLT(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldAPIKeyID, v)) +} + +// APIKeyIDLTE applies the LTE predicate on the "api_key_id" field. +func APIKeyIDLTE(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldAPIKeyID, v)) +} + +// APIKeyIDIsNil applies the IsNil predicate on the "api_key_id" field. +func APIKeyIDIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldAPIKeyID)) +} + +// APIKeyIDNotNil applies the NotNil predicate on the "api_key_id" field. +func APIKeyIDNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldAPIKeyID)) +} + +// AccountIDEQ applies the EQ predicate on the "account_id" field. +func AccountIDEQ(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldAccountID, v)) +} + +// AccountIDNEQ applies the NEQ predicate on the "account_id" field. +func AccountIDNEQ(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldAccountID, v)) +} + +// AccountIDIn applies the In predicate on the "account_id" field. +func AccountIDIn(vs ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldAccountID, vs...)) +} + +// AccountIDNotIn applies the NotIn predicate on the "account_id" field. +func AccountIDNotIn(vs ...int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldAccountID, vs...)) +} + +// AccountIDGT applies the GT predicate on the "account_id" field. +func AccountIDGT(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldAccountID, v)) +} + +// AccountIDGTE applies the GTE predicate on the "account_id" field. +func AccountIDGTE(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldAccountID, v)) +} + +// AccountIDLT applies the LT predicate on the "account_id" field. +func AccountIDLT(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldAccountID, v)) +} + +// AccountIDLTE applies the LTE predicate on the "account_id" field. +func AccountIDLTE(v int64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldAccountID, v)) +} + +// AccountIDIsNil applies the IsNil predicate on the "account_id" field. +func AccountIDIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldAccountID)) +} + +// AccountIDNotNil applies the NotNil predicate on the "account_id" field. +func AccountIDNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldAccountID)) +} + +// ProviderEQ applies the EQ predicate on the "provider" field. +func ProviderEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProvider, v)) +} + +// ProviderNEQ applies the NEQ predicate on the "provider" field. +func ProviderNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldProvider, v)) +} + +// ProviderIn applies the In predicate on the "provider" field. +func ProviderIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldProvider, vs...)) +} + +// ProviderNotIn applies the NotIn predicate on the "provider" field. +func ProviderNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldProvider, vs...)) +} + +// ProviderGT applies the GT predicate on the "provider" field. +func ProviderGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldProvider, v)) +} + +// ProviderGTE applies the GTE predicate on the "provider" field. +func ProviderGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldProvider, v)) +} + +// ProviderLT applies the LT predicate on the "provider" field. +func ProviderLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldProvider, v)) +} + +// ProviderLTE applies the LTE predicate on the "provider" field. +func ProviderLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldProvider, v)) +} + +// ProviderContains applies the Contains predicate on the "provider" field. +func ProviderContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldProvider, v)) +} + +// ProviderHasPrefix applies the HasPrefix predicate on the "provider" field. +func ProviderHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldProvider, v)) +} + +// ProviderHasSuffix applies the HasSuffix predicate on the "provider" field. +func ProviderHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldProvider, v)) +} + +// ProviderEqualFold applies the EqualFold predicate on the "provider" field. +func ProviderEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldProvider, v)) +} + +// ProviderContainsFold applies the ContainsFold predicate on the "provider" field. +func ProviderContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldProvider, v)) +} + +// ModelEQ applies the EQ predicate on the "model" field. +func ModelEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldModel, v)) +} + +// ModelNEQ applies the NEQ predicate on the "model" field. +func ModelNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldModel, v)) +} + +// ModelIn applies the In predicate on the "model" field. +func ModelIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldModel, vs...)) +} + +// ModelNotIn applies the NotIn predicate on the "model" field. +func ModelNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldModel, vs...)) +} + +// ModelGT applies the GT predicate on the "model" field. +func ModelGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldModel, v)) +} + +// ModelGTE applies the GTE predicate on the "model" field. +func ModelGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldModel, v)) +} + +// ModelLT applies the LT predicate on the "model" field. +func ModelLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldModel, v)) +} + +// ModelLTE applies the LTE predicate on the "model" field. +func ModelLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldModel, v)) +} + +// ModelContains applies the Contains predicate on the "model" field. +func ModelContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldModel, v)) +} + +// ModelHasPrefix applies the HasPrefix predicate on the "model" field. +func ModelHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldModel, v)) +} + +// ModelHasSuffix applies the HasSuffix predicate on the "model" field. +func ModelHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldModel, v)) +} + +// ModelEqualFold applies the EqualFold predicate on the "model" field. +func ModelEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldModel, v)) +} + +// ModelContainsFold applies the ContainsFold predicate on the "model" field. +func ModelContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldModel, v)) +} + +// TaskNameEQ applies the EQ predicate on the "task_name" field. +func TaskNameEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldTaskName, v)) +} + +// TaskNameNEQ applies the NEQ predicate on the "task_name" field. +func TaskNameNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldTaskName, v)) +} + +// TaskNameIn applies the In predicate on the "task_name" field. +func TaskNameIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldTaskName, vs...)) +} + +// TaskNameNotIn applies the NotIn predicate on the "task_name" field. +func TaskNameNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldTaskName, vs...)) +} + +// TaskNameGT applies the GT predicate on the "task_name" field. +func TaskNameGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldTaskName, v)) +} + +// TaskNameGTE applies the GTE predicate on the "task_name" field. +func TaskNameGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldTaskName, v)) +} + +// TaskNameLT applies the LT predicate on the "task_name" field. +func TaskNameLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldTaskName, v)) +} + +// TaskNameLTE applies the LTE predicate on the "task_name" field. +func TaskNameLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldTaskName, v)) +} + +// TaskNameContains applies the Contains predicate on the "task_name" field. +func TaskNameContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldTaskName, v)) +} + +// TaskNameHasPrefix applies the HasPrefix predicate on the "task_name" field. +func TaskNameHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldTaskName, v)) +} + +// TaskNameHasSuffix applies the HasSuffix predicate on the "task_name" field. +func TaskNameHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldTaskName, v)) +} + +// TaskNameEqualFold applies the EqualFold predicate on the "task_name" field. +func TaskNameEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldTaskName, v)) +} + +// TaskNameContainsFold applies the ContainsFold predicate on the "task_name" field. +func TaskNameContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldTaskName, v)) +} + +// StatusEQ applies the EQ predicate on the "status" field. +func StatusEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldStatus, v)) +} + +// StatusNEQ applies the NEQ predicate on the "status" field. +func StatusNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldStatus, v)) +} + +// StatusIn applies the In predicate on the "status" field. +func StatusIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldStatus, vs...)) +} + +// StatusNotIn applies the NotIn predicate on the "status" field. +func StatusNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldStatus, vs...)) +} + +// StatusGT applies the GT predicate on the "status" field. +func StatusGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldStatus, v)) +} + +// StatusGTE applies the GTE predicate on the "status" field. +func StatusGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldStatus, v)) +} + +// StatusLT applies the LT predicate on the "status" field. +func StatusLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldStatus, v)) +} + +// StatusLTE applies the LTE predicate on the "status" field. +func StatusLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldStatus, v)) +} + +// StatusContains applies the Contains predicate on the "status" field. +func StatusContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldStatus, v)) +} + +// StatusHasPrefix applies the HasPrefix predicate on the "status" field. +func StatusHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldStatus, v)) +} + +// StatusHasSuffix applies the HasSuffix predicate on the "status" field. +func StatusHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldStatus, v)) +} + +// StatusEqualFold applies the EqualFold predicate on the "status" field. +func StatusEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldStatus, v)) +} + +// StatusContainsFold applies the ContainsFold predicate on the "status" field. +func StatusContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldStatus, v)) +} + +// ProviderJobNameEQ applies the EQ predicate on the "provider_job_name" field. +func ProviderJobNameEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProviderJobName, v)) +} + +// ProviderJobNameNEQ applies the NEQ predicate on the "provider_job_name" field. +func ProviderJobNameNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldProviderJobName, v)) +} + +// ProviderJobNameIn applies the In predicate on the "provider_job_name" field. +func ProviderJobNameIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldProviderJobName, vs...)) +} + +// ProviderJobNameNotIn applies the NotIn predicate on the "provider_job_name" field. +func ProviderJobNameNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldProviderJobName, vs...)) +} + +// ProviderJobNameGT applies the GT predicate on the "provider_job_name" field. +func ProviderJobNameGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldProviderJobName, v)) +} + +// ProviderJobNameGTE applies the GTE predicate on the "provider_job_name" field. +func ProviderJobNameGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldProviderJobName, v)) +} + +// ProviderJobNameLT applies the LT predicate on the "provider_job_name" field. +func ProviderJobNameLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldProviderJobName, v)) +} + +// ProviderJobNameLTE applies the LTE predicate on the "provider_job_name" field. +func ProviderJobNameLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldProviderJobName, v)) +} + +// ProviderJobNameContains applies the Contains predicate on the "provider_job_name" field. +func ProviderJobNameContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldProviderJobName, v)) +} + +// ProviderJobNameHasPrefix applies the HasPrefix predicate on the "provider_job_name" field. +func ProviderJobNameHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldProviderJobName, v)) +} + +// ProviderJobNameHasSuffix applies the HasSuffix predicate on the "provider_job_name" field. +func ProviderJobNameHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldProviderJobName, v)) +} + +// ProviderJobNameIsNil applies the IsNil predicate on the "provider_job_name" field. +func ProviderJobNameIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldProviderJobName)) +} + +// ProviderJobNameNotNil applies the NotNil predicate on the "provider_job_name" field. +func ProviderJobNameNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldProviderJobName)) +} + +// ProviderJobNameEqualFold applies the EqualFold predicate on the "provider_job_name" field. +func ProviderJobNameEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldProviderJobName, v)) +} + +// ProviderJobNameContainsFold applies the ContainsFold predicate on the "provider_job_name" field. +func ProviderJobNameContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldProviderJobName, v)) +} + +// ProviderInputRefEQ applies the EQ predicate on the "provider_input_ref" field. +func ProviderInputRefEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProviderInputRef, v)) +} + +// ProviderInputRefNEQ applies the NEQ predicate on the "provider_input_ref" field. +func ProviderInputRefNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldProviderInputRef, v)) +} + +// ProviderInputRefIn applies the In predicate on the "provider_input_ref" field. +func ProviderInputRefIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldProviderInputRef, vs...)) +} + +// ProviderInputRefNotIn applies the NotIn predicate on the "provider_input_ref" field. +func ProviderInputRefNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldProviderInputRef, vs...)) +} + +// ProviderInputRefGT applies the GT predicate on the "provider_input_ref" field. +func ProviderInputRefGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldProviderInputRef, v)) +} + +// ProviderInputRefGTE applies the GTE predicate on the "provider_input_ref" field. +func ProviderInputRefGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldProviderInputRef, v)) +} + +// ProviderInputRefLT applies the LT predicate on the "provider_input_ref" field. +func ProviderInputRefLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldProviderInputRef, v)) +} + +// ProviderInputRefLTE applies the LTE predicate on the "provider_input_ref" field. +func ProviderInputRefLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldProviderInputRef, v)) +} + +// ProviderInputRefContains applies the Contains predicate on the "provider_input_ref" field. +func ProviderInputRefContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldProviderInputRef, v)) +} + +// ProviderInputRefHasPrefix applies the HasPrefix predicate on the "provider_input_ref" field. +func ProviderInputRefHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldProviderInputRef, v)) +} + +// ProviderInputRefHasSuffix applies the HasSuffix predicate on the "provider_input_ref" field. +func ProviderInputRefHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldProviderInputRef, v)) +} + +// ProviderInputRefIsNil applies the IsNil predicate on the "provider_input_ref" field. +func ProviderInputRefIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldProviderInputRef)) +} + +// ProviderInputRefNotNil applies the NotNil predicate on the "provider_input_ref" field. +func ProviderInputRefNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldProviderInputRef)) +} + +// ProviderInputRefEqualFold applies the EqualFold predicate on the "provider_input_ref" field. +func ProviderInputRefEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldProviderInputRef, v)) +} + +// ProviderInputRefContainsFold applies the ContainsFold predicate on the "provider_input_ref" field. +func ProviderInputRefContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldProviderInputRef, v)) +} + +// ProviderOutputRefEQ applies the EQ predicate on the "provider_output_ref" field. +func ProviderOutputRefEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefNEQ applies the NEQ predicate on the "provider_output_ref" field. +func ProviderOutputRefNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefIn applies the In predicate on the "provider_output_ref" field. +func ProviderOutputRefIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldProviderOutputRef, vs...)) +} + +// ProviderOutputRefNotIn applies the NotIn predicate on the "provider_output_ref" field. +func ProviderOutputRefNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldProviderOutputRef, vs...)) +} + +// ProviderOutputRefGT applies the GT predicate on the "provider_output_ref" field. +func ProviderOutputRefGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefGTE applies the GTE predicate on the "provider_output_ref" field. +func ProviderOutputRefGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefLT applies the LT predicate on the "provider_output_ref" field. +func ProviderOutputRefLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefLTE applies the LTE predicate on the "provider_output_ref" field. +func ProviderOutputRefLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefContains applies the Contains predicate on the "provider_output_ref" field. +func ProviderOutputRefContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefHasPrefix applies the HasPrefix predicate on the "provider_output_ref" field. +func ProviderOutputRefHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefHasSuffix applies the HasSuffix predicate on the "provider_output_ref" field. +func ProviderOutputRefHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefIsNil applies the IsNil predicate on the "provider_output_ref" field. +func ProviderOutputRefIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldProviderOutputRef)) +} + +// ProviderOutputRefNotNil applies the NotNil predicate on the "provider_output_ref" field. +func ProviderOutputRefNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldProviderOutputRef)) +} + +// ProviderOutputRefEqualFold applies the EqualFold predicate on the "provider_output_ref" field. +func ProviderOutputRefEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldProviderOutputRef, v)) +} + +// ProviderOutputRefContainsFold applies the ContainsFold predicate on the "provider_output_ref" field. +func ProviderOutputRefContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldProviderOutputRef, v)) +} + +// GcsInputURIEQ applies the EQ predicate on the "gcs_input_uri" field. +func GcsInputURIEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldGcsInputURI, v)) +} + +// GcsInputURINEQ applies the NEQ predicate on the "gcs_input_uri" field. +func GcsInputURINEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldGcsInputURI, v)) +} + +// GcsInputURIIn applies the In predicate on the "gcs_input_uri" field. +func GcsInputURIIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldGcsInputURI, vs...)) +} + +// GcsInputURINotIn applies the NotIn predicate on the "gcs_input_uri" field. +func GcsInputURINotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldGcsInputURI, vs...)) +} + +// GcsInputURIGT applies the GT predicate on the "gcs_input_uri" field. +func GcsInputURIGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldGcsInputURI, v)) +} + +// GcsInputURIGTE applies the GTE predicate on the "gcs_input_uri" field. +func GcsInputURIGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldGcsInputURI, v)) +} + +// GcsInputURILT applies the LT predicate on the "gcs_input_uri" field. +func GcsInputURILT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldGcsInputURI, v)) +} + +// GcsInputURILTE applies the LTE predicate on the "gcs_input_uri" field. +func GcsInputURILTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldGcsInputURI, v)) +} + +// GcsInputURIContains applies the Contains predicate on the "gcs_input_uri" field. +func GcsInputURIContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldGcsInputURI, v)) +} + +// GcsInputURIHasPrefix applies the HasPrefix predicate on the "gcs_input_uri" field. +func GcsInputURIHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldGcsInputURI, v)) +} + +// GcsInputURIHasSuffix applies the HasSuffix predicate on the "gcs_input_uri" field. +func GcsInputURIHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldGcsInputURI, v)) +} + +// GcsInputURIIsNil applies the IsNil predicate on the "gcs_input_uri" field. +func GcsInputURIIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldGcsInputURI)) +} + +// GcsInputURINotNil applies the NotNil predicate on the "gcs_input_uri" field. +func GcsInputURINotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldGcsInputURI)) +} + +// GcsInputURIEqualFold applies the EqualFold predicate on the "gcs_input_uri" field. +func GcsInputURIEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldGcsInputURI, v)) +} + +// GcsInputURIContainsFold applies the ContainsFold predicate on the "gcs_input_uri" field. +func GcsInputURIContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldGcsInputURI, v)) +} + +// GcsOutputURIEQ applies the EQ predicate on the "gcs_output_uri" field. +func GcsOutputURIEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldGcsOutputURI, v)) +} + +// GcsOutputURINEQ applies the NEQ predicate on the "gcs_output_uri" field. +func GcsOutputURINEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldGcsOutputURI, v)) +} + +// GcsOutputURIIn applies the In predicate on the "gcs_output_uri" field. +func GcsOutputURIIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldGcsOutputURI, vs...)) +} + +// GcsOutputURINotIn applies the NotIn predicate on the "gcs_output_uri" field. +func GcsOutputURINotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldGcsOutputURI, vs...)) +} + +// GcsOutputURIGT applies the GT predicate on the "gcs_output_uri" field. +func GcsOutputURIGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldGcsOutputURI, v)) +} + +// GcsOutputURIGTE applies the GTE predicate on the "gcs_output_uri" field. +func GcsOutputURIGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldGcsOutputURI, v)) +} + +// GcsOutputURILT applies the LT predicate on the "gcs_output_uri" field. +func GcsOutputURILT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldGcsOutputURI, v)) +} + +// GcsOutputURILTE applies the LTE predicate on the "gcs_output_uri" field. +func GcsOutputURILTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldGcsOutputURI, v)) +} + +// GcsOutputURIContains applies the Contains predicate on the "gcs_output_uri" field. +func GcsOutputURIContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldGcsOutputURI, v)) +} + +// GcsOutputURIHasPrefix applies the HasPrefix predicate on the "gcs_output_uri" field. +func GcsOutputURIHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldGcsOutputURI, v)) +} + +// GcsOutputURIHasSuffix applies the HasSuffix predicate on the "gcs_output_uri" field. +func GcsOutputURIHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldGcsOutputURI, v)) +} + +// GcsOutputURIIsNil applies the IsNil predicate on the "gcs_output_uri" field. +func GcsOutputURIIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldGcsOutputURI)) +} + +// GcsOutputURINotNil applies the NotNil predicate on the "gcs_output_uri" field. +func GcsOutputURINotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldGcsOutputURI)) +} + +// GcsOutputURIEqualFold applies the EqualFold predicate on the "gcs_output_uri" field. +func GcsOutputURIEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldGcsOutputURI, v)) +} + +// GcsOutputURIContainsFold applies the ContainsFold predicate on the "gcs_output_uri" field. +func GcsOutputURIContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldGcsOutputURI, v)) +} + +// ItemCountEQ applies the EQ predicate on the "item_count" field. +func ItemCountEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldItemCount, v)) +} + +// ItemCountNEQ applies the NEQ predicate on the "item_count" field. +func ItemCountNEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldItemCount, v)) +} + +// ItemCountIn applies the In predicate on the "item_count" field. +func ItemCountIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldItemCount, vs...)) +} + +// ItemCountNotIn applies the NotIn predicate on the "item_count" field. +func ItemCountNotIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldItemCount, vs...)) +} + +// ItemCountGT applies the GT predicate on the "item_count" field. +func ItemCountGT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldItemCount, v)) +} + +// ItemCountGTE applies the GTE predicate on the "item_count" field. +func ItemCountGTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldItemCount, v)) +} + +// ItemCountLT applies the LT predicate on the "item_count" field. +func ItemCountLT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldItemCount, v)) +} + +// ItemCountLTE applies the LTE predicate on the "item_count" field. +func ItemCountLTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldItemCount, v)) +} + +// SuccessCountEQ applies the EQ predicate on the "success_count" field. +func SuccessCountEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldSuccessCount, v)) +} + +// SuccessCountNEQ applies the NEQ predicate on the "success_count" field. +func SuccessCountNEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldSuccessCount, v)) +} + +// SuccessCountIn applies the In predicate on the "success_count" field. +func SuccessCountIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldSuccessCount, vs...)) +} + +// SuccessCountNotIn applies the NotIn predicate on the "success_count" field. +func SuccessCountNotIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldSuccessCount, vs...)) +} + +// SuccessCountGT applies the GT predicate on the "success_count" field. +func SuccessCountGT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldSuccessCount, v)) +} + +// SuccessCountGTE applies the GTE predicate on the "success_count" field. +func SuccessCountGTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldSuccessCount, v)) +} + +// SuccessCountLT applies the LT predicate on the "success_count" field. +func SuccessCountLT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldSuccessCount, v)) +} + +// SuccessCountLTE applies the LTE predicate on the "success_count" field. +func SuccessCountLTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldSuccessCount, v)) +} + +// FailCountEQ applies the EQ predicate on the "fail_count" field. +func FailCountEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldFailCount, v)) +} + +// FailCountNEQ applies the NEQ predicate on the "fail_count" field. +func FailCountNEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldFailCount, v)) +} + +// FailCountIn applies the In predicate on the "fail_count" field. +func FailCountIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldFailCount, vs...)) +} + +// FailCountNotIn applies the NotIn predicate on the "fail_count" field. +func FailCountNotIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldFailCount, vs...)) +} + +// FailCountGT applies the GT predicate on the "fail_count" field. +func FailCountGT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldFailCount, v)) +} + +// FailCountGTE applies the GTE predicate on the "fail_count" field. +func FailCountGTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldFailCount, v)) +} + +// FailCountLT applies the LT predicate on the "fail_count" field. +func FailCountLT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldFailCount, v)) +} + +// FailCountLTE applies the LTE predicate on the "fail_count" field. +func FailCountLTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldFailCount, v)) +} + +// CancelledCountEQ applies the EQ predicate on the "cancelled_count" field. +func CancelledCountEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldCancelledCount, v)) +} + +// CancelledCountNEQ applies the NEQ predicate on the "cancelled_count" field. +func CancelledCountNEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldCancelledCount, v)) +} + +// CancelledCountIn applies the In predicate on the "cancelled_count" field. +func CancelledCountIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldCancelledCount, vs...)) +} + +// CancelledCountNotIn applies the NotIn predicate on the "cancelled_count" field. +func CancelledCountNotIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldCancelledCount, vs...)) +} + +// CancelledCountGT applies the GT predicate on the "cancelled_count" field. +func CancelledCountGT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldCancelledCount, v)) +} + +// CancelledCountGTE applies the GTE predicate on the "cancelled_count" field. +func CancelledCountGTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldCancelledCount, v)) +} + +// CancelledCountLT applies the LT predicate on the "cancelled_count" field. +func CancelledCountLT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldCancelledCount, v)) +} + +// CancelledCountLTE applies the LTE predicate on the "cancelled_count" field. +func CancelledCountLTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldCancelledCount, v)) +} + +// EstimatedCostEQ applies the EQ predicate on the "estimated_cost" field. +func EstimatedCostEQ(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldEstimatedCost, v)) +} + +// EstimatedCostNEQ applies the NEQ predicate on the "estimated_cost" field. +func EstimatedCostNEQ(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldEstimatedCost, v)) +} + +// EstimatedCostIn applies the In predicate on the "estimated_cost" field. +func EstimatedCostIn(vs ...float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldEstimatedCost, vs...)) +} + +// EstimatedCostNotIn applies the NotIn predicate on the "estimated_cost" field. +func EstimatedCostNotIn(vs ...float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldEstimatedCost, vs...)) +} + +// EstimatedCostGT applies the GT predicate on the "estimated_cost" field. +func EstimatedCostGT(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldEstimatedCost, v)) +} + +// EstimatedCostGTE applies the GTE predicate on the "estimated_cost" field. +func EstimatedCostGTE(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldEstimatedCost, v)) +} + +// EstimatedCostLT applies the LT predicate on the "estimated_cost" field. +func EstimatedCostLT(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldEstimatedCost, v)) +} + +// EstimatedCostLTE applies the LTE predicate on the "estimated_cost" field. +func EstimatedCostLTE(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldEstimatedCost, v)) +} + +// HoldAmountEQ applies the EQ predicate on the "hold_amount" field. +func HoldAmountEQ(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldHoldAmount, v)) +} + +// HoldAmountNEQ applies the NEQ predicate on the "hold_amount" field. +func HoldAmountNEQ(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldHoldAmount, v)) +} + +// HoldAmountIn applies the In predicate on the "hold_amount" field. +func HoldAmountIn(vs ...float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldHoldAmount, vs...)) +} + +// HoldAmountNotIn applies the NotIn predicate on the "hold_amount" field. +func HoldAmountNotIn(vs ...float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldHoldAmount, vs...)) +} + +// HoldAmountGT applies the GT predicate on the "hold_amount" field. +func HoldAmountGT(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldHoldAmount, v)) +} + +// HoldAmountGTE applies the GTE predicate on the "hold_amount" field. +func HoldAmountGTE(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldHoldAmount, v)) +} + +// HoldAmountLT applies the LT predicate on the "hold_amount" field. +func HoldAmountLT(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldHoldAmount, v)) +} + +// HoldAmountLTE applies the LTE predicate on the "hold_amount" field. +func HoldAmountLTE(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldHoldAmount, v)) +} + +// HoldAmountIsNil applies the IsNil predicate on the "hold_amount" field. +func HoldAmountIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldHoldAmount)) +} + +// HoldAmountNotNil applies the NotNil predicate on the "hold_amount" field. +func HoldAmountNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldHoldAmount)) +} + +// ActualCostEQ applies the EQ predicate on the "actual_cost" field. +func ActualCostEQ(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldActualCost, v)) +} + +// ActualCostNEQ applies the NEQ predicate on the "actual_cost" field. +func ActualCostNEQ(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldActualCost, v)) +} + +// ActualCostIn applies the In predicate on the "actual_cost" field. +func ActualCostIn(vs ...float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldActualCost, vs...)) +} + +// ActualCostNotIn applies the NotIn predicate on the "actual_cost" field. +func ActualCostNotIn(vs ...float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldActualCost, vs...)) +} + +// ActualCostGT applies the GT predicate on the "actual_cost" field. +func ActualCostGT(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldActualCost, v)) +} + +// ActualCostGTE applies the GTE predicate on the "actual_cost" field. +func ActualCostGTE(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldActualCost, v)) +} + +// ActualCostLT applies the LT predicate on the "actual_cost" field. +func ActualCostLT(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldActualCost, v)) +} + +// ActualCostLTE applies the LTE predicate on the "actual_cost" field. +func ActualCostLTE(v float64) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldActualCost, v)) +} + +// ActualCostIsNil applies the IsNil predicate on the "actual_cost" field. +func ActualCostIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldActualCost)) +} + +// ActualCostNotNil applies the NotNil predicate on the "actual_cost" field. +func ActualCostNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldActualCost)) +} + +// CurrencyEQ applies the EQ predicate on the "currency" field. +func CurrencyEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldCurrency, v)) +} + +// CurrencyNEQ applies the NEQ predicate on the "currency" field. +func CurrencyNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldCurrency, v)) +} + +// CurrencyIn applies the In predicate on the "currency" field. +func CurrencyIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldCurrency, vs...)) +} + +// CurrencyNotIn applies the NotIn predicate on the "currency" field. +func CurrencyNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldCurrency, vs...)) +} + +// CurrencyGT applies the GT predicate on the "currency" field. +func CurrencyGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldCurrency, v)) +} + +// CurrencyGTE applies the GTE predicate on the "currency" field. +func CurrencyGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldCurrency, v)) +} + +// CurrencyLT applies the LT predicate on the "currency" field. +func CurrencyLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldCurrency, v)) +} + +// CurrencyLTE applies the LTE predicate on the "currency" field. +func CurrencyLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldCurrency, v)) +} + +// CurrencyContains applies the Contains predicate on the "currency" field. +func CurrencyContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldCurrency, v)) +} + +// CurrencyHasPrefix applies the HasPrefix predicate on the "currency" field. +func CurrencyHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldCurrency, v)) +} + +// CurrencyHasSuffix applies the HasSuffix predicate on the "currency" field. +func CurrencyHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldCurrency, v)) +} + +// CurrencyEqualFold applies the EqualFold predicate on the "currency" field. +func CurrencyEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldCurrency, v)) +} + +// CurrencyContainsFold applies the ContainsFold predicate on the "currency" field. +func CurrencyContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldCurrency, v)) +} + +// HoldIDEQ applies the EQ predicate on the "hold_id" field. +func HoldIDEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldHoldID, v)) +} + +// HoldIDNEQ applies the NEQ predicate on the "hold_id" field. +func HoldIDNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldHoldID, v)) +} + +// HoldIDIn applies the In predicate on the "hold_id" field. +func HoldIDIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldHoldID, vs...)) +} + +// HoldIDNotIn applies the NotIn predicate on the "hold_id" field. +func HoldIDNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldHoldID, vs...)) +} + +// HoldIDGT applies the GT predicate on the "hold_id" field. +func HoldIDGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldHoldID, v)) +} + +// HoldIDGTE applies the GTE predicate on the "hold_id" field. +func HoldIDGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldHoldID, v)) +} + +// HoldIDLT applies the LT predicate on the "hold_id" field. +func HoldIDLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldHoldID, v)) +} + +// HoldIDLTE applies the LTE predicate on the "hold_id" field. +func HoldIDLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldHoldID, v)) +} + +// HoldIDContains applies the Contains predicate on the "hold_id" field. +func HoldIDContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldHoldID, v)) +} + +// HoldIDHasPrefix applies the HasPrefix predicate on the "hold_id" field. +func HoldIDHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldHoldID, v)) +} + +// HoldIDHasSuffix applies the HasSuffix predicate on the "hold_id" field. +func HoldIDHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldHoldID, v)) +} + +// HoldIDIsNil applies the IsNil predicate on the "hold_id" field. +func HoldIDIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldHoldID)) +} + +// HoldIDNotNil applies the NotNil predicate on the "hold_id" field. +func HoldIDNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldHoldID)) +} + +// HoldIDEqualFold applies the EqualFold predicate on the "hold_id" field. +func HoldIDEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldHoldID, v)) +} + +// HoldIDContainsFold applies the ContainsFold predicate on the "hold_id" field. +func HoldIDContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldHoldID, v)) +} + +// IdempotencyKeyEQ applies the EQ predicate on the "idempotency_key" field. +func IdempotencyKeyEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyNEQ applies the NEQ predicate on the "idempotency_key" field. +func IdempotencyKeyNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyIn applies the In predicate on the "idempotency_key" field. +func IdempotencyKeyIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldIdempotencyKey, vs...)) +} + +// IdempotencyKeyNotIn applies the NotIn predicate on the "idempotency_key" field. +func IdempotencyKeyNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldIdempotencyKey, vs...)) +} + +// IdempotencyKeyGT applies the GT predicate on the "idempotency_key" field. +func IdempotencyKeyGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyGTE applies the GTE predicate on the "idempotency_key" field. +func IdempotencyKeyGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyLT applies the LT predicate on the "idempotency_key" field. +func IdempotencyKeyLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyLTE applies the LTE predicate on the "idempotency_key" field. +func IdempotencyKeyLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyContains applies the Contains predicate on the "idempotency_key" field. +func IdempotencyKeyContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyHasPrefix applies the HasPrefix predicate on the "idempotency_key" field. +func IdempotencyKeyHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyHasSuffix applies the HasSuffix predicate on the "idempotency_key" field. +func IdempotencyKeyHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyIsNil applies the IsNil predicate on the "idempotency_key" field. +func IdempotencyKeyIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldIdempotencyKey)) +} + +// IdempotencyKeyNotNil applies the NotNil predicate on the "idempotency_key" field. +func IdempotencyKeyNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldIdempotencyKey)) +} + +// IdempotencyKeyEqualFold applies the EqualFold predicate on the "idempotency_key" field. +func IdempotencyKeyEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldIdempotencyKey, v)) +} + +// IdempotencyKeyContainsFold applies the ContainsFold predicate on the "idempotency_key" field. +func IdempotencyKeyContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldIdempotencyKey, v)) +} + +// RequestHashEQ applies the EQ predicate on the "request_hash" field. +func RequestHashEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldRequestHash, v)) +} + +// RequestHashNEQ applies the NEQ predicate on the "request_hash" field. +func RequestHashNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldRequestHash, v)) +} + +// RequestHashIn applies the In predicate on the "request_hash" field. +func RequestHashIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldRequestHash, vs...)) +} + +// RequestHashNotIn applies the NotIn predicate on the "request_hash" field. +func RequestHashNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldRequestHash, vs...)) +} + +// RequestHashGT applies the GT predicate on the "request_hash" field. +func RequestHashGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldRequestHash, v)) +} + +// RequestHashGTE applies the GTE predicate on the "request_hash" field. +func RequestHashGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldRequestHash, v)) +} + +// RequestHashLT applies the LT predicate on the "request_hash" field. +func RequestHashLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldRequestHash, v)) +} + +// RequestHashLTE applies the LTE predicate on the "request_hash" field. +func RequestHashLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldRequestHash, v)) +} + +// RequestHashContains applies the Contains predicate on the "request_hash" field. +func RequestHashContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldRequestHash, v)) +} + +// RequestHashHasPrefix applies the HasPrefix predicate on the "request_hash" field. +func RequestHashHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldRequestHash, v)) +} + +// RequestHashHasSuffix applies the HasSuffix predicate on the "request_hash" field. +func RequestHashHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldRequestHash, v)) +} + +// RequestHashIsNil applies the IsNil predicate on the "request_hash" field. +func RequestHashIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldRequestHash)) +} + +// RequestHashNotNil applies the NotNil predicate on the "request_hash" field. +func RequestHashNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldRequestHash)) +} + +// RequestHashEqualFold applies the EqualFold predicate on the "request_hash" field. +func RequestHashEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldRequestHash, v)) +} + +// RequestHashContainsFold applies the ContainsFold predicate on the "request_hash" field. +func RequestHashContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldRequestHash, v)) +} + +// ManifestHashEQ applies the EQ predicate on the "manifest_hash" field. +func ManifestHashEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldManifestHash, v)) +} + +// ManifestHashNEQ applies the NEQ predicate on the "manifest_hash" field. +func ManifestHashNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldManifestHash, v)) +} + +// ManifestHashIn applies the In predicate on the "manifest_hash" field. +func ManifestHashIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldManifestHash, vs...)) +} + +// ManifestHashNotIn applies the NotIn predicate on the "manifest_hash" field. +func ManifestHashNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldManifestHash, vs...)) +} + +// ManifestHashGT applies the GT predicate on the "manifest_hash" field. +func ManifestHashGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldManifestHash, v)) +} + +// ManifestHashGTE applies the GTE predicate on the "manifest_hash" field. +func ManifestHashGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldManifestHash, v)) +} + +// ManifestHashLT applies the LT predicate on the "manifest_hash" field. +func ManifestHashLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldManifestHash, v)) +} + +// ManifestHashLTE applies the LTE predicate on the "manifest_hash" field. +func ManifestHashLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldManifestHash, v)) +} + +// ManifestHashContains applies the Contains predicate on the "manifest_hash" field. +func ManifestHashContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldManifestHash, v)) +} + +// ManifestHashHasPrefix applies the HasPrefix predicate on the "manifest_hash" field. +func ManifestHashHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldManifestHash, v)) +} + +// ManifestHashHasSuffix applies the HasSuffix predicate on the "manifest_hash" field. +func ManifestHashHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldManifestHash, v)) +} + +// ManifestHashIsNil applies the IsNil predicate on the "manifest_hash" field. +func ManifestHashIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldManifestHash)) +} + +// ManifestHashNotNil applies the NotNil predicate on the "manifest_hash" field. +func ManifestHashNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldManifestHash)) +} + +// ManifestHashEqualFold applies the EqualFold predicate on the "manifest_hash" field. +func ManifestHashEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldManifestHash, v)) +} + +// ManifestHashContainsFold applies the ContainsFold predicate on the "manifest_hash" field. +func ManifestHashContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldManifestHash, v)) +} + +// RetryCountEQ applies the EQ predicate on the "retry_count" field. +func RetryCountEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldRetryCount, v)) +} + +// RetryCountNEQ applies the NEQ predicate on the "retry_count" field. +func RetryCountNEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldRetryCount, v)) +} + +// RetryCountIn applies the In predicate on the "retry_count" field. +func RetryCountIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldRetryCount, vs...)) +} + +// RetryCountNotIn applies the NotIn predicate on the "retry_count" field. +func RetryCountNotIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldRetryCount, vs...)) +} + +// RetryCountGT applies the GT predicate on the "retry_count" field. +func RetryCountGT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldRetryCount, v)) +} + +// RetryCountGTE applies the GTE predicate on the "retry_count" field. +func RetryCountGTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldRetryCount, v)) +} + +// RetryCountLT applies the LT predicate on the "retry_count" field. +func RetryCountLT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldRetryCount, v)) +} + +// RetryCountLTE applies the LTE predicate on the "retry_count" field. +func RetryCountLTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldRetryCount, v)) +} + +// VersionEQ applies the EQ predicate on the "version" field. +func VersionEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldVersion, v)) +} + +// VersionNEQ applies the NEQ predicate on the "version" field. +func VersionNEQ(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldVersion, v)) +} + +// VersionIn applies the In predicate on the "version" field. +func VersionIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldVersion, vs...)) +} + +// VersionNotIn applies the NotIn predicate on the "version" field. +func VersionNotIn(vs ...int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldVersion, vs...)) +} + +// VersionGT applies the GT predicate on the "version" field. +func VersionGT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldVersion, v)) +} + +// VersionGTE applies the GTE predicate on the "version" field. +func VersionGTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldVersion, v)) +} + +// VersionLT applies the LT predicate on the "version" field. +func VersionLT(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldVersion, v)) +} + +// VersionLTE applies the LTE predicate on the "version" field. +func VersionLTE(v int) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldVersion, v)) +} + +// OutputExpiresAtEQ applies the EQ predicate on the "output_expires_at" field. +func OutputExpiresAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldOutputExpiresAt, v)) +} + +// OutputExpiresAtNEQ applies the NEQ predicate on the "output_expires_at" field. +func OutputExpiresAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldOutputExpiresAt, v)) +} + +// OutputExpiresAtIn applies the In predicate on the "output_expires_at" field. +func OutputExpiresAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldOutputExpiresAt, vs...)) +} + +// OutputExpiresAtNotIn applies the NotIn predicate on the "output_expires_at" field. +func OutputExpiresAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldOutputExpiresAt, vs...)) +} + +// OutputExpiresAtGT applies the GT predicate on the "output_expires_at" field. +func OutputExpiresAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldOutputExpiresAt, v)) +} + +// OutputExpiresAtGTE applies the GTE predicate on the "output_expires_at" field. +func OutputExpiresAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldOutputExpiresAt, v)) +} + +// OutputExpiresAtLT applies the LT predicate on the "output_expires_at" field. +func OutputExpiresAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldOutputExpiresAt, v)) +} + +// OutputExpiresAtLTE applies the LTE predicate on the "output_expires_at" field. +func OutputExpiresAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldOutputExpiresAt, v)) +} + +// OutputExpiresAtIsNil applies the IsNil predicate on the "output_expires_at" field. +func OutputExpiresAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldOutputExpiresAt)) +} + +// OutputExpiresAtNotNil applies the NotNil predicate on the "output_expires_at" field. +func OutputExpiresAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldOutputExpiresAt)) +} + +// InputDeletedAtEQ applies the EQ predicate on the "input_deleted_at" field. +func InputDeletedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldInputDeletedAt, v)) +} + +// InputDeletedAtNEQ applies the NEQ predicate on the "input_deleted_at" field. +func InputDeletedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldInputDeletedAt, v)) +} + +// InputDeletedAtIn applies the In predicate on the "input_deleted_at" field. +func InputDeletedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldInputDeletedAt, vs...)) +} + +// InputDeletedAtNotIn applies the NotIn predicate on the "input_deleted_at" field. +func InputDeletedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldInputDeletedAt, vs...)) +} + +// InputDeletedAtGT applies the GT predicate on the "input_deleted_at" field. +func InputDeletedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldInputDeletedAt, v)) +} + +// InputDeletedAtGTE applies the GTE predicate on the "input_deleted_at" field. +func InputDeletedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldInputDeletedAt, v)) +} + +// InputDeletedAtLT applies the LT predicate on the "input_deleted_at" field. +func InputDeletedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldInputDeletedAt, v)) +} + +// InputDeletedAtLTE applies the LTE predicate on the "input_deleted_at" field. +func InputDeletedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldInputDeletedAt, v)) +} + +// InputDeletedAtIsNil applies the IsNil predicate on the "input_deleted_at" field. +func InputDeletedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldInputDeletedAt)) +} + +// InputDeletedAtNotNil applies the NotNil predicate on the "input_deleted_at" field. +func InputDeletedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldInputDeletedAt)) +} + +// OutputDeletedAtEQ applies the EQ predicate on the "output_deleted_at" field. +func OutputDeletedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldOutputDeletedAt, v)) +} + +// OutputDeletedAtNEQ applies the NEQ predicate on the "output_deleted_at" field. +func OutputDeletedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldOutputDeletedAt, v)) +} + +// OutputDeletedAtIn applies the In predicate on the "output_deleted_at" field. +func OutputDeletedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldOutputDeletedAt, vs...)) +} + +// OutputDeletedAtNotIn applies the NotIn predicate on the "output_deleted_at" field. +func OutputDeletedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldOutputDeletedAt, vs...)) +} + +// OutputDeletedAtGT applies the GT predicate on the "output_deleted_at" field. +func OutputDeletedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldOutputDeletedAt, v)) +} + +// OutputDeletedAtGTE applies the GTE predicate on the "output_deleted_at" field. +func OutputDeletedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldOutputDeletedAt, v)) +} + +// OutputDeletedAtLT applies the LT predicate on the "output_deleted_at" field. +func OutputDeletedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldOutputDeletedAt, v)) +} + +// OutputDeletedAtLTE applies the LTE predicate on the "output_deleted_at" field. +func OutputDeletedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldOutputDeletedAt, v)) +} + +// OutputDeletedAtIsNil applies the IsNil predicate on the "output_deleted_at" field. +func OutputDeletedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldOutputDeletedAt)) +} + +// OutputDeletedAtNotNil applies the NotNil predicate on the "output_deleted_at" field. +func OutputDeletedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldOutputDeletedAt)) +} + +// DownloadedAtEQ applies the EQ predicate on the "downloaded_at" field. +func DownloadedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldDownloadedAt, v)) +} + +// DownloadedAtNEQ applies the NEQ predicate on the "downloaded_at" field. +func DownloadedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldDownloadedAt, v)) +} + +// DownloadedAtIn applies the In predicate on the "downloaded_at" field. +func DownloadedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldDownloadedAt, vs...)) +} + +// DownloadedAtNotIn applies the NotIn predicate on the "downloaded_at" field. +func DownloadedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldDownloadedAt, vs...)) +} + +// DownloadedAtGT applies the GT predicate on the "downloaded_at" field. +func DownloadedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldDownloadedAt, v)) +} + +// DownloadedAtGTE applies the GTE predicate on the "downloaded_at" field. +func DownloadedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldDownloadedAt, v)) +} + +// DownloadedAtLT applies the LT predicate on the "downloaded_at" field. +func DownloadedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldDownloadedAt, v)) +} + +// DownloadedAtLTE applies the LTE predicate on the "downloaded_at" field. +func DownloadedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldDownloadedAt, v)) +} + +// DownloadedAtIsNil applies the IsNil predicate on the "downloaded_at" field. +func DownloadedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldDownloadedAt)) +} + +// DownloadedAtNotNil applies the NotNil predicate on the "downloaded_at" field. +func DownloadedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldDownloadedAt)) +} + +// UserDeletedAtEQ applies the EQ predicate on the "user_deleted_at" field. +func UserDeletedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldUserDeletedAt, v)) +} + +// UserDeletedAtNEQ applies the NEQ predicate on the "user_deleted_at" field. +func UserDeletedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldUserDeletedAt, v)) +} + +// UserDeletedAtIn applies the In predicate on the "user_deleted_at" field. +func UserDeletedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldUserDeletedAt, vs...)) +} + +// UserDeletedAtNotIn applies the NotIn predicate on the "user_deleted_at" field. +func UserDeletedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldUserDeletedAt, vs...)) +} + +// UserDeletedAtGT applies the GT predicate on the "user_deleted_at" field. +func UserDeletedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldUserDeletedAt, v)) +} + +// UserDeletedAtGTE applies the GTE predicate on the "user_deleted_at" field. +func UserDeletedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldUserDeletedAt, v)) +} + +// UserDeletedAtLT applies the LT predicate on the "user_deleted_at" field. +func UserDeletedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldUserDeletedAt, v)) +} + +// UserDeletedAtLTE applies the LTE predicate on the "user_deleted_at" field. +func UserDeletedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldUserDeletedAt, v)) +} + +// UserDeletedAtIsNil applies the IsNil predicate on the "user_deleted_at" field. +func UserDeletedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldUserDeletedAt)) +} + +// UserDeletedAtNotNil applies the NotNil predicate on the "user_deleted_at" field. +func UserDeletedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldUserDeletedAt)) +} + +// LastErrorCodeEQ applies the EQ predicate on the "last_error_code" field. +func LastErrorCodeEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldLastErrorCode, v)) +} + +// LastErrorCodeNEQ applies the NEQ predicate on the "last_error_code" field. +func LastErrorCodeNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldLastErrorCode, v)) +} + +// LastErrorCodeIn applies the In predicate on the "last_error_code" field. +func LastErrorCodeIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldLastErrorCode, vs...)) +} + +// LastErrorCodeNotIn applies the NotIn predicate on the "last_error_code" field. +func LastErrorCodeNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldLastErrorCode, vs...)) +} + +// LastErrorCodeGT applies the GT predicate on the "last_error_code" field. +func LastErrorCodeGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldLastErrorCode, v)) +} + +// LastErrorCodeGTE applies the GTE predicate on the "last_error_code" field. +func LastErrorCodeGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldLastErrorCode, v)) +} + +// LastErrorCodeLT applies the LT predicate on the "last_error_code" field. +func LastErrorCodeLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldLastErrorCode, v)) +} + +// LastErrorCodeLTE applies the LTE predicate on the "last_error_code" field. +func LastErrorCodeLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldLastErrorCode, v)) +} + +// LastErrorCodeContains applies the Contains predicate on the "last_error_code" field. +func LastErrorCodeContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldLastErrorCode, v)) +} + +// LastErrorCodeHasPrefix applies the HasPrefix predicate on the "last_error_code" field. +func LastErrorCodeHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldLastErrorCode, v)) +} + +// LastErrorCodeHasSuffix applies the HasSuffix predicate on the "last_error_code" field. +func LastErrorCodeHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldLastErrorCode, v)) +} + +// LastErrorCodeIsNil applies the IsNil predicate on the "last_error_code" field. +func LastErrorCodeIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldLastErrorCode)) +} + +// LastErrorCodeNotNil applies the NotNil predicate on the "last_error_code" field. +func LastErrorCodeNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldLastErrorCode)) +} + +// LastErrorCodeEqualFold applies the EqualFold predicate on the "last_error_code" field. +func LastErrorCodeEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldLastErrorCode, v)) +} + +// LastErrorCodeContainsFold applies the ContainsFold predicate on the "last_error_code" field. +func LastErrorCodeContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldLastErrorCode, v)) +} + +// LastErrorMessageEQ applies the EQ predicate on the "last_error_message" field. +func LastErrorMessageEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldLastErrorMessage, v)) +} + +// LastErrorMessageNEQ applies the NEQ predicate on the "last_error_message" field. +func LastErrorMessageNEQ(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldLastErrorMessage, v)) +} + +// LastErrorMessageIn applies the In predicate on the "last_error_message" field. +func LastErrorMessageIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldLastErrorMessage, vs...)) +} + +// LastErrorMessageNotIn applies the NotIn predicate on the "last_error_message" field. +func LastErrorMessageNotIn(vs ...string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldLastErrorMessage, vs...)) +} + +// LastErrorMessageGT applies the GT predicate on the "last_error_message" field. +func LastErrorMessageGT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldLastErrorMessage, v)) +} + +// LastErrorMessageGTE applies the GTE predicate on the "last_error_message" field. +func LastErrorMessageGTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldLastErrorMessage, v)) +} + +// LastErrorMessageLT applies the LT predicate on the "last_error_message" field. +func LastErrorMessageLT(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldLastErrorMessage, v)) +} + +// LastErrorMessageLTE applies the LTE predicate on the "last_error_message" field. +func LastErrorMessageLTE(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldLastErrorMessage, v)) +} + +// LastErrorMessageContains applies the Contains predicate on the "last_error_message" field. +func LastErrorMessageContains(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContains(FieldLastErrorMessage, v)) +} + +// LastErrorMessageHasPrefix applies the HasPrefix predicate on the "last_error_message" field. +func LastErrorMessageHasPrefix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasPrefix(FieldLastErrorMessage, v)) +} + +// LastErrorMessageHasSuffix applies the HasSuffix predicate on the "last_error_message" field. +func LastErrorMessageHasSuffix(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldHasSuffix(FieldLastErrorMessage, v)) +} + +// LastErrorMessageIsNil applies the IsNil predicate on the "last_error_message" field. +func LastErrorMessageIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldLastErrorMessage)) +} + +// LastErrorMessageNotNil applies the NotNil predicate on the "last_error_message" field. +func LastErrorMessageNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldLastErrorMessage)) +} + +// LastErrorMessageEqualFold applies the EqualFold predicate on the "last_error_message" field. +func LastErrorMessageEqualFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEqualFold(FieldLastErrorMessage, v)) +} + +// LastErrorMessageContainsFold applies the ContainsFold predicate on the "last_error_message" field. +func LastErrorMessageContainsFold(v string) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldContainsFold(FieldLastErrorMessage, v)) +} + +// CreatedAtEQ applies the EQ predicate on the "created_at" field. +func CreatedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldCreatedAt, v)) +} + +// CreatedAtNEQ applies the NEQ predicate on the "created_at" field. +func CreatedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldCreatedAt, v)) +} + +// CreatedAtIn applies the In predicate on the "created_at" field. +func CreatedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldCreatedAt, vs...)) +} + +// CreatedAtNotIn applies the NotIn predicate on the "created_at" field. +func CreatedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldCreatedAt, vs...)) +} + +// CreatedAtGT applies the GT predicate on the "created_at" field. +func CreatedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldCreatedAt, v)) +} + +// CreatedAtGTE applies the GTE predicate on the "created_at" field. +func CreatedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldCreatedAt, v)) +} + +// CreatedAtLT applies the LT predicate on the "created_at" field. +func CreatedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldCreatedAt, v)) +} + +// CreatedAtLTE applies the LTE predicate on the "created_at" field. +func CreatedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldCreatedAt, v)) +} + +// UpdatedAtEQ applies the EQ predicate on the "updated_at" field. +func UpdatedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldUpdatedAt, v)) +} + +// UpdatedAtNEQ applies the NEQ predicate on the "updated_at" field. +func UpdatedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldUpdatedAt, v)) +} + +// UpdatedAtIn applies the In predicate on the "updated_at" field. +func UpdatedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldUpdatedAt, vs...)) +} + +// UpdatedAtNotIn applies the NotIn predicate on the "updated_at" field. +func UpdatedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldUpdatedAt, vs...)) +} + +// UpdatedAtGT applies the GT predicate on the "updated_at" field. +func UpdatedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldUpdatedAt, v)) +} + +// UpdatedAtGTE applies the GTE predicate on the "updated_at" field. +func UpdatedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldUpdatedAt, v)) +} + +// UpdatedAtLT applies the LT predicate on the "updated_at" field. +func UpdatedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldUpdatedAt, v)) +} + +// UpdatedAtLTE applies the LTE predicate on the "updated_at" field. +func UpdatedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldUpdatedAt, v)) +} + +// SubmittedAtEQ applies the EQ predicate on the "submitted_at" field. +func SubmittedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldSubmittedAt, v)) +} + +// SubmittedAtNEQ applies the NEQ predicate on the "submitted_at" field. +func SubmittedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldSubmittedAt, v)) +} + +// SubmittedAtIn applies the In predicate on the "submitted_at" field. +func SubmittedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldSubmittedAt, vs...)) +} + +// SubmittedAtNotIn applies the NotIn predicate on the "submitted_at" field. +func SubmittedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldSubmittedAt, vs...)) +} + +// SubmittedAtGT applies the GT predicate on the "submitted_at" field. +func SubmittedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldSubmittedAt, v)) +} + +// SubmittedAtGTE applies the GTE predicate on the "submitted_at" field. +func SubmittedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldSubmittedAt, v)) +} + +// SubmittedAtLT applies the LT predicate on the "submitted_at" field. +func SubmittedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldSubmittedAt, v)) +} + +// SubmittedAtLTE applies the LTE predicate on the "submitted_at" field. +func SubmittedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldSubmittedAt, v)) +} + +// SubmittedAtIsNil applies the IsNil predicate on the "submitted_at" field. +func SubmittedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldSubmittedAt)) +} + +// SubmittedAtNotNil applies the NotNil predicate on the "submitted_at" field. +func SubmittedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldSubmittedAt)) +} + +// StartedAtEQ applies the EQ predicate on the "started_at" field. +func StartedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldStartedAt, v)) +} + +// StartedAtNEQ applies the NEQ predicate on the "started_at" field. +func StartedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldStartedAt, v)) +} + +// StartedAtIn applies the In predicate on the "started_at" field. +func StartedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldStartedAt, vs...)) +} + +// StartedAtNotIn applies the NotIn predicate on the "started_at" field. +func StartedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldStartedAt, vs...)) +} + +// StartedAtGT applies the GT predicate on the "started_at" field. +func StartedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldStartedAt, v)) +} + +// StartedAtGTE applies the GTE predicate on the "started_at" field. +func StartedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldStartedAt, v)) +} + +// StartedAtLT applies the LT predicate on the "started_at" field. +func StartedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldStartedAt, v)) +} + +// StartedAtLTE applies the LTE predicate on the "started_at" field. +func StartedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldStartedAt, v)) +} + +// StartedAtIsNil applies the IsNil predicate on the "started_at" field. +func StartedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldStartedAt)) +} + +// StartedAtNotNil applies the NotNil predicate on the "started_at" field. +func StartedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldStartedAt)) +} + +// FinishedAtEQ applies the EQ predicate on the "finished_at" field. +func FinishedAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldFinishedAt, v)) +} + +// FinishedAtNEQ applies the NEQ predicate on the "finished_at" field. +func FinishedAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldFinishedAt, v)) +} + +// FinishedAtIn applies the In predicate on the "finished_at" field. +func FinishedAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldFinishedAt, vs...)) +} + +// FinishedAtNotIn applies the NotIn predicate on the "finished_at" field. +func FinishedAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldFinishedAt, vs...)) +} + +// FinishedAtGT applies the GT predicate on the "finished_at" field. +func FinishedAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldFinishedAt, v)) +} + +// FinishedAtGTE applies the GTE predicate on the "finished_at" field. +func FinishedAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldFinishedAt, v)) +} + +// FinishedAtLT applies the LT predicate on the "finished_at" field. +func FinishedAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldFinishedAt, v)) +} + +// FinishedAtLTE applies the LTE predicate on the "finished_at" field. +func FinishedAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldFinishedAt, v)) +} + +// FinishedAtIsNil applies the IsNil predicate on the "finished_at" field. +func FinishedAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldFinishedAt)) +} + +// FinishedAtNotNil applies the NotNil predicate on the "finished_at" field. +func FinishedAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldFinishedAt)) +} + +// SettledAtEQ applies the EQ predicate on the "settled_at" field. +func SettledAtEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldEQ(FieldSettledAt, v)) +} + +// SettledAtNEQ applies the NEQ predicate on the "settled_at" field. +func SettledAtNEQ(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNEQ(FieldSettledAt, v)) +} + +// SettledAtIn applies the In predicate on the "settled_at" field. +func SettledAtIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIn(FieldSettledAt, vs...)) +} + +// SettledAtNotIn applies the NotIn predicate on the "settled_at" field. +func SettledAtNotIn(vs ...time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotIn(FieldSettledAt, vs...)) +} + +// SettledAtGT applies the GT predicate on the "settled_at" field. +func SettledAtGT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGT(FieldSettledAt, v)) +} + +// SettledAtGTE applies the GTE predicate on the "settled_at" field. +func SettledAtGTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldGTE(FieldSettledAt, v)) +} + +// SettledAtLT applies the LT predicate on the "settled_at" field. +func SettledAtLT(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLT(FieldSettledAt, v)) +} + +// SettledAtLTE applies the LTE predicate on the "settled_at" field. +func SettledAtLTE(v time.Time) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldLTE(FieldSettledAt, v)) +} + +// SettledAtIsNil applies the IsNil predicate on the "settled_at" field. +func SettledAtIsNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldIsNull(FieldSettledAt)) +} + +// SettledAtNotNil applies the NotNil predicate on the "settled_at" field. +func SettledAtNotNil() predicate.BatchImageJob { + return predicate.BatchImageJob(sql.FieldNotNull(FieldSettledAt)) +} + +// And groups predicates with the AND operator between them. +func And(predicates ...predicate.BatchImageJob) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.AndPredicates(predicates...)) +} + +// Or groups predicates with the OR operator between them. +func Or(predicates ...predicate.BatchImageJob) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.OrPredicates(predicates...)) +} + +// Not applies the not operator on the given predicate. +func Not(p predicate.BatchImageJob) predicate.BatchImageJob { + return predicate.BatchImageJob(sql.NotPredicates(p)) +} diff --git a/backend/ent/batchimagejob_create.go b/backend/ent/batchimagejob_create.go new file mode 100644 index 0000000000..88c1197b15 --- /dev/null +++ b/backend/ent/batchimagejob_create.go @@ -0,0 +1,3518 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "errors" + "fmt" + "time" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" +) + +// BatchImageJobCreate is the builder for creating a BatchImageJob entity. +type BatchImageJobCreate struct { + config + mutation *BatchImageJobMutation + hooks []Hook + conflict []sql.ConflictOption +} + +// SetBatchID sets the "batch_id" field. +func (_c *BatchImageJobCreate) SetBatchID(v string) *BatchImageJobCreate { + _c.mutation.SetBatchID(v) + return _c +} + +// SetUserID sets the "user_id" field. +func (_c *BatchImageJobCreate) SetUserID(v int64) *BatchImageJobCreate { + _c.mutation.SetUserID(v) + return _c +} + +// SetAPIKeyID sets the "api_key_id" field. +func (_c *BatchImageJobCreate) SetAPIKeyID(v int64) *BatchImageJobCreate { + _c.mutation.SetAPIKeyID(v) + return _c +} + +// SetNillableAPIKeyID sets the "api_key_id" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableAPIKeyID(v *int64) *BatchImageJobCreate { + if v != nil { + _c.SetAPIKeyID(*v) + } + return _c +} + +// SetAccountID sets the "account_id" field. +func (_c *BatchImageJobCreate) SetAccountID(v int64) *BatchImageJobCreate { + _c.mutation.SetAccountID(v) + return _c +} + +// SetNillableAccountID sets the "account_id" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableAccountID(v *int64) *BatchImageJobCreate { + if v != nil { + _c.SetAccountID(*v) + } + return _c +} + +// SetProvider sets the "provider" field. +func (_c *BatchImageJobCreate) SetProvider(v string) *BatchImageJobCreate { + _c.mutation.SetProvider(v) + return _c +} + +// SetModel sets the "model" field. +func (_c *BatchImageJobCreate) SetModel(v string) *BatchImageJobCreate { + _c.mutation.SetModel(v) + return _c +} + +// SetTaskName sets the "task_name" field. +func (_c *BatchImageJobCreate) SetTaskName(v string) *BatchImageJobCreate { + _c.mutation.SetTaskName(v) + return _c +} + +// SetNillableTaskName sets the "task_name" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableTaskName(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetTaskName(*v) + } + return _c +} + +// SetStatus sets the "status" field. +func (_c *BatchImageJobCreate) SetStatus(v string) *BatchImageJobCreate { + _c.mutation.SetStatus(v) + return _c +} + +// SetNillableStatus sets the "status" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableStatus(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetStatus(*v) + } + return _c +} + +// SetProviderJobName sets the "provider_job_name" field. +func (_c *BatchImageJobCreate) SetProviderJobName(v string) *BatchImageJobCreate { + _c.mutation.SetProviderJobName(v) + return _c +} + +// SetNillableProviderJobName sets the "provider_job_name" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableProviderJobName(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetProviderJobName(*v) + } + return _c +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (_c *BatchImageJobCreate) SetProviderInputRef(v string) *BatchImageJobCreate { + _c.mutation.SetProviderInputRef(v) + return _c +} + +// SetNillableProviderInputRef sets the "provider_input_ref" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableProviderInputRef(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetProviderInputRef(*v) + } + return _c +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (_c *BatchImageJobCreate) SetProviderOutputRef(v string) *BatchImageJobCreate { + _c.mutation.SetProviderOutputRef(v) + return _c +} + +// SetNillableProviderOutputRef sets the "provider_output_ref" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableProviderOutputRef(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetProviderOutputRef(*v) + } + return _c +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (_c *BatchImageJobCreate) SetGcsInputURI(v string) *BatchImageJobCreate { + _c.mutation.SetGcsInputURI(v) + return _c +} + +// SetNillableGcsInputURI sets the "gcs_input_uri" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableGcsInputURI(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetGcsInputURI(*v) + } + return _c +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (_c *BatchImageJobCreate) SetGcsOutputURI(v string) *BatchImageJobCreate { + _c.mutation.SetGcsOutputURI(v) + return _c +} + +// SetNillableGcsOutputURI sets the "gcs_output_uri" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableGcsOutputURI(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetGcsOutputURI(*v) + } + return _c +} + +// SetItemCount sets the "item_count" field. +func (_c *BatchImageJobCreate) SetItemCount(v int) *BatchImageJobCreate { + _c.mutation.SetItemCount(v) + return _c +} + +// SetSuccessCount sets the "success_count" field. +func (_c *BatchImageJobCreate) SetSuccessCount(v int) *BatchImageJobCreate { + _c.mutation.SetSuccessCount(v) + return _c +} + +// SetNillableSuccessCount sets the "success_count" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableSuccessCount(v *int) *BatchImageJobCreate { + if v != nil { + _c.SetSuccessCount(*v) + } + return _c +} + +// SetFailCount sets the "fail_count" field. +func (_c *BatchImageJobCreate) SetFailCount(v int) *BatchImageJobCreate { + _c.mutation.SetFailCount(v) + return _c +} + +// SetNillableFailCount sets the "fail_count" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableFailCount(v *int) *BatchImageJobCreate { + if v != nil { + _c.SetFailCount(*v) + } + return _c +} + +// SetCancelledCount sets the "cancelled_count" field. +func (_c *BatchImageJobCreate) SetCancelledCount(v int) *BatchImageJobCreate { + _c.mutation.SetCancelledCount(v) + return _c +} + +// SetNillableCancelledCount sets the "cancelled_count" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableCancelledCount(v *int) *BatchImageJobCreate { + if v != nil { + _c.SetCancelledCount(*v) + } + return _c +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (_c *BatchImageJobCreate) SetEstimatedCost(v float64) *BatchImageJobCreate { + _c.mutation.SetEstimatedCost(v) + return _c +} + +// SetNillableEstimatedCost sets the "estimated_cost" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableEstimatedCost(v *float64) *BatchImageJobCreate { + if v != nil { + _c.SetEstimatedCost(*v) + } + return _c +} + +// SetHoldAmount sets the "hold_amount" field. +func (_c *BatchImageJobCreate) SetHoldAmount(v float64) *BatchImageJobCreate { + _c.mutation.SetHoldAmount(v) + return _c +} + +// SetNillableHoldAmount sets the "hold_amount" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableHoldAmount(v *float64) *BatchImageJobCreate { + if v != nil { + _c.SetHoldAmount(*v) + } + return _c +} + +// SetActualCost sets the "actual_cost" field. +func (_c *BatchImageJobCreate) SetActualCost(v float64) *BatchImageJobCreate { + _c.mutation.SetActualCost(v) + return _c +} + +// SetNillableActualCost sets the "actual_cost" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableActualCost(v *float64) *BatchImageJobCreate { + if v != nil { + _c.SetActualCost(*v) + } + return _c +} + +// SetCurrency sets the "currency" field. +func (_c *BatchImageJobCreate) SetCurrency(v string) *BatchImageJobCreate { + _c.mutation.SetCurrency(v) + return _c +} + +// SetNillableCurrency sets the "currency" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableCurrency(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetCurrency(*v) + } + return _c +} + +// SetHoldID sets the "hold_id" field. +func (_c *BatchImageJobCreate) SetHoldID(v string) *BatchImageJobCreate { + _c.mutation.SetHoldID(v) + return _c +} + +// SetNillableHoldID sets the "hold_id" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableHoldID(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetHoldID(*v) + } + return _c +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (_c *BatchImageJobCreate) SetIdempotencyKey(v string) *BatchImageJobCreate { + _c.mutation.SetIdempotencyKey(v) + return _c +} + +// SetNillableIdempotencyKey sets the "idempotency_key" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableIdempotencyKey(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetIdempotencyKey(*v) + } + return _c +} + +// SetRequestHash sets the "request_hash" field. +func (_c *BatchImageJobCreate) SetRequestHash(v string) *BatchImageJobCreate { + _c.mutation.SetRequestHash(v) + return _c +} + +// SetNillableRequestHash sets the "request_hash" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableRequestHash(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetRequestHash(*v) + } + return _c +} + +// SetManifestHash sets the "manifest_hash" field. +func (_c *BatchImageJobCreate) SetManifestHash(v string) *BatchImageJobCreate { + _c.mutation.SetManifestHash(v) + return _c +} + +// SetNillableManifestHash sets the "manifest_hash" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableManifestHash(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetManifestHash(*v) + } + return _c +} + +// SetRetryCount sets the "retry_count" field. +func (_c *BatchImageJobCreate) SetRetryCount(v int) *BatchImageJobCreate { + _c.mutation.SetRetryCount(v) + return _c +} + +// SetNillableRetryCount sets the "retry_count" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableRetryCount(v *int) *BatchImageJobCreate { + if v != nil { + _c.SetRetryCount(*v) + } + return _c +} + +// SetVersion sets the "version" field. +func (_c *BatchImageJobCreate) SetVersion(v int) *BatchImageJobCreate { + _c.mutation.SetVersion(v) + return _c +} + +// SetNillableVersion sets the "version" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableVersion(v *int) *BatchImageJobCreate { + if v != nil { + _c.SetVersion(*v) + } + return _c +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (_c *BatchImageJobCreate) SetOutputExpiresAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetOutputExpiresAt(v) + return _c +} + +// SetNillableOutputExpiresAt sets the "output_expires_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableOutputExpiresAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetOutputExpiresAt(*v) + } + return _c +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (_c *BatchImageJobCreate) SetInputDeletedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetInputDeletedAt(v) + return _c +} + +// SetNillableInputDeletedAt sets the "input_deleted_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableInputDeletedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetInputDeletedAt(*v) + } + return _c +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (_c *BatchImageJobCreate) SetOutputDeletedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetOutputDeletedAt(v) + return _c +} + +// SetNillableOutputDeletedAt sets the "output_deleted_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableOutputDeletedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetOutputDeletedAt(*v) + } + return _c +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (_c *BatchImageJobCreate) SetDownloadedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetDownloadedAt(v) + return _c +} + +// SetNillableDownloadedAt sets the "downloaded_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableDownloadedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetDownloadedAt(*v) + } + return _c +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (_c *BatchImageJobCreate) SetUserDeletedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetUserDeletedAt(v) + return _c +} + +// SetNillableUserDeletedAt sets the "user_deleted_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableUserDeletedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetUserDeletedAt(*v) + } + return _c +} + +// SetLastErrorCode sets the "last_error_code" field. +func (_c *BatchImageJobCreate) SetLastErrorCode(v string) *BatchImageJobCreate { + _c.mutation.SetLastErrorCode(v) + return _c +} + +// SetNillableLastErrorCode sets the "last_error_code" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableLastErrorCode(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetLastErrorCode(*v) + } + return _c +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (_c *BatchImageJobCreate) SetLastErrorMessage(v string) *BatchImageJobCreate { + _c.mutation.SetLastErrorMessage(v) + return _c +} + +// SetNillableLastErrorMessage sets the "last_error_message" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableLastErrorMessage(v *string) *BatchImageJobCreate { + if v != nil { + _c.SetLastErrorMessage(*v) + } + return _c +} + +// SetCreatedAt sets the "created_at" field. +func (_c *BatchImageJobCreate) SetCreatedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetCreatedAt(v) + return _c +} + +// SetNillableCreatedAt sets the "created_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableCreatedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetCreatedAt(*v) + } + return _c +} + +// SetUpdatedAt sets the "updated_at" field. +func (_c *BatchImageJobCreate) SetUpdatedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetUpdatedAt(v) + return _c +} + +// SetNillableUpdatedAt sets the "updated_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableUpdatedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetUpdatedAt(*v) + } + return _c +} + +// SetSubmittedAt sets the "submitted_at" field. +func (_c *BatchImageJobCreate) SetSubmittedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetSubmittedAt(v) + return _c +} + +// SetNillableSubmittedAt sets the "submitted_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableSubmittedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetSubmittedAt(*v) + } + return _c +} + +// SetStartedAt sets the "started_at" field. +func (_c *BatchImageJobCreate) SetStartedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetStartedAt(v) + return _c +} + +// SetNillableStartedAt sets the "started_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableStartedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetStartedAt(*v) + } + return _c +} + +// SetFinishedAt sets the "finished_at" field. +func (_c *BatchImageJobCreate) SetFinishedAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetFinishedAt(v) + return _c +} + +// SetNillableFinishedAt sets the "finished_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableFinishedAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetFinishedAt(*v) + } + return _c +} + +// SetSettledAt sets the "settled_at" field. +func (_c *BatchImageJobCreate) SetSettledAt(v time.Time) *BatchImageJobCreate { + _c.mutation.SetSettledAt(v) + return _c +} + +// SetNillableSettledAt sets the "settled_at" field if the given value is not nil. +func (_c *BatchImageJobCreate) SetNillableSettledAt(v *time.Time) *BatchImageJobCreate { + if v != nil { + _c.SetSettledAt(*v) + } + return _c +} + +// Mutation returns the BatchImageJobMutation object of the builder. +func (_c *BatchImageJobCreate) Mutation() *BatchImageJobMutation { + return _c.mutation +} + +// Save creates the BatchImageJob in the database. +func (_c *BatchImageJobCreate) Save(ctx context.Context) (*BatchImageJob, error) { + _c.defaults() + return withHooks(ctx, _c.sqlSave, _c.mutation, _c.hooks) +} + +// SaveX calls Save and panics if Save returns an error. +func (_c *BatchImageJobCreate) SaveX(ctx context.Context) *BatchImageJob { + v, err := _c.Save(ctx) + if err != nil { + panic(err) + } + return v +} + +// Exec executes the query. +func (_c *BatchImageJobCreate) Exec(ctx context.Context) error { + _, err := _c.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_c *BatchImageJobCreate) ExecX(ctx context.Context) { + if err := _c.Exec(ctx); err != nil { + panic(err) + } +} + +// defaults sets the default values of the builder before save. +func (_c *BatchImageJobCreate) defaults() { + if _, ok := _c.mutation.TaskName(); !ok { + v := batchimagejob.DefaultTaskName + _c.mutation.SetTaskName(v) + } + if _, ok := _c.mutation.Status(); !ok { + v := batchimagejob.DefaultStatus + _c.mutation.SetStatus(v) + } + if _, ok := _c.mutation.SuccessCount(); !ok { + v := batchimagejob.DefaultSuccessCount + _c.mutation.SetSuccessCount(v) + } + if _, ok := _c.mutation.FailCount(); !ok { + v := batchimagejob.DefaultFailCount + _c.mutation.SetFailCount(v) + } + if _, ok := _c.mutation.CancelledCount(); !ok { + v := batchimagejob.DefaultCancelledCount + _c.mutation.SetCancelledCount(v) + } + if _, ok := _c.mutation.EstimatedCost(); !ok { + v := batchimagejob.DefaultEstimatedCost + _c.mutation.SetEstimatedCost(v) + } + if _, ok := _c.mutation.Currency(); !ok { + v := batchimagejob.DefaultCurrency + _c.mutation.SetCurrency(v) + } + if _, ok := _c.mutation.RetryCount(); !ok { + v := batchimagejob.DefaultRetryCount + _c.mutation.SetRetryCount(v) + } + if _, ok := _c.mutation.Version(); !ok { + v := batchimagejob.DefaultVersion + _c.mutation.SetVersion(v) + } + if _, ok := _c.mutation.CreatedAt(); !ok { + v := batchimagejob.DefaultCreatedAt() + _c.mutation.SetCreatedAt(v) + } + if _, ok := _c.mutation.UpdatedAt(); !ok { + v := batchimagejob.DefaultUpdatedAt() + _c.mutation.SetUpdatedAt(v) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_c *BatchImageJobCreate) check() error { + if _, ok := _c.mutation.BatchID(); !ok { + return &ValidationError{Name: "batch_id", err: errors.New(`ent: missing required field "BatchImageJob.batch_id"`)} + } + if v, ok := _c.mutation.BatchID(); ok { + if err := batchimagejob.BatchIDValidator(v); err != nil { + return &ValidationError{Name: "batch_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.batch_id": %w`, err)} + } + } + if _, ok := _c.mutation.UserID(); !ok { + return &ValidationError{Name: "user_id", err: errors.New(`ent: missing required field "BatchImageJob.user_id"`)} + } + if _, ok := _c.mutation.Provider(); !ok { + return &ValidationError{Name: "provider", err: errors.New(`ent: missing required field "BatchImageJob.provider"`)} + } + if v, ok := _c.mutation.Provider(); ok { + if err := batchimagejob.ProviderValidator(v); err != nil { + return &ValidationError{Name: "provider", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider": %w`, err)} + } + } + if _, ok := _c.mutation.Model(); !ok { + return &ValidationError{Name: "model", err: errors.New(`ent: missing required field "BatchImageJob.model"`)} + } + if v, ok := _c.mutation.Model(); ok { + if err := batchimagejob.ModelValidator(v); err != nil { + return &ValidationError{Name: "model", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.model": %w`, err)} + } + } + if _, ok := _c.mutation.TaskName(); !ok { + return &ValidationError{Name: "task_name", err: errors.New(`ent: missing required field "BatchImageJob.task_name"`)} + } + if v, ok := _c.mutation.TaskName(); ok { + if err := batchimagejob.TaskNameValidator(v); err != nil { + return &ValidationError{Name: "task_name", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.task_name": %w`, err)} + } + } + if _, ok := _c.mutation.Status(); !ok { + return &ValidationError{Name: "status", err: errors.New(`ent: missing required field "BatchImageJob.status"`)} + } + if v, ok := _c.mutation.Status(); ok { + if err := batchimagejob.StatusValidator(v); err != nil { + return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.status": %w`, err)} + } + } + if v, ok := _c.mutation.ProviderJobName(); ok { + if err := batchimagejob.ProviderJobNameValidator(v); err != nil { + return &ValidationError{Name: "provider_job_name", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_job_name": %w`, err)} + } + } + if v, ok := _c.mutation.ProviderInputRef(); ok { + if err := batchimagejob.ProviderInputRefValidator(v); err != nil { + return &ValidationError{Name: "provider_input_ref", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_input_ref": %w`, err)} + } + } + if v, ok := _c.mutation.ProviderOutputRef(); ok { + if err := batchimagejob.ProviderOutputRefValidator(v); err != nil { + return &ValidationError{Name: "provider_output_ref", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_output_ref": %w`, err)} + } + } + if v, ok := _c.mutation.GcsInputURI(); ok { + if err := batchimagejob.GcsInputURIValidator(v); err != nil { + return &ValidationError{Name: "gcs_input_uri", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.gcs_input_uri": %w`, err)} + } + } + if v, ok := _c.mutation.GcsOutputURI(); ok { + if err := batchimagejob.GcsOutputURIValidator(v); err != nil { + return &ValidationError{Name: "gcs_output_uri", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.gcs_output_uri": %w`, err)} + } + } + if _, ok := _c.mutation.ItemCount(); !ok { + return &ValidationError{Name: "item_count", err: errors.New(`ent: missing required field "BatchImageJob.item_count"`)} + } + if _, ok := _c.mutation.SuccessCount(); !ok { + return &ValidationError{Name: "success_count", err: errors.New(`ent: missing required field "BatchImageJob.success_count"`)} + } + if _, ok := _c.mutation.FailCount(); !ok { + return &ValidationError{Name: "fail_count", err: errors.New(`ent: missing required field "BatchImageJob.fail_count"`)} + } + if _, ok := _c.mutation.CancelledCount(); !ok { + return &ValidationError{Name: "cancelled_count", err: errors.New(`ent: missing required field "BatchImageJob.cancelled_count"`)} + } + if _, ok := _c.mutation.EstimatedCost(); !ok { + return &ValidationError{Name: "estimated_cost", err: errors.New(`ent: missing required field "BatchImageJob.estimated_cost"`)} + } + if _, ok := _c.mutation.Currency(); !ok { + return &ValidationError{Name: "currency", err: errors.New(`ent: missing required field "BatchImageJob.currency"`)} + } + if v, ok := _c.mutation.Currency(); ok { + if err := batchimagejob.CurrencyValidator(v); err != nil { + return &ValidationError{Name: "currency", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.currency": %w`, err)} + } + } + if v, ok := _c.mutation.HoldID(); ok { + if err := batchimagejob.HoldIDValidator(v); err != nil { + return &ValidationError{Name: "hold_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.hold_id": %w`, err)} + } + } + if v, ok := _c.mutation.IdempotencyKey(); ok { + if err := batchimagejob.IdempotencyKeyValidator(v); err != nil { + return &ValidationError{Name: "idempotency_key", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.idempotency_key": %w`, err)} + } + } + if v, ok := _c.mutation.RequestHash(); ok { + if err := batchimagejob.RequestHashValidator(v); err != nil { + return &ValidationError{Name: "request_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.request_hash": %w`, err)} + } + } + if v, ok := _c.mutation.ManifestHash(); ok { + if err := batchimagejob.ManifestHashValidator(v); err != nil { + return &ValidationError{Name: "manifest_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.manifest_hash": %w`, err)} + } + } + if _, ok := _c.mutation.RetryCount(); !ok { + return &ValidationError{Name: "retry_count", err: errors.New(`ent: missing required field "BatchImageJob.retry_count"`)} + } + if _, ok := _c.mutation.Version(); !ok { + return &ValidationError{Name: "version", err: errors.New(`ent: missing required field "BatchImageJob.version"`)} + } + if v, ok := _c.mutation.LastErrorCode(); ok { + if err := batchimagejob.LastErrorCodeValidator(v); err != nil { + return &ValidationError{Name: "last_error_code", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.last_error_code": %w`, err)} + } + } + if _, ok := _c.mutation.CreatedAt(); !ok { + return &ValidationError{Name: "created_at", err: errors.New(`ent: missing required field "BatchImageJob.created_at"`)} + } + if _, ok := _c.mutation.UpdatedAt(); !ok { + return &ValidationError{Name: "updated_at", err: errors.New(`ent: missing required field "BatchImageJob.updated_at"`)} + } + return nil +} + +func (_c *BatchImageJobCreate) sqlSave(ctx context.Context) (*BatchImageJob, error) { + if err := _c.check(); err != nil { + return nil, err + } + _node, _spec := _c.createSpec() + if err := sqlgraph.CreateNode(ctx, _c.driver, _spec); err != nil { + if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return nil, err + } + id := _spec.ID.Value.(int64) + _node.ID = int64(id) + _c.mutation.id = &_node.ID + _c.mutation.done = true + return _node, nil +} + +func (_c *BatchImageJobCreate) createSpec() (*BatchImageJob, *sqlgraph.CreateSpec) { + var ( + _node = &BatchImageJob{config: _c.config} + _spec = sqlgraph.NewCreateSpec(batchimagejob.Table, sqlgraph.NewFieldSpec(batchimagejob.FieldID, field.TypeInt64)) + ) + _spec.OnConflict = _c.conflict + if value, ok := _c.mutation.BatchID(); ok { + _spec.SetField(batchimagejob.FieldBatchID, field.TypeString, value) + _node.BatchID = value + } + if value, ok := _c.mutation.UserID(); ok { + _spec.SetField(batchimagejob.FieldUserID, field.TypeInt64, value) + _node.UserID = value + } + if value, ok := _c.mutation.APIKeyID(); ok { + _spec.SetField(batchimagejob.FieldAPIKeyID, field.TypeInt64, value) + _node.APIKeyID = &value + } + if value, ok := _c.mutation.AccountID(); ok { + _spec.SetField(batchimagejob.FieldAccountID, field.TypeInt64, value) + _node.AccountID = &value + } + if value, ok := _c.mutation.Provider(); ok { + _spec.SetField(batchimagejob.FieldProvider, field.TypeString, value) + _node.Provider = value + } + if value, ok := _c.mutation.Model(); ok { + _spec.SetField(batchimagejob.FieldModel, field.TypeString, value) + _node.Model = value + } + if value, ok := _c.mutation.TaskName(); ok { + _spec.SetField(batchimagejob.FieldTaskName, field.TypeString, value) + _node.TaskName = value + } + if value, ok := _c.mutation.Status(); ok { + _spec.SetField(batchimagejob.FieldStatus, field.TypeString, value) + _node.Status = value + } + if value, ok := _c.mutation.ProviderJobName(); ok { + _spec.SetField(batchimagejob.FieldProviderJobName, field.TypeString, value) + _node.ProviderJobName = &value + } + if value, ok := _c.mutation.ProviderInputRef(); ok { + _spec.SetField(batchimagejob.FieldProviderInputRef, field.TypeString, value) + _node.ProviderInputRef = &value + } + if value, ok := _c.mutation.ProviderOutputRef(); ok { + _spec.SetField(batchimagejob.FieldProviderOutputRef, field.TypeString, value) + _node.ProviderOutputRef = &value + } + if value, ok := _c.mutation.GcsInputURI(); ok { + _spec.SetField(batchimagejob.FieldGcsInputURI, field.TypeString, value) + _node.GcsInputURI = &value + } + if value, ok := _c.mutation.GcsOutputURI(); ok { + _spec.SetField(batchimagejob.FieldGcsOutputURI, field.TypeString, value) + _node.GcsOutputURI = &value + } + if value, ok := _c.mutation.ItemCount(); ok { + _spec.SetField(batchimagejob.FieldItemCount, field.TypeInt, value) + _node.ItemCount = value + } + if value, ok := _c.mutation.SuccessCount(); ok { + _spec.SetField(batchimagejob.FieldSuccessCount, field.TypeInt, value) + _node.SuccessCount = value + } + if value, ok := _c.mutation.FailCount(); ok { + _spec.SetField(batchimagejob.FieldFailCount, field.TypeInt, value) + _node.FailCount = value + } + if value, ok := _c.mutation.CancelledCount(); ok { + _spec.SetField(batchimagejob.FieldCancelledCount, field.TypeInt, value) + _node.CancelledCount = value + } + if value, ok := _c.mutation.EstimatedCost(); ok { + _spec.SetField(batchimagejob.FieldEstimatedCost, field.TypeFloat64, value) + _node.EstimatedCost = value + } + if value, ok := _c.mutation.HoldAmount(); ok { + _spec.SetField(batchimagejob.FieldHoldAmount, field.TypeFloat64, value) + _node.HoldAmount = &value + } + if value, ok := _c.mutation.ActualCost(); ok { + _spec.SetField(batchimagejob.FieldActualCost, field.TypeFloat64, value) + _node.ActualCost = &value + } + if value, ok := _c.mutation.Currency(); ok { + _spec.SetField(batchimagejob.FieldCurrency, field.TypeString, value) + _node.Currency = value + } + if value, ok := _c.mutation.HoldID(); ok { + _spec.SetField(batchimagejob.FieldHoldID, field.TypeString, value) + _node.HoldID = &value + } + if value, ok := _c.mutation.IdempotencyKey(); ok { + _spec.SetField(batchimagejob.FieldIdempotencyKey, field.TypeString, value) + _node.IdempotencyKey = &value + } + if value, ok := _c.mutation.RequestHash(); ok { + _spec.SetField(batchimagejob.FieldRequestHash, field.TypeString, value) + _node.RequestHash = &value + } + if value, ok := _c.mutation.ManifestHash(); ok { + _spec.SetField(batchimagejob.FieldManifestHash, field.TypeString, value) + _node.ManifestHash = &value + } + if value, ok := _c.mutation.RetryCount(); ok { + _spec.SetField(batchimagejob.FieldRetryCount, field.TypeInt, value) + _node.RetryCount = value + } + if value, ok := _c.mutation.Version(); ok { + _spec.SetField(batchimagejob.FieldVersion, field.TypeInt, value) + _node.Version = value + } + if value, ok := _c.mutation.OutputExpiresAt(); ok { + _spec.SetField(batchimagejob.FieldOutputExpiresAt, field.TypeTime, value) + _node.OutputExpiresAt = &value + } + if value, ok := _c.mutation.InputDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldInputDeletedAt, field.TypeTime, value) + _node.InputDeletedAt = &value + } + if value, ok := _c.mutation.OutputDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldOutputDeletedAt, field.TypeTime, value) + _node.OutputDeletedAt = &value + } + if value, ok := _c.mutation.DownloadedAt(); ok { + _spec.SetField(batchimagejob.FieldDownloadedAt, field.TypeTime, value) + _node.DownloadedAt = &value + } + if value, ok := _c.mutation.UserDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldUserDeletedAt, field.TypeTime, value) + _node.UserDeletedAt = &value + } + if value, ok := _c.mutation.LastErrorCode(); ok { + _spec.SetField(batchimagejob.FieldLastErrorCode, field.TypeString, value) + _node.LastErrorCode = &value + } + if value, ok := _c.mutation.LastErrorMessage(); ok { + _spec.SetField(batchimagejob.FieldLastErrorMessage, field.TypeString, value) + _node.LastErrorMessage = &value + } + if value, ok := _c.mutation.CreatedAt(); ok { + _spec.SetField(batchimagejob.FieldCreatedAt, field.TypeTime, value) + _node.CreatedAt = value + } + if value, ok := _c.mutation.UpdatedAt(); ok { + _spec.SetField(batchimagejob.FieldUpdatedAt, field.TypeTime, value) + _node.UpdatedAt = value + } + if value, ok := _c.mutation.SubmittedAt(); ok { + _spec.SetField(batchimagejob.FieldSubmittedAt, field.TypeTime, value) + _node.SubmittedAt = &value + } + if value, ok := _c.mutation.StartedAt(); ok { + _spec.SetField(batchimagejob.FieldStartedAt, field.TypeTime, value) + _node.StartedAt = &value + } + if value, ok := _c.mutation.FinishedAt(); ok { + _spec.SetField(batchimagejob.FieldFinishedAt, field.TypeTime, value) + _node.FinishedAt = &value + } + if value, ok := _c.mutation.SettledAt(); ok { + _spec.SetField(batchimagejob.FieldSettledAt, field.TypeTime, value) + _node.SettledAt = &value + } + return _node, _spec +} + +// OnConflict allows configuring the `ON CONFLICT` / `ON DUPLICATE KEY` clause +// of the `INSERT` statement. For example: +// +// client.BatchImageJob.Create(). +// SetBatchID(v). +// OnConflict( +// // Update the row with the new values +// // the was proposed for insertion. +// sql.ResolveWithNewValues(), +// ). +// // Override some of the fields with custom +// // update values. +// Update(func(u *ent.BatchImageJobUpsert) { +// SetBatchID(v+v). +// }). +// Exec(ctx) +func (_c *BatchImageJobCreate) OnConflict(opts ...sql.ConflictOption) *BatchImageJobUpsertOne { + _c.conflict = opts + return &BatchImageJobUpsertOne{ + create: _c, + } +} + +// OnConflictColumns calls `OnConflict` and configures the columns +// as conflict target. Using this option is equivalent to using: +// +// client.BatchImageJob.Create(). +// OnConflict(sql.ConflictColumns(columns...)). +// Exec(ctx) +func (_c *BatchImageJobCreate) OnConflictColumns(columns ...string) *BatchImageJobUpsertOne { + _c.conflict = append(_c.conflict, sql.ConflictColumns(columns...)) + return &BatchImageJobUpsertOne{ + create: _c, + } +} + +type ( + // BatchImageJobUpsertOne is the builder for "upsert"-ing + // one BatchImageJob node. + BatchImageJobUpsertOne struct { + create *BatchImageJobCreate + } + + // BatchImageJobUpsert is the "OnConflict" setter. + BatchImageJobUpsert struct { + *sql.UpdateSet + } +) + +// SetUserID sets the "user_id" field. +func (u *BatchImageJobUpsert) SetUserID(v int64) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldUserID, v) + return u +} + +// UpdateUserID sets the "user_id" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateUserID() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldUserID) + return u +} + +// AddUserID adds v to the "user_id" field. +func (u *BatchImageJobUpsert) AddUserID(v int64) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldUserID, v) + return u +} + +// SetAPIKeyID sets the "api_key_id" field. +func (u *BatchImageJobUpsert) SetAPIKeyID(v int64) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldAPIKeyID, v) + return u +} + +// UpdateAPIKeyID sets the "api_key_id" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateAPIKeyID() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldAPIKeyID) + return u +} + +// AddAPIKeyID adds v to the "api_key_id" field. +func (u *BatchImageJobUpsert) AddAPIKeyID(v int64) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldAPIKeyID, v) + return u +} + +// ClearAPIKeyID clears the value of the "api_key_id" field. +func (u *BatchImageJobUpsert) ClearAPIKeyID() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldAPIKeyID) + return u +} + +// SetAccountID sets the "account_id" field. +func (u *BatchImageJobUpsert) SetAccountID(v int64) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldAccountID, v) + return u +} + +// UpdateAccountID sets the "account_id" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateAccountID() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldAccountID) + return u +} + +// AddAccountID adds v to the "account_id" field. +func (u *BatchImageJobUpsert) AddAccountID(v int64) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldAccountID, v) + return u +} + +// ClearAccountID clears the value of the "account_id" field. +func (u *BatchImageJobUpsert) ClearAccountID() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldAccountID) + return u +} + +// SetProvider sets the "provider" field. +func (u *BatchImageJobUpsert) SetProvider(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldProvider, v) + return u +} + +// UpdateProvider sets the "provider" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateProvider() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldProvider) + return u +} + +// SetModel sets the "model" field. +func (u *BatchImageJobUpsert) SetModel(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldModel, v) + return u +} + +// UpdateModel sets the "model" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateModel() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldModel) + return u +} + +// SetTaskName sets the "task_name" field. +func (u *BatchImageJobUpsert) SetTaskName(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldTaskName, v) + return u +} + +// UpdateTaskName sets the "task_name" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateTaskName() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldTaskName) + return u +} + +// SetStatus sets the "status" field. +func (u *BatchImageJobUpsert) SetStatus(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldStatus, v) + return u +} + +// UpdateStatus sets the "status" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateStatus() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldStatus) + return u +} + +// SetProviderJobName sets the "provider_job_name" field. +func (u *BatchImageJobUpsert) SetProviderJobName(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldProviderJobName, v) + return u +} + +// UpdateProviderJobName sets the "provider_job_name" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateProviderJobName() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldProviderJobName) + return u +} + +// ClearProviderJobName clears the value of the "provider_job_name" field. +func (u *BatchImageJobUpsert) ClearProviderJobName() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldProviderJobName) + return u +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (u *BatchImageJobUpsert) SetProviderInputRef(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldProviderInputRef, v) + return u +} + +// UpdateProviderInputRef sets the "provider_input_ref" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateProviderInputRef() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldProviderInputRef) + return u +} + +// ClearProviderInputRef clears the value of the "provider_input_ref" field. +func (u *BatchImageJobUpsert) ClearProviderInputRef() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldProviderInputRef) + return u +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (u *BatchImageJobUpsert) SetProviderOutputRef(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldProviderOutputRef, v) + return u +} + +// UpdateProviderOutputRef sets the "provider_output_ref" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateProviderOutputRef() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldProviderOutputRef) + return u +} + +// ClearProviderOutputRef clears the value of the "provider_output_ref" field. +func (u *BatchImageJobUpsert) ClearProviderOutputRef() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldProviderOutputRef) + return u +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (u *BatchImageJobUpsert) SetGcsInputURI(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldGcsInputURI, v) + return u +} + +// UpdateGcsInputURI sets the "gcs_input_uri" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateGcsInputURI() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldGcsInputURI) + return u +} + +// ClearGcsInputURI clears the value of the "gcs_input_uri" field. +func (u *BatchImageJobUpsert) ClearGcsInputURI() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldGcsInputURI) + return u +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (u *BatchImageJobUpsert) SetGcsOutputURI(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldGcsOutputURI, v) + return u +} + +// UpdateGcsOutputURI sets the "gcs_output_uri" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateGcsOutputURI() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldGcsOutputURI) + return u +} + +// ClearGcsOutputURI clears the value of the "gcs_output_uri" field. +func (u *BatchImageJobUpsert) ClearGcsOutputURI() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldGcsOutputURI) + return u +} + +// SetItemCount sets the "item_count" field. +func (u *BatchImageJobUpsert) SetItemCount(v int) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldItemCount, v) + return u +} + +// UpdateItemCount sets the "item_count" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateItemCount() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldItemCount) + return u +} + +// AddItemCount adds v to the "item_count" field. +func (u *BatchImageJobUpsert) AddItemCount(v int) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldItemCount, v) + return u +} + +// SetSuccessCount sets the "success_count" field. +func (u *BatchImageJobUpsert) SetSuccessCount(v int) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldSuccessCount, v) + return u +} + +// UpdateSuccessCount sets the "success_count" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateSuccessCount() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldSuccessCount) + return u +} + +// AddSuccessCount adds v to the "success_count" field. +func (u *BatchImageJobUpsert) AddSuccessCount(v int) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldSuccessCount, v) + return u +} + +// SetFailCount sets the "fail_count" field. +func (u *BatchImageJobUpsert) SetFailCount(v int) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldFailCount, v) + return u +} + +// UpdateFailCount sets the "fail_count" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateFailCount() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldFailCount) + return u +} + +// AddFailCount adds v to the "fail_count" field. +func (u *BatchImageJobUpsert) AddFailCount(v int) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldFailCount, v) + return u +} + +// SetCancelledCount sets the "cancelled_count" field. +func (u *BatchImageJobUpsert) SetCancelledCount(v int) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldCancelledCount, v) + return u +} + +// UpdateCancelledCount sets the "cancelled_count" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateCancelledCount() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldCancelledCount) + return u +} + +// AddCancelledCount adds v to the "cancelled_count" field. +func (u *BatchImageJobUpsert) AddCancelledCount(v int) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldCancelledCount, v) + return u +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (u *BatchImageJobUpsert) SetEstimatedCost(v float64) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldEstimatedCost, v) + return u +} + +// UpdateEstimatedCost sets the "estimated_cost" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateEstimatedCost() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldEstimatedCost) + return u +} + +// AddEstimatedCost adds v to the "estimated_cost" field. +func (u *BatchImageJobUpsert) AddEstimatedCost(v float64) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldEstimatedCost, v) + return u +} + +// SetHoldAmount sets the "hold_amount" field. +func (u *BatchImageJobUpsert) SetHoldAmount(v float64) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldHoldAmount, v) + return u +} + +// UpdateHoldAmount sets the "hold_amount" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateHoldAmount() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldHoldAmount) + return u +} + +// AddHoldAmount adds v to the "hold_amount" field. +func (u *BatchImageJobUpsert) AddHoldAmount(v float64) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldHoldAmount, v) + return u +} + +// ClearHoldAmount clears the value of the "hold_amount" field. +func (u *BatchImageJobUpsert) ClearHoldAmount() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldHoldAmount) + return u +} + +// SetActualCost sets the "actual_cost" field. +func (u *BatchImageJobUpsert) SetActualCost(v float64) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldActualCost, v) + return u +} + +// UpdateActualCost sets the "actual_cost" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateActualCost() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldActualCost) + return u +} + +// AddActualCost adds v to the "actual_cost" field. +func (u *BatchImageJobUpsert) AddActualCost(v float64) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldActualCost, v) + return u +} + +// ClearActualCost clears the value of the "actual_cost" field. +func (u *BatchImageJobUpsert) ClearActualCost() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldActualCost) + return u +} + +// SetCurrency sets the "currency" field. +func (u *BatchImageJobUpsert) SetCurrency(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldCurrency, v) + return u +} + +// UpdateCurrency sets the "currency" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateCurrency() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldCurrency) + return u +} + +// SetHoldID sets the "hold_id" field. +func (u *BatchImageJobUpsert) SetHoldID(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldHoldID, v) + return u +} + +// UpdateHoldID sets the "hold_id" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateHoldID() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldHoldID) + return u +} + +// ClearHoldID clears the value of the "hold_id" field. +func (u *BatchImageJobUpsert) ClearHoldID() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldHoldID) + return u +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (u *BatchImageJobUpsert) SetIdempotencyKey(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldIdempotencyKey, v) + return u +} + +// UpdateIdempotencyKey sets the "idempotency_key" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateIdempotencyKey() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldIdempotencyKey) + return u +} + +// ClearIdempotencyKey clears the value of the "idempotency_key" field. +func (u *BatchImageJobUpsert) ClearIdempotencyKey() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldIdempotencyKey) + return u +} + +// SetRequestHash sets the "request_hash" field. +func (u *BatchImageJobUpsert) SetRequestHash(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldRequestHash, v) + return u +} + +// UpdateRequestHash sets the "request_hash" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateRequestHash() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldRequestHash) + return u +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (u *BatchImageJobUpsert) ClearRequestHash() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldRequestHash) + return u +} + +// SetManifestHash sets the "manifest_hash" field. +func (u *BatchImageJobUpsert) SetManifestHash(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldManifestHash, v) + return u +} + +// UpdateManifestHash sets the "manifest_hash" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateManifestHash() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldManifestHash) + return u +} + +// ClearManifestHash clears the value of the "manifest_hash" field. +func (u *BatchImageJobUpsert) ClearManifestHash() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldManifestHash) + return u +} + +// SetRetryCount sets the "retry_count" field. +func (u *BatchImageJobUpsert) SetRetryCount(v int) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldRetryCount, v) + return u +} + +// UpdateRetryCount sets the "retry_count" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateRetryCount() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldRetryCount) + return u +} + +// AddRetryCount adds v to the "retry_count" field. +func (u *BatchImageJobUpsert) AddRetryCount(v int) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldRetryCount, v) + return u +} + +// SetVersion sets the "version" field. +func (u *BatchImageJobUpsert) SetVersion(v int) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldVersion, v) + return u +} + +// UpdateVersion sets the "version" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateVersion() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldVersion) + return u +} + +// AddVersion adds v to the "version" field. +func (u *BatchImageJobUpsert) AddVersion(v int) *BatchImageJobUpsert { + u.Add(batchimagejob.FieldVersion, v) + return u +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (u *BatchImageJobUpsert) SetOutputExpiresAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldOutputExpiresAt, v) + return u +} + +// UpdateOutputExpiresAt sets the "output_expires_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateOutputExpiresAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldOutputExpiresAt) + return u +} + +// ClearOutputExpiresAt clears the value of the "output_expires_at" field. +func (u *BatchImageJobUpsert) ClearOutputExpiresAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldOutputExpiresAt) + return u +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (u *BatchImageJobUpsert) SetInputDeletedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldInputDeletedAt, v) + return u +} + +// UpdateInputDeletedAt sets the "input_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateInputDeletedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldInputDeletedAt) + return u +} + +// ClearInputDeletedAt clears the value of the "input_deleted_at" field. +func (u *BatchImageJobUpsert) ClearInputDeletedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldInputDeletedAt) + return u +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (u *BatchImageJobUpsert) SetOutputDeletedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldOutputDeletedAt, v) + return u +} + +// UpdateOutputDeletedAt sets the "output_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateOutputDeletedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldOutputDeletedAt) + return u +} + +// ClearOutputDeletedAt clears the value of the "output_deleted_at" field. +func (u *BatchImageJobUpsert) ClearOutputDeletedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldOutputDeletedAt) + return u +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (u *BatchImageJobUpsert) SetDownloadedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldDownloadedAt, v) + return u +} + +// UpdateDownloadedAt sets the "downloaded_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateDownloadedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldDownloadedAt) + return u +} + +// ClearDownloadedAt clears the value of the "downloaded_at" field. +func (u *BatchImageJobUpsert) ClearDownloadedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldDownloadedAt) + return u +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (u *BatchImageJobUpsert) SetUserDeletedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldUserDeletedAt, v) + return u +} + +// UpdateUserDeletedAt sets the "user_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateUserDeletedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldUserDeletedAt) + return u +} + +// ClearUserDeletedAt clears the value of the "user_deleted_at" field. +func (u *BatchImageJobUpsert) ClearUserDeletedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldUserDeletedAt) + return u +} + +// SetLastErrorCode sets the "last_error_code" field. +func (u *BatchImageJobUpsert) SetLastErrorCode(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldLastErrorCode, v) + return u +} + +// UpdateLastErrorCode sets the "last_error_code" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateLastErrorCode() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldLastErrorCode) + return u +} + +// ClearLastErrorCode clears the value of the "last_error_code" field. +func (u *BatchImageJobUpsert) ClearLastErrorCode() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldLastErrorCode) + return u +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (u *BatchImageJobUpsert) SetLastErrorMessage(v string) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldLastErrorMessage, v) + return u +} + +// UpdateLastErrorMessage sets the "last_error_message" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateLastErrorMessage() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldLastErrorMessage) + return u +} + +// ClearLastErrorMessage clears the value of the "last_error_message" field. +func (u *BatchImageJobUpsert) ClearLastErrorMessage() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldLastErrorMessage) + return u +} + +// SetUpdatedAt sets the "updated_at" field. +func (u *BatchImageJobUpsert) SetUpdatedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldUpdatedAt, v) + return u +} + +// UpdateUpdatedAt sets the "updated_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateUpdatedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldUpdatedAt) + return u +} + +// SetSubmittedAt sets the "submitted_at" field. +func (u *BatchImageJobUpsert) SetSubmittedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldSubmittedAt, v) + return u +} + +// UpdateSubmittedAt sets the "submitted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateSubmittedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldSubmittedAt) + return u +} + +// ClearSubmittedAt clears the value of the "submitted_at" field. +func (u *BatchImageJobUpsert) ClearSubmittedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldSubmittedAt) + return u +} + +// SetStartedAt sets the "started_at" field. +func (u *BatchImageJobUpsert) SetStartedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldStartedAt, v) + return u +} + +// UpdateStartedAt sets the "started_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateStartedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldStartedAt) + return u +} + +// ClearStartedAt clears the value of the "started_at" field. +func (u *BatchImageJobUpsert) ClearStartedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldStartedAt) + return u +} + +// SetFinishedAt sets the "finished_at" field. +func (u *BatchImageJobUpsert) SetFinishedAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldFinishedAt, v) + return u +} + +// UpdateFinishedAt sets the "finished_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateFinishedAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldFinishedAt) + return u +} + +// ClearFinishedAt clears the value of the "finished_at" field. +func (u *BatchImageJobUpsert) ClearFinishedAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldFinishedAt) + return u +} + +// SetSettledAt sets the "settled_at" field. +func (u *BatchImageJobUpsert) SetSettledAt(v time.Time) *BatchImageJobUpsert { + u.Set(batchimagejob.FieldSettledAt, v) + return u +} + +// UpdateSettledAt sets the "settled_at" field to the value that was provided on create. +func (u *BatchImageJobUpsert) UpdateSettledAt() *BatchImageJobUpsert { + u.SetExcluded(batchimagejob.FieldSettledAt) + return u +} + +// ClearSettledAt clears the value of the "settled_at" field. +func (u *BatchImageJobUpsert) ClearSettledAt() *BatchImageJobUpsert { + u.SetNull(batchimagejob.FieldSettledAt) + return u +} + +// UpdateNewValues updates the mutable fields using the new values that were set on create. +// Using this option is equivalent to using: +// +// client.BatchImageJob.Create(). +// OnConflict( +// sql.ResolveWithNewValues(), +// ). +// Exec(ctx) +func (u *BatchImageJobUpsertOne) UpdateNewValues() *BatchImageJobUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWithNewValues()) + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(s *sql.UpdateSet) { + if _, exists := u.create.mutation.BatchID(); exists { + s.SetIgnore(batchimagejob.FieldBatchID) + } + if _, exists := u.create.mutation.CreatedAt(); exists { + s.SetIgnore(batchimagejob.FieldCreatedAt) + } + })) + return u +} + +// Ignore sets each column to itself in case of conflict. +// Using this option is equivalent to using: +// +// client.BatchImageJob.Create(). +// OnConflict(sql.ResolveWithIgnore()). +// Exec(ctx) +func (u *BatchImageJobUpsertOne) Ignore() *BatchImageJobUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWithIgnore()) + return u +} + +// DoNothing configures the conflict_action to `DO NOTHING`. +// Supported only by SQLite and PostgreSQL. +func (u *BatchImageJobUpsertOne) DoNothing() *BatchImageJobUpsertOne { + u.create.conflict = append(u.create.conflict, sql.DoNothing()) + return u +} + +// Update allows overriding fields `UPDATE` values. See the BatchImageJobCreate.OnConflict +// documentation for more info. +func (u *BatchImageJobUpsertOne) Update(set func(*BatchImageJobUpsert)) *BatchImageJobUpsertOne { + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(update *sql.UpdateSet) { + set(&BatchImageJobUpsert{UpdateSet: update}) + })) + return u +} + +// SetUserID sets the "user_id" field. +func (u *BatchImageJobUpsertOne) SetUserID(v int64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetUserID(v) + }) +} + +// AddUserID adds v to the "user_id" field. +func (u *BatchImageJobUpsertOne) AddUserID(v int64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddUserID(v) + }) +} + +// UpdateUserID sets the "user_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateUserID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateUserID() + }) +} + +// SetAPIKeyID sets the "api_key_id" field. +func (u *BatchImageJobUpsertOne) SetAPIKeyID(v int64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetAPIKeyID(v) + }) +} + +// AddAPIKeyID adds v to the "api_key_id" field. +func (u *BatchImageJobUpsertOne) AddAPIKeyID(v int64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddAPIKeyID(v) + }) +} + +// UpdateAPIKeyID sets the "api_key_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateAPIKeyID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateAPIKeyID() + }) +} + +// ClearAPIKeyID clears the value of the "api_key_id" field. +func (u *BatchImageJobUpsertOne) ClearAPIKeyID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearAPIKeyID() + }) +} + +// SetAccountID sets the "account_id" field. +func (u *BatchImageJobUpsertOne) SetAccountID(v int64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetAccountID(v) + }) +} + +// AddAccountID adds v to the "account_id" field. +func (u *BatchImageJobUpsertOne) AddAccountID(v int64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddAccountID(v) + }) +} + +// UpdateAccountID sets the "account_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateAccountID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateAccountID() + }) +} + +// ClearAccountID clears the value of the "account_id" field. +func (u *BatchImageJobUpsertOne) ClearAccountID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearAccountID() + }) +} + +// SetProvider sets the "provider" field. +func (u *BatchImageJobUpsertOne) SetProvider(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProvider(v) + }) +} + +// UpdateProvider sets the "provider" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateProvider() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProvider() + }) +} + +// SetModel sets the "model" field. +func (u *BatchImageJobUpsertOne) SetModel(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetModel(v) + }) +} + +// UpdateModel sets the "model" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateModel() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateModel() + }) +} + +// SetTaskName sets the "task_name" field. +func (u *BatchImageJobUpsertOne) SetTaskName(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetTaskName(v) + }) +} + +// UpdateTaskName sets the "task_name" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateTaskName() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateTaskName() + }) +} + +// SetStatus sets the "status" field. +func (u *BatchImageJobUpsertOne) SetStatus(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetStatus(v) + }) +} + +// UpdateStatus sets the "status" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateStatus() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateStatus() + }) +} + +// SetProviderJobName sets the "provider_job_name" field. +func (u *BatchImageJobUpsertOne) SetProviderJobName(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProviderJobName(v) + }) +} + +// UpdateProviderJobName sets the "provider_job_name" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateProviderJobName() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProviderJobName() + }) +} + +// ClearProviderJobName clears the value of the "provider_job_name" field. +func (u *BatchImageJobUpsertOne) ClearProviderJobName() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearProviderJobName() + }) +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (u *BatchImageJobUpsertOne) SetProviderInputRef(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProviderInputRef(v) + }) +} + +// UpdateProviderInputRef sets the "provider_input_ref" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateProviderInputRef() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProviderInputRef() + }) +} + +// ClearProviderInputRef clears the value of the "provider_input_ref" field. +func (u *BatchImageJobUpsertOne) ClearProviderInputRef() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearProviderInputRef() + }) +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (u *BatchImageJobUpsertOne) SetProviderOutputRef(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProviderOutputRef(v) + }) +} + +// UpdateProviderOutputRef sets the "provider_output_ref" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateProviderOutputRef() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProviderOutputRef() + }) +} + +// ClearProviderOutputRef clears the value of the "provider_output_ref" field. +func (u *BatchImageJobUpsertOne) ClearProviderOutputRef() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearProviderOutputRef() + }) +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (u *BatchImageJobUpsertOne) SetGcsInputURI(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetGcsInputURI(v) + }) +} + +// UpdateGcsInputURI sets the "gcs_input_uri" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateGcsInputURI() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateGcsInputURI() + }) +} + +// ClearGcsInputURI clears the value of the "gcs_input_uri" field. +func (u *BatchImageJobUpsertOne) ClearGcsInputURI() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearGcsInputURI() + }) +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (u *BatchImageJobUpsertOne) SetGcsOutputURI(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetGcsOutputURI(v) + }) +} + +// UpdateGcsOutputURI sets the "gcs_output_uri" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateGcsOutputURI() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateGcsOutputURI() + }) +} + +// ClearGcsOutputURI clears the value of the "gcs_output_uri" field. +func (u *BatchImageJobUpsertOne) ClearGcsOutputURI() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearGcsOutputURI() + }) +} + +// SetItemCount sets the "item_count" field. +func (u *BatchImageJobUpsertOne) SetItemCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetItemCount(v) + }) +} + +// AddItemCount adds v to the "item_count" field. +func (u *BatchImageJobUpsertOne) AddItemCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddItemCount(v) + }) +} + +// UpdateItemCount sets the "item_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateItemCount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateItemCount() + }) +} + +// SetSuccessCount sets the "success_count" field. +func (u *BatchImageJobUpsertOne) SetSuccessCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetSuccessCount(v) + }) +} + +// AddSuccessCount adds v to the "success_count" field. +func (u *BatchImageJobUpsertOne) AddSuccessCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddSuccessCount(v) + }) +} + +// UpdateSuccessCount sets the "success_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateSuccessCount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateSuccessCount() + }) +} + +// SetFailCount sets the "fail_count" field. +func (u *BatchImageJobUpsertOne) SetFailCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetFailCount(v) + }) +} + +// AddFailCount adds v to the "fail_count" field. +func (u *BatchImageJobUpsertOne) AddFailCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddFailCount(v) + }) +} + +// UpdateFailCount sets the "fail_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateFailCount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateFailCount() + }) +} + +// SetCancelledCount sets the "cancelled_count" field. +func (u *BatchImageJobUpsertOne) SetCancelledCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetCancelledCount(v) + }) +} + +// AddCancelledCount adds v to the "cancelled_count" field. +func (u *BatchImageJobUpsertOne) AddCancelledCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddCancelledCount(v) + }) +} + +// UpdateCancelledCount sets the "cancelled_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateCancelledCount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateCancelledCount() + }) +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (u *BatchImageJobUpsertOne) SetEstimatedCost(v float64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetEstimatedCost(v) + }) +} + +// AddEstimatedCost adds v to the "estimated_cost" field. +func (u *BatchImageJobUpsertOne) AddEstimatedCost(v float64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddEstimatedCost(v) + }) +} + +// UpdateEstimatedCost sets the "estimated_cost" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateEstimatedCost() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateEstimatedCost() + }) +} + +// SetHoldAmount sets the "hold_amount" field. +func (u *BatchImageJobUpsertOne) SetHoldAmount(v float64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetHoldAmount(v) + }) +} + +// AddHoldAmount adds v to the "hold_amount" field. +func (u *BatchImageJobUpsertOne) AddHoldAmount(v float64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddHoldAmount(v) + }) +} + +// UpdateHoldAmount sets the "hold_amount" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateHoldAmount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateHoldAmount() + }) +} + +// ClearHoldAmount clears the value of the "hold_amount" field. +func (u *BatchImageJobUpsertOne) ClearHoldAmount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearHoldAmount() + }) +} + +// SetActualCost sets the "actual_cost" field. +func (u *BatchImageJobUpsertOne) SetActualCost(v float64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetActualCost(v) + }) +} + +// AddActualCost adds v to the "actual_cost" field. +func (u *BatchImageJobUpsertOne) AddActualCost(v float64) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddActualCost(v) + }) +} + +// UpdateActualCost sets the "actual_cost" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateActualCost() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateActualCost() + }) +} + +// ClearActualCost clears the value of the "actual_cost" field. +func (u *BatchImageJobUpsertOne) ClearActualCost() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearActualCost() + }) +} + +// SetCurrency sets the "currency" field. +func (u *BatchImageJobUpsertOne) SetCurrency(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetCurrency(v) + }) +} + +// UpdateCurrency sets the "currency" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateCurrency() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateCurrency() + }) +} + +// SetHoldID sets the "hold_id" field. +func (u *BatchImageJobUpsertOne) SetHoldID(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetHoldID(v) + }) +} + +// UpdateHoldID sets the "hold_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateHoldID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateHoldID() + }) +} + +// ClearHoldID clears the value of the "hold_id" field. +func (u *BatchImageJobUpsertOne) ClearHoldID() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearHoldID() + }) +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (u *BatchImageJobUpsertOne) SetIdempotencyKey(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetIdempotencyKey(v) + }) +} + +// UpdateIdempotencyKey sets the "idempotency_key" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateIdempotencyKey() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateIdempotencyKey() + }) +} + +// ClearIdempotencyKey clears the value of the "idempotency_key" field. +func (u *BatchImageJobUpsertOne) ClearIdempotencyKey() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearIdempotencyKey() + }) +} + +// SetRequestHash sets the "request_hash" field. +func (u *BatchImageJobUpsertOne) SetRequestHash(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetRequestHash(v) + }) +} + +// UpdateRequestHash sets the "request_hash" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateRequestHash() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateRequestHash() + }) +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (u *BatchImageJobUpsertOne) ClearRequestHash() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearRequestHash() + }) +} + +// SetManifestHash sets the "manifest_hash" field. +func (u *BatchImageJobUpsertOne) SetManifestHash(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetManifestHash(v) + }) +} + +// UpdateManifestHash sets the "manifest_hash" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateManifestHash() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateManifestHash() + }) +} + +// ClearManifestHash clears the value of the "manifest_hash" field. +func (u *BatchImageJobUpsertOne) ClearManifestHash() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearManifestHash() + }) +} + +// SetRetryCount sets the "retry_count" field. +func (u *BatchImageJobUpsertOne) SetRetryCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetRetryCount(v) + }) +} + +// AddRetryCount adds v to the "retry_count" field. +func (u *BatchImageJobUpsertOne) AddRetryCount(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddRetryCount(v) + }) +} + +// UpdateRetryCount sets the "retry_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateRetryCount() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateRetryCount() + }) +} + +// SetVersion sets the "version" field. +func (u *BatchImageJobUpsertOne) SetVersion(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetVersion(v) + }) +} + +// AddVersion adds v to the "version" field. +func (u *BatchImageJobUpsertOne) AddVersion(v int) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddVersion(v) + }) +} + +// UpdateVersion sets the "version" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateVersion() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateVersion() + }) +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (u *BatchImageJobUpsertOne) SetOutputExpiresAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetOutputExpiresAt(v) + }) +} + +// UpdateOutputExpiresAt sets the "output_expires_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateOutputExpiresAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateOutputExpiresAt() + }) +} + +// ClearOutputExpiresAt clears the value of the "output_expires_at" field. +func (u *BatchImageJobUpsertOne) ClearOutputExpiresAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearOutputExpiresAt() + }) +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (u *BatchImageJobUpsertOne) SetInputDeletedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetInputDeletedAt(v) + }) +} + +// UpdateInputDeletedAt sets the "input_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateInputDeletedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateInputDeletedAt() + }) +} + +// ClearInputDeletedAt clears the value of the "input_deleted_at" field. +func (u *BatchImageJobUpsertOne) ClearInputDeletedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearInputDeletedAt() + }) +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (u *BatchImageJobUpsertOne) SetOutputDeletedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetOutputDeletedAt(v) + }) +} + +// UpdateOutputDeletedAt sets the "output_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateOutputDeletedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateOutputDeletedAt() + }) +} + +// ClearOutputDeletedAt clears the value of the "output_deleted_at" field. +func (u *BatchImageJobUpsertOne) ClearOutputDeletedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearOutputDeletedAt() + }) +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (u *BatchImageJobUpsertOne) SetDownloadedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetDownloadedAt(v) + }) +} + +// UpdateDownloadedAt sets the "downloaded_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateDownloadedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateDownloadedAt() + }) +} + +// ClearDownloadedAt clears the value of the "downloaded_at" field. +func (u *BatchImageJobUpsertOne) ClearDownloadedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearDownloadedAt() + }) +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (u *BatchImageJobUpsertOne) SetUserDeletedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetUserDeletedAt(v) + }) +} + +// UpdateUserDeletedAt sets the "user_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateUserDeletedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateUserDeletedAt() + }) +} + +// ClearUserDeletedAt clears the value of the "user_deleted_at" field. +func (u *BatchImageJobUpsertOne) ClearUserDeletedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearUserDeletedAt() + }) +} + +// SetLastErrorCode sets the "last_error_code" field. +func (u *BatchImageJobUpsertOne) SetLastErrorCode(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetLastErrorCode(v) + }) +} + +// UpdateLastErrorCode sets the "last_error_code" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateLastErrorCode() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateLastErrorCode() + }) +} + +// ClearLastErrorCode clears the value of the "last_error_code" field. +func (u *BatchImageJobUpsertOne) ClearLastErrorCode() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearLastErrorCode() + }) +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (u *BatchImageJobUpsertOne) SetLastErrorMessage(v string) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetLastErrorMessage(v) + }) +} + +// UpdateLastErrorMessage sets the "last_error_message" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateLastErrorMessage() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateLastErrorMessage() + }) +} + +// ClearLastErrorMessage clears the value of the "last_error_message" field. +func (u *BatchImageJobUpsertOne) ClearLastErrorMessage() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearLastErrorMessage() + }) +} + +// SetUpdatedAt sets the "updated_at" field. +func (u *BatchImageJobUpsertOne) SetUpdatedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetUpdatedAt(v) + }) +} + +// UpdateUpdatedAt sets the "updated_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateUpdatedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateUpdatedAt() + }) +} + +// SetSubmittedAt sets the "submitted_at" field. +func (u *BatchImageJobUpsertOne) SetSubmittedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetSubmittedAt(v) + }) +} + +// UpdateSubmittedAt sets the "submitted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateSubmittedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateSubmittedAt() + }) +} + +// ClearSubmittedAt clears the value of the "submitted_at" field. +func (u *BatchImageJobUpsertOne) ClearSubmittedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearSubmittedAt() + }) +} + +// SetStartedAt sets the "started_at" field. +func (u *BatchImageJobUpsertOne) SetStartedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetStartedAt(v) + }) +} + +// UpdateStartedAt sets the "started_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateStartedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateStartedAt() + }) +} + +// ClearStartedAt clears the value of the "started_at" field. +func (u *BatchImageJobUpsertOne) ClearStartedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearStartedAt() + }) +} + +// SetFinishedAt sets the "finished_at" field. +func (u *BatchImageJobUpsertOne) SetFinishedAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetFinishedAt(v) + }) +} + +// UpdateFinishedAt sets the "finished_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateFinishedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateFinishedAt() + }) +} + +// ClearFinishedAt clears the value of the "finished_at" field. +func (u *BatchImageJobUpsertOne) ClearFinishedAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearFinishedAt() + }) +} + +// SetSettledAt sets the "settled_at" field. +func (u *BatchImageJobUpsertOne) SetSettledAt(v time.Time) *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetSettledAt(v) + }) +} + +// UpdateSettledAt sets the "settled_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertOne) UpdateSettledAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateSettledAt() + }) +} + +// ClearSettledAt clears the value of the "settled_at" field. +func (u *BatchImageJobUpsertOne) ClearSettledAt() *BatchImageJobUpsertOne { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearSettledAt() + }) +} + +// Exec executes the query. +func (u *BatchImageJobUpsertOne) Exec(ctx context.Context) error { + if len(u.create.conflict) == 0 { + return errors.New("ent: missing options for BatchImageJobCreate.OnConflict") + } + return u.create.Exec(ctx) +} + +// ExecX is like Exec, but panics if an error occurs. +func (u *BatchImageJobUpsertOne) ExecX(ctx context.Context) { + if err := u.create.Exec(ctx); err != nil { + panic(err) + } +} + +// Exec executes the UPSERT query and returns the inserted/updated ID. +func (u *BatchImageJobUpsertOne) ID(ctx context.Context) (id int64, err error) { + node, err := u.create.Save(ctx) + if err != nil { + return id, err + } + return node.ID, nil +} + +// IDX is like ID, but panics if an error occurs. +func (u *BatchImageJobUpsertOne) IDX(ctx context.Context) int64 { + id, err := u.ID(ctx) + if err != nil { + panic(err) + } + return id +} + +// BatchImageJobCreateBulk is the builder for creating many BatchImageJob entities in bulk. +type BatchImageJobCreateBulk struct { + config + err error + builders []*BatchImageJobCreate + conflict []sql.ConflictOption +} + +// Save creates the BatchImageJob entities in the database. +func (_c *BatchImageJobCreateBulk) Save(ctx context.Context) ([]*BatchImageJob, error) { + if _c.err != nil { + return nil, _c.err + } + specs := make([]*sqlgraph.CreateSpec, len(_c.builders)) + nodes := make([]*BatchImageJob, len(_c.builders)) + mutators := make([]Mutator, len(_c.builders)) + for i := range _c.builders { + func(i int, root context.Context) { + builder := _c.builders[i] + builder.defaults() + var mut Mutator = MutateFunc(func(ctx context.Context, m Mutation) (Value, error) { + mutation, ok := m.(*BatchImageJobMutation) + if !ok { + return nil, fmt.Errorf("unexpected mutation type %T", m) + } + if err := builder.check(); err != nil { + return nil, err + } + builder.mutation = mutation + var err error + nodes[i], specs[i] = builder.createSpec() + if i < len(mutators)-1 { + _, err = mutators[i+1].Mutate(root, _c.builders[i+1].mutation) + } else { + spec := &sqlgraph.BatchCreateSpec{Nodes: specs} + spec.OnConflict = _c.conflict + // Invoke the actual operation on the latest mutation in the chain. + if err = sqlgraph.BatchCreate(ctx, _c.driver, spec); err != nil { + if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + } + } + if err != nil { + return nil, err + } + mutation.id = &nodes[i].ID + if specs[i].ID.Value != nil { + id := specs[i].ID.Value.(int64) + nodes[i].ID = int64(id) + } + mutation.done = true + return nodes[i], nil + }) + for i := len(builder.hooks) - 1; i >= 0; i-- { + mut = builder.hooks[i](mut) + } + mutators[i] = mut + }(i, ctx) + } + if len(mutators) > 0 { + if _, err := mutators[0].Mutate(ctx, _c.builders[0].mutation); err != nil { + return nil, err + } + } + return nodes, nil +} + +// SaveX is like Save, but panics if an error occurs. +func (_c *BatchImageJobCreateBulk) SaveX(ctx context.Context) []*BatchImageJob { + v, err := _c.Save(ctx) + if err != nil { + panic(err) + } + return v +} + +// Exec executes the query. +func (_c *BatchImageJobCreateBulk) Exec(ctx context.Context) error { + _, err := _c.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_c *BatchImageJobCreateBulk) ExecX(ctx context.Context) { + if err := _c.Exec(ctx); err != nil { + panic(err) + } +} + +// OnConflict allows configuring the `ON CONFLICT` / `ON DUPLICATE KEY` clause +// of the `INSERT` statement. For example: +// +// client.BatchImageJob.CreateBulk(builders...). +// OnConflict( +// // Update the row with the new values +// // the was proposed for insertion. +// sql.ResolveWithNewValues(), +// ). +// // Override some of the fields with custom +// // update values. +// Update(func(u *ent.BatchImageJobUpsert) { +// SetBatchID(v+v). +// }). +// Exec(ctx) +func (_c *BatchImageJobCreateBulk) OnConflict(opts ...sql.ConflictOption) *BatchImageJobUpsertBulk { + _c.conflict = opts + return &BatchImageJobUpsertBulk{ + create: _c, + } +} + +// OnConflictColumns calls `OnConflict` and configures the columns +// as conflict target. Using this option is equivalent to using: +// +// client.BatchImageJob.Create(). +// OnConflict(sql.ConflictColumns(columns...)). +// Exec(ctx) +func (_c *BatchImageJobCreateBulk) OnConflictColumns(columns ...string) *BatchImageJobUpsertBulk { + _c.conflict = append(_c.conflict, sql.ConflictColumns(columns...)) + return &BatchImageJobUpsertBulk{ + create: _c, + } +} + +// BatchImageJobUpsertBulk is the builder for "upsert"-ing +// a bulk of BatchImageJob nodes. +type BatchImageJobUpsertBulk struct { + create *BatchImageJobCreateBulk +} + +// UpdateNewValues updates the mutable fields using the new values that +// were set on create. Using this option is equivalent to using: +// +// client.BatchImageJob.Create(). +// OnConflict( +// sql.ResolveWithNewValues(), +// ). +// Exec(ctx) +func (u *BatchImageJobUpsertBulk) UpdateNewValues() *BatchImageJobUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWithNewValues()) + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(s *sql.UpdateSet) { + for _, b := range u.create.builders { + if _, exists := b.mutation.BatchID(); exists { + s.SetIgnore(batchimagejob.FieldBatchID) + } + if _, exists := b.mutation.CreatedAt(); exists { + s.SetIgnore(batchimagejob.FieldCreatedAt) + } + } + })) + return u +} + +// Ignore sets each column to itself in case of conflict. +// Using this option is equivalent to using: +// +// client.BatchImageJob.Create(). +// OnConflict(sql.ResolveWithIgnore()). +// Exec(ctx) +func (u *BatchImageJobUpsertBulk) Ignore() *BatchImageJobUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWithIgnore()) + return u +} + +// DoNothing configures the conflict_action to `DO NOTHING`. +// Supported only by SQLite and PostgreSQL. +func (u *BatchImageJobUpsertBulk) DoNothing() *BatchImageJobUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.DoNothing()) + return u +} + +// Update allows overriding fields `UPDATE` values. See the BatchImageJobCreateBulk.OnConflict +// documentation for more info. +func (u *BatchImageJobUpsertBulk) Update(set func(*BatchImageJobUpsert)) *BatchImageJobUpsertBulk { + u.create.conflict = append(u.create.conflict, sql.ResolveWith(func(update *sql.UpdateSet) { + set(&BatchImageJobUpsert{UpdateSet: update}) + })) + return u +} + +// SetUserID sets the "user_id" field. +func (u *BatchImageJobUpsertBulk) SetUserID(v int64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetUserID(v) + }) +} + +// AddUserID adds v to the "user_id" field. +func (u *BatchImageJobUpsertBulk) AddUserID(v int64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddUserID(v) + }) +} + +// UpdateUserID sets the "user_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateUserID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateUserID() + }) +} + +// SetAPIKeyID sets the "api_key_id" field. +func (u *BatchImageJobUpsertBulk) SetAPIKeyID(v int64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetAPIKeyID(v) + }) +} + +// AddAPIKeyID adds v to the "api_key_id" field. +func (u *BatchImageJobUpsertBulk) AddAPIKeyID(v int64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddAPIKeyID(v) + }) +} + +// UpdateAPIKeyID sets the "api_key_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateAPIKeyID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateAPIKeyID() + }) +} + +// ClearAPIKeyID clears the value of the "api_key_id" field. +func (u *BatchImageJobUpsertBulk) ClearAPIKeyID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearAPIKeyID() + }) +} + +// SetAccountID sets the "account_id" field. +func (u *BatchImageJobUpsertBulk) SetAccountID(v int64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetAccountID(v) + }) +} + +// AddAccountID adds v to the "account_id" field. +func (u *BatchImageJobUpsertBulk) AddAccountID(v int64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddAccountID(v) + }) +} + +// UpdateAccountID sets the "account_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateAccountID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateAccountID() + }) +} + +// ClearAccountID clears the value of the "account_id" field. +func (u *BatchImageJobUpsertBulk) ClearAccountID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearAccountID() + }) +} + +// SetProvider sets the "provider" field. +func (u *BatchImageJobUpsertBulk) SetProvider(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProvider(v) + }) +} + +// UpdateProvider sets the "provider" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateProvider() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProvider() + }) +} + +// SetModel sets the "model" field. +func (u *BatchImageJobUpsertBulk) SetModel(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetModel(v) + }) +} + +// UpdateModel sets the "model" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateModel() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateModel() + }) +} + +// SetTaskName sets the "task_name" field. +func (u *BatchImageJobUpsertBulk) SetTaskName(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetTaskName(v) + }) +} + +// UpdateTaskName sets the "task_name" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateTaskName() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateTaskName() + }) +} + +// SetStatus sets the "status" field. +func (u *BatchImageJobUpsertBulk) SetStatus(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetStatus(v) + }) +} + +// UpdateStatus sets the "status" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateStatus() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateStatus() + }) +} + +// SetProviderJobName sets the "provider_job_name" field. +func (u *BatchImageJobUpsertBulk) SetProviderJobName(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProviderJobName(v) + }) +} + +// UpdateProviderJobName sets the "provider_job_name" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateProviderJobName() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProviderJobName() + }) +} + +// ClearProviderJobName clears the value of the "provider_job_name" field. +func (u *BatchImageJobUpsertBulk) ClearProviderJobName() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearProviderJobName() + }) +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (u *BatchImageJobUpsertBulk) SetProviderInputRef(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProviderInputRef(v) + }) +} + +// UpdateProviderInputRef sets the "provider_input_ref" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateProviderInputRef() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProviderInputRef() + }) +} + +// ClearProviderInputRef clears the value of the "provider_input_ref" field. +func (u *BatchImageJobUpsertBulk) ClearProviderInputRef() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearProviderInputRef() + }) +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (u *BatchImageJobUpsertBulk) SetProviderOutputRef(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetProviderOutputRef(v) + }) +} + +// UpdateProviderOutputRef sets the "provider_output_ref" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateProviderOutputRef() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateProviderOutputRef() + }) +} + +// ClearProviderOutputRef clears the value of the "provider_output_ref" field. +func (u *BatchImageJobUpsertBulk) ClearProviderOutputRef() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearProviderOutputRef() + }) +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (u *BatchImageJobUpsertBulk) SetGcsInputURI(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetGcsInputURI(v) + }) +} + +// UpdateGcsInputURI sets the "gcs_input_uri" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateGcsInputURI() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateGcsInputURI() + }) +} + +// ClearGcsInputURI clears the value of the "gcs_input_uri" field. +func (u *BatchImageJobUpsertBulk) ClearGcsInputURI() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearGcsInputURI() + }) +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (u *BatchImageJobUpsertBulk) SetGcsOutputURI(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetGcsOutputURI(v) + }) +} + +// UpdateGcsOutputURI sets the "gcs_output_uri" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateGcsOutputURI() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateGcsOutputURI() + }) +} + +// ClearGcsOutputURI clears the value of the "gcs_output_uri" field. +func (u *BatchImageJobUpsertBulk) ClearGcsOutputURI() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearGcsOutputURI() + }) +} + +// SetItemCount sets the "item_count" field. +func (u *BatchImageJobUpsertBulk) SetItemCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetItemCount(v) + }) +} + +// AddItemCount adds v to the "item_count" field. +func (u *BatchImageJobUpsertBulk) AddItemCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddItemCount(v) + }) +} + +// UpdateItemCount sets the "item_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateItemCount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateItemCount() + }) +} + +// SetSuccessCount sets the "success_count" field. +func (u *BatchImageJobUpsertBulk) SetSuccessCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetSuccessCount(v) + }) +} + +// AddSuccessCount adds v to the "success_count" field. +func (u *BatchImageJobUpsertBulk) AddSuccessCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddSuccessCount(v) + }) +} + +// UpdateSuccessCount sets the "success_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateSuccessCount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateSuccessCount() + }) +} + +// SetFailCount sets the "fail_count" field. +func (u *BatchImageJobUpsertBulk) SetFailCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetFailCount(v) + }) +} + +// AddFailCount adds v to the "fail_count" field. +func (u *BatchImageJobUpsertBulk) AddFailCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddFailCount(v) + }) +} + +// UpdateFailCount sets the "fail_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateFailCount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateFailCount() + }) +} + +// SetCancelledCount sets the "cancelled_count" field. +func (u *BatchImageJobUpsertBulk) SetCancelledCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetCancelledCount(v) + }) +} + +// AddCancelledCount adds v to the "cancelled_count" field. +func (u *BatchImageJobUpsertBulk) AddCancelledCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddCancelledCount(v) + }) +} + +// UpdateCancelledCount sets the "cancelled_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateCancelledCount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateCancelledCount() + }) +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (u *BatchImageJobUpsertBulk) SetEstimatedCost(v float64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetEstimatedCost(v) + }) +} + +// AddEstimatedCost adds v to the "estimated_cost" field. +func (u *BatchImageJobUpsertBulk) AddEstimatedCost(v float64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddEstimatedCost(v) + }) +} + +// UpdateEstimatedCost sets the "estimated_cost" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateEstimatedCost() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateEstimatedCost() + }) +} + +// SetHoldAmount sets the "hold_amount" field. +func (u *BatchImageJobUpsertBulk) SetHoldAmount(v float64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetHoldAmount(v) + }) +} + +// AddHoldAmount adds v to the "hold_amount" field. +func (u *BatchImageJobUpsertBulk) AddHoldAmount(v float64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddHoldAmount(v) + }) +} + +// UpdateHoldAmount sets the "hold_amount" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateHoldAmount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateHoldAmount() + }) +} + +// ClearHoldAmount clears the value of the "hold_amount" field. +func (u *BatchImageJobUpsertBulk) ClearHoldAmount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearHoldAmount() + }) +} + +// SetActualCost sets the "actual_cost" field. +func (u *BatchImageJobUpsertBulk) SetActualCost(v float64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetActualCost(v) + }) +} + +// AddActualCost adds v to the "actual_cost" field. +func (u *BatchImageJobUpsertBulk) AddActualCost(v float64) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddActualCost(v) + }) +} + +// UpdateActualCost sets the "actual_cost" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateActualCost() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateActualCost() + }) +} + +// ClearActualCost clears the value of the "actual_cost" field. +func (u *BatchImageJobUpsertBulk) ClearActualCost() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearActualCost() + }) +} + +// SetCurrency sets the "currency" field. +func (u *BatchImageJobUpsertBulk) SetCurrency(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetCurrency(v) + }) +} + +// UpdateCurrency sets the "currency" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateCurrency() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateCurrency() + }) +} + +// SetHoldID sets the "hold_id" field. +func (u *BatchImageJobUpsertBulk) SetHoldID(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetHoldID(v) + }) +} + +// UpdateHoldID sets the "hold_id" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateHoldID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateHoldID() + }) +} + +// ClearHoldID clears the value of the "hold_id" field. +func (u *BatchImageJobUpsertBulk) ClearHoldID() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearHoldID() + }) +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (u *BatchImageJobUpsertBulk) SetIdempotencyKey(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetIdempotencyKey(v) + }) +} + +// UpdateIdempotencyKey sets the "idempotency_key" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateIdempotencyKey() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateIdempotencyKey() + }) +} + +// ClearIdempotencyKey clears the value of the "idempotency_key" field. +func (u *BatchImageJobUpsertBulk) ClearIdempotencyKey() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearIdempotencyKey() + }) +} + +// SetRequestHash sets the "request_hash" field. +func (u *BatchImageJobUpsertBulk) SetRequestHash(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetRequestHash(v) + }) +} + +// UpdateRequestHash sets the "request_hash" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateRequestHash() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateRequestHash() + }) +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (u *BatchImageJobUpsertBulk) ClearRequestHash() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearRequestHash() + }) +} + +// SetManifestHash sets the "manifest_hash" field. +func (u *BatchImageJobUpsertBulk) SetManifestHash(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetManifestHash(v) + }) +} + +// UpdateManifestHash sets the "manifest_hash" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateManifestHash() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateManifestHash() + }) +} + +// ClearManifestHash clears the value of the "manifest_hash" field. +func (u *BatchImageJobUpsertBulk) ClearManifestHash() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearManifestHash() + }) +} + +// SetRetryCount sets the "retry_count" field. +func (u *BatchImageJobUpsertBulk) SetRetryCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetRetryCount(v) + }) +} + +// AddRetryCount adds v to the "retry_count" field. +func (u *BatchImageJobUpsertBulk) AddRetryCount(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddRetryCount(v) + }) +} + +// UpdateRetryCount sets the "retry_count" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateRetryCount() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateRetryCount() + }) +} + +// SetVersion sets the "version" field. +func (u *BatchImageJobUpsertBulk) SetVersion(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetVersion(v) + }) +} + +// AddVersion adds v to the "version" field. +func (u *BatchImageJobUpsertBulk) AddVersion(v int) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.AddVersion(v) + }) +} + +// UpdateVersion sets the "version" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateVersion() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateVersion() + }) +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (u *BatchImageJobUpsertBulk) SetOutputExpiresAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetOutputExpiresAt(v) + }) +} + +// UpdateOutputExpiresAt sets the "output_expires_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateOutputExpiresAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateOutputExpiresAt() + }) +} + +// ClearOutputExpiresAt clears the value of the "output_expires_at" field. +func (u *BatchImageJobUpsertBulk) ClearOutputExpiresAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearOutputExpiresAt() + }) +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (u *BatchImageJobUpsertBulk) SetInputDeletedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetInputDeletedAt(v) + }) +} + +// UpdateInputDeletedAt sets the "input_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateInputDeletedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateInputDeletedAt() + }) +} + +// ClearInputDeletedAt clears the value of the "input_deleted_at" field. +func (u *BatchImageJobUpsertBulk) ClearInputDeletedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearInputDeletedAt() + }) +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (u *BatchImageJobUpsertBulk) SetOutputDeletedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetOutputDeletedAt(v) + }) +} + +// UpdateOutputDeletedAt sets the "output_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateOutputDeletedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateOutputDeletedAt() + }) +} + +// ClearOutputDeletedAt clears the value of the "output_deleted_at" field. +func (u *BatchImageJobUpsertBulk) ClearOutputDeletedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearOutputDeletedAt() + }) +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (u *BatchImageJobUpsertBulk) SetDownloadedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetDownloadedAt(v) + }) +} + +// UpdateDownloadedAt sets the "downloaded_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateDownloadedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateDownloadedAt() + }) +} + +// ClearDownloadedAt clears the value of the "downloaded_at" field. +func (u *BatchImageJobUpsertBulk) ClearDownloadedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearDownloadedAt() + }) +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (u *BatchImageJobUpsertBulk) SetUserDeletedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetUserDeletedAt(v) + }) +} + +// UpdateUserDeletedAt sets the "user_deleted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateUserDeletedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateUserDeletedAt() + }) +} + +// ClearUserDeletedAt clears the value of the "user_deleted_at" field. +func (u *BatchImageJobUpsertBulk) ClearUserDeletedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearUserDeletedAt() + }) +} + +// SetLastErrorCode sets the "last_error_code" field. +func (u *BatchImageJobUpsertBulk) SetLastErrorCode(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetLastErrorCode(v) + }) +} + +// UpdateLastErrorCode sets the "last_error_code" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateLastErrorCode() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateLastErrorCode() + }) +} + +// ClearLastErrorCode clears the value of the "last_error_code" field. +func (u *BatchImageJobUpsertBulk) ClearLastErrorCode() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearLastErrorCode() + }) +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (u *BatchImageJobUpsertBulk) SetLastErrorMessage(v string) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetLastErrorMessage(v) + }) +} + +// UpdateLastErrorMessage sets the "last_error_message" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateLastErrorMessage() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateLastErrorMessage() + }) +} + +// ClearLastErrorMessage clears the value of the "last_error_message" field. +func (u *BatchImageJobUpsertBulk) ClearLastErrorMessage() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearLastErrorMessage() + }) +} + +// SetUpdatedAt sets the "updated_at" field. +func (u *BatchImageJobUpsertBulk) SetUpdatedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetUpdatedAt(v) + }) +} + +// UpdateUpdatedAt sets the "updated_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateUpdatedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateUpdatedAt() + }) +} + +// SetSubmittedAt sets the "submitted_at" field. +func (u *BatchImageJobUpsertBulk) SetSubmittedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetSubmittedAt(v) + }) +} + +// UpdateSubmittedAt sets the "submitted_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateSubmittedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateSubmittedAt() + }) +} + +// ClearSubmittedAt clears the value of the "submitted_at" field. +func (u *BatchImageJobUpsertBulk) ClearSubmittedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearSubmittedAt() + }) +} + +// SetStartedAt sets the "started_at" field. +func (u *BatchImageJobUpsertBulk) SetStartedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetStartedAt(v) + }) +} + +// UpdateStartedAt sets the "started_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateStartedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateStartedAt() + }) +} + +// ClearStartedAt clears the value of the "started_at" field. +func (u *BatchImageJobUpsertBulk) ClearStartedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearStartedAt() + }) +} + +// SetFinishedAt sets the "finished_at" field. +func (u *BatchImageJobUpsertBulk) SetFinishedAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetFinishedAt(v) + }) +} + +// UpdateFinishedAt sets the "finished_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateFinishedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateFinishedAt() + }) +} + +// ClearFinishedAt clears the value of the "finished_at" field. +func (u *BatchImageJobUpsertBulk) ClearFinishedAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearFinishedAt() + }) +} + +// SetSettledAt sets the "settled_at" field. +func (u *BatchImageJobUpsertBulk) SetSettledAt(v time.Time) *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.SetSettledAt(v) + }) +} + +// UpdateSettledAt sets the "settled_at" field to the value that was provided on create. +func (u *BatchImageJobUpsertBulk) UpdateSettledAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.UpdateSettledAt() + }) +} + +// ClearSettledAt clears the value of the "settled_at" field. +func (u *BatchImageJobUpsertBulk) ClearSettledAt() *BatchImageJobUpsertBulk { + return u.Update(func(s *BatchImageJobUpsert) { + s.ClearSettledAt() + }) +} + +// Exec executes the query. +func (u *BatchImageJobUpsertBulk) Exec(ctx context.Context) error { + if u.create.err != nil { + return u.create.err + } + for i, b := range u.create.builders { + if len(b.conflict) != 0 { + return fmt.Errorf("ent: OnConflict was set for builder %d. Set it on the BatchImageJobCreateBulk instead", i) + } + } + if len(u.create.conflict) == 0 { + return errors.New("ent: missing options for BatchImageJobCreateBulk.OnConflict") + } + return u.create.Exec(ctx) +} + +// ExecX is like Exec, but panics if an error occurs. +func (u *BatchImageJobUpsertBulk) ExecX(ctx context.Context) { + if err := u.create.Exec(ctx); err != nil { + panic(err) + } +} diff --git a/backend/ent/batchimagejob_delete.go b/backend/ent/batchimagejob_delete.go new file mode 100644 index 0000000000..da3dec2109 --- /dev/null +++ b/backend/ent/batchimagejob_delete.go @@ -0,0 +1,88 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageJobDelete is the builder for deleting a BatchImageJob entity. +type BatchImageJobDelete struct { + config + hooks []Hook + mutation *BatchImageJobMutation +} + +// Where appends a list predicates to the BatchImageJobDelete builder. +func (_d *BatchImageJobDelete) Where(ps ...predicate.BatchImageJob) *BatchImageJobDelete { + _d.mutation.Where(ps...) + return _d +} + +// Exec executes the deletion query and returns how many vertices were deleted. +func (_d *BatchImageJobDelete) Exec(ctx context.Context) (int, error) { + return withHooks(ctx, _d.sqlExec, _d.mutation, _d.hooks) +} + +// ExecX is like Exec, but panics if an error occurs. +func (_d *BatchImageJobDelete) ExecX(ctx context.Context) int { + n, err := _d.Exec(ctx) + if err != nil { + panic(err) + } + return n +} + +func (_d *BatchImageJobDelete) sqlExec(ctx context.Context) (int, error) { + _spec := sqlgraph.NewDeleteSpec(batchimagejob.Table, sqlgraph.NewFieldSpec(batchimagejob.FieldID, field.TypeInt64)) + if ps := _d.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + affected, err := sqlgraph.DeleteNodes(ctx, _d.driver, _spec) + if err != nil && sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + _d.mutation.done = true + return affected, err +} + +// BatchImageJobDeleteOne is the builder for deleting a single BatchImageJob entity. +type BatchImageJobDeleteOne struct { + _d *BatchImageJobDelete +} + +// Where appends a list predicates to the BatchImageJobDelete builder. +func (_d *BatchImageJobDeleteOne) Where(ps ...predicate.BatchImageJob) *BatchImageJobDeleteOne { + _d._d.mutation.Where(ps...) + return _d +} + +// Exec executes the deletion query. +func (_d *BatchImageJobDeleteOne) Exec(ctx context.Context) error { + n, err := _d._d.Exec(ctx) + switch { + case err != nil: + return err + case n == 0: + return &NotFoundError{batchimagejob.Label} + default: + return nil + } +} + +// ExecX is like Exec, but panics if an error occurs. +func (_d *BatchImageJobDeleteOne) ExecX(ctx context.Context) { + if err := _d.Exec(ctx); err != nil { + panic(err) + } +} diff --git a/backend/ent/batchimagejob_query.go b/backend/ent/batchimagejob_query.go new file mode 100644 index 0000000000..5ea4af42b2 --- /dev/null +++ b/backend/ent/batchimagejob_query.go @@ -0,0 +1,564 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "fmt" + "math" + + "entgo.io/ent" + "entgo.io/ent/dialect" + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageJobQuery is the builder for querying BatchImageJob entities. +type BatchImageJobQuery struct { + config + ctx *QueryContext + order []batchimagejob.OrderOption + inters []Interceptor + predicates []predicate.BatchImageJob + modifiers []func(*sql.Selector) + // intermediate query (i.e. traversal path). + sql *sql.Selector + path func(context.Context) (*sql.Selector, error) +} + +// Where adds a new predicate for the BatchImageJobQuery builder. +func (_q *BatchImageJobQuery) Where(ps ...predicate.BatchImageJob) *BatchImageJobQuery { + _q.predicates = append(_q.predicates, ps...) + return _q +} + +// Limit the number of records to be returned by this query. +func (_q *BatchImageJobQuery) Limit(limit int) *BatchImageJobQuery { + _q.ctx.Limit = &limit + return _q +} + +// Offset to start from. +func (_q *BatchImageJobQuery) Offset(offset int) *BatchImageJobQuery { + _q.ctx.Offset = &offset + return _q +} + +// Unique configures the query builder to filter duplicate records on query. +// By default, unique is set to true, and can be disabled using this method. +func (_q *BatchImageJobQuery) Unique(unique bool) *BatchImageJobQuery { + _q.ctx.Unique = &unique + return _q +} + +// Order specifies how the records should be ordered. +func (_q *BatchImageJobQuery) Order(o ...batchimagejob.OrderOption) *BatchImageJobQuery { + _q.order = append(_q.order, o...) + return _q +} + +// First returns the first BatchImageJob entity from the query. +// Returns a *NotFoundError when no BatchImageJob was found. +func (_q *BatchImageJobQuery) First(ctx context.Context) (*BatchImageJob, error) { + nodes, err := _q.Limit(1).All(setContextOp(ctx, _q.ctx, ent.OpQueryFirst)) + if err != nil { + return nil, err + } + if len(nodes) == 0 { + return nil, &NotFoundError{batchimagejob.Label} + } + return nodes[0], nil +} + +// FirstX is like First, but panics if an error occurs. +func (_q *BatchImageJobQuery) FirstX(ctx context.Context) *BatchImageJob { + node, err := _q.First(ctx) + if err != nil && !IsNotFound(err) { + panic(err) + } + return node +} + +// FirstID returns the first BatchImageJob ID from the query. +// Returns a *NotFoundError when no BatchImageJob ID was found. +func (_q *BatchImageJobQuery) FirstID(ctx context.Context) (id int64, err error) { + var ids []int64 + if ids, err = _q.Limit(1).IDs(setContextOp(ctx, _q.ctx, ent.OpQueryFirstID)); err != nil { + return + } + if len(ids) == 0 { + err = &NotFoundError{batchimagejob.Label} + return + } + return ids[0], nil +} + +// FirstIDX is like FirstID, but panics if an error occurs. +func (_q *BatchImageJobQuery) FirstIDX(ctx context.Context) int64 { + id, err := _q.FirstID(ctx) + if err != nil && !IsNotFound(err) { + panic(err) + } + return id +} + +// Only returns a single BatchImageJob entity found by the query, ensuring it only returns one. +// Returns a *NotSingularError when more than one BatchImageJob entity is found. +// Returns a *NotFoundError when no BatchImageJob entities are found. +func (_q *BatchImageJobQuery) Only(ctx context.Context) (*BatchImageJob, error) { + nodes, err := _q.Limit(2).All(setContextOp(ctx, _q.ctx, ent.OpQueryOnly)) + if err != nil { + return nil, err + } + switch len(nodes) { + case 1: + return nodes[0], nil + case 0: + return nil, &NotFoundError{batchimagejob.Label} + default: + return nil, &NotSingularError{batchimagejob.Label} + } +} + +// OnlyX is like Only, but panics if an error occurs. +func (_q *BatchImageJobQuery) OnlyX(ctx context.Context) *BatchImageJob { + node, err := _q.Only(ctx) + if err != nil { + panic(err) + } + return node +} + +// OnlyID is like Only, but returns the only BatchImageJob ID in the query. +// Returns a *NotSingularError when more than one BatchImageJob ID is found. +// Returns a *NotFoundError when no entities are found. +func (_q *BatchImageJobQuery) OnlyID(ctx context.Context) (id int64, err error) { + var ids []int64 + if ids, err = _q.Limit(2).IDs(setContextOp(ctx, _q.ctx, ent.OpQueryOnlyID)); err != nil { + return + } + switch len(ids) { + case 1: + id = ids[0] + case 0: + err = &NotFoundError{batchimagejob.Label} + default: + err = &NotSingularError{batchimagejob.Label} + } + return +} + +// OnlyIDX is like OnlyID, but panics if an error occurs. +func (_q *BatchImageJobQuery) OnlyIDX(ctx context.Context) int64 { + id, err := _q.OnlyID(ctx) + if err != nil { + panic(err) + } + return id +} + +// All executes the query and returns a list of BatchImageJobs. +func (_q *BatchImageJobQuery) All(ctx context.Context) ([]*BatchImageJob, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryAll) + if err := _q.prepareQuery(ctx); err != nil { + return nil, err + } + qr := querierAll[[]*BatchImageJob, *BatchImageJobQuery]() + return withInterceptors[[]*BatchImageJob](ctx, _q, qr, _q.inters) +} + +// AllX is like All, but panics if an error occurs. +func (_q *BatchImageJobQuery) AllX(ctx context.Context) []*BatchImageJob { + nodes, err := _q.All(ctx) + if err != nil { + panic(err) + } + return nodes +} + +// IDs executes the query and returns a list of BatchImageJob IDs. +func (_q *BatchImageJobQuery) IDs(ctx context.Context) (ids []int64, err error) { + if _q.ctx.Unique == nil && _q.path != nil { + _q.Unique(true) + } + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryIDs) + if err = _q.Select(batchimagejob.FieldID).Scan(ctx, &ids); err != nil { + return nil, err + } + return ids, nil +} + +// IDsX is like IDs, but panics if an error occurs. +func (_q *BatchImageJobQuery) IDsX(ctx context.Context) []int64 { + ids, err := _q.IDs(ctx) + if err != nil { + panic(err) + } + return ids +} + +// Count returns the count of the given query. +func (_q *BatchImageJobQuery) Count(ctx context.Context) (int, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryCount) + if err := _q.prepareQuery(ctx); err != nil { + return 0, err + } + return withInterceptors[int](ctx, _q, querierCount[*BatchImageJobQuery](), _q.inters) +} + +// CountX is like Count, but panics if an error occurs. +func (_q *BatchImageJobQuery) CountX(ctx context.Context) int { + count, err := _q.Count(ctx) + if err != nil { + panic(err) + } + return count +} + +// Exist returns true if the query has elements in the graph. +func (_q *BatchImageJobQuery) Exist(ctx context.Context) (bool, error) { + ctx = setContextOp(ctx, _q.ctx, ent.OpQueryExist) + switch _, err := _q.FirstID(ctx); { + case IsNotFound(err): + return false, nil + case err != nil: + return false, fmt.Errorf("ent: check existence: %w", err) + default: + return true, nil + } +} + +// ExistX is like Exist, but panics if an error occurs. +func (_q *BatchImageJobQuery) ExistX(ctx context.Context) bool { + exist, err := _q.Exist(ctx) + if err != nil { + panic(err) + } + return exist +} + +// Clone returns a duplicate of the BatchImageJobQuery builder, including all associated steps. It can be +// used to prepare common query builders and use them differently after the clone is made. +func (_q *BatchImageJobQuery) Clone() *BatchImageJobQuery { + if _q == nil { + return nil + } + return &BatchImageJobQuery{ + config: _q.config, + ctx: _q.ctx.Clone(), + order: append([]batchimagejob.OrderOption{}, _q.order...), + inters: append([]Interceptor{}, _q.inters...), + predicates: append([]predicate.BatchImageJob{}, _q.predicates...), + // clone intermediate query. + sql: _q.sql.Clone(), + path: _q.path, + } +} + +// GroupBy is used to group vertices by one or more fields/columns. +// It is often used with aggregate functions, like: count, max, mean, min, sum. +// +// Example: +// +// var v []struct { +// BatchID string `json:"batch_id,omitempty"` +// Count int `json:"count,omitempty"` +// } +// +// client.BatchImageJob.Query(). +// GroupBy(batchimagejob.FieldBatchID). +// Aggregate(ent.Count()). +// Scan(ctx, &v) +func (_q *BatchImageJobQuery) GroupBy(field string, fields ...string) *BatchImageJobGroupBy { + _q.ctx.Fields = append([]string{field}, fields...) + grbuild := &BatchImageJobGroupBy{build: _q} + grbuild.flds = &_q.ctx.Fields + grbuild.label = batchimagejob.Label + grbuild.scan = grbuild.Scan + return grbuild +} + +// Select allows the selection one or more fields/columns for the given query, +// instead of selecting all fields in the entity. +// +// Example: +// +// var v []struct { +// BatchID string `json:"batch_id,omitempty"` +// } +// +// client.BatchImageJob.Query(). +// Select(batchimagejob.FieldBatchID). +// Scan(ctx, &v) +func (_q *BatchImageJobQuery) Select(fields ...string) *BatchImageJobSelect { + _q.ctx.Fields = append(_q.ctx.Fields, fields...) + sbuild := &BatchImageJobSelect{BatchImageJobQuery: _q} + sbuild.label = batchimagejob.Label + sbuild.flds, sbuild.scan = &_q.ctx.Fields, sbuild.Scan + return sbuild +} + +// Aggregate returns a BatchImageJobSelect configured with the given aggregations. +func (_q *BatchImageJobQuery) Aggregate(fns ...AggregateFunc) *BatchImageJobSelect { + return _q.Select().Aggregate(fns...) +} + +func (_q *BatchImageJobQuery) prepareQuery(ctx context.Context) error { + for _, inter := range _q.inters { + if inter == nil { + return fmt.Errorf("ent: uninitialized interceptor (forgotten import ent/runtime?)") + } + if trv, ok := inter.(Traverser); ok { + if err := trv.Traverse(ctx, _q); err != nil { + return err + } + } + } + for _, f := range _q.ctx.Fields { + if !batchimagejob.ValidColumn(f) { + return &ValidationError{Name: f, err: fmt.Errorf("ent: invalid field %q for query", f)} + } + } + if _q.path != nil { + prev, err := _q.path(ctx) + if err != nil { + return err + } + _q.sql = prev + } + return nil +} + +func (_q *BatchImageJobQuery) sqlAll(ctx context.Context, hooks ...queryHook) ([]*BatchImageJob, error) { + var ( + nodes = []*BatchImageJob{} + _spec = _q.querySpec() + ) + _spec.ScanValues = func(columns []string) ([]any, error) { + return (*BatchImageJob).scanValues(nil, columns) + } + _spec.Assign = func(columns []string, values []any) error { + node := &BatchImageJob{config: _q.config} + nodes = append(nodes, node) + return node.assignValues(columns, values) + } + if len(_q.modifiers) > 0 { + _spec.Modifiers = _q.modifiers + } + for i := range hooks { + hooks[i](ctx, _spec) + } + if err := sqlgraph.QueryNodes(ctx, _q.driver, _spec); err != nil { + return nil, err + } + if len(nodes) == 0 { + return nodes, nil + } + return nodes, nil +} + +func (_q *BatchImageJobQuery) sqlCount(ctx context.Context) (int, error) { + _spec := _q.querySpec() + if len(_q.modifiers) > 0 { + _spec.Modifiers = _q.modifiers + } + _spec.Node.Columns = _q.ctx.Fields + if len(_q.ctx.Fields) > 0 { + _spec.Unique = _q.ctx.Unique != nil && *_q.ctx.Unique + } + return sqlgraph.CountNodes(ctx, _q.driver, _spec) +} + +func (_q *BatchImageJobQuery) querySpec() *sqlgraph.QuerySpec { + _spec := sqlgraph.NewQuerySpec(batchimagejob.Table, batchimagejob.Columns, sqlgraph.NewFieldSpec(batchimagejob.FieldID, field.TypeInt64)) + _spec.From = _q.sql + if unique := _q.ctx.Unique; unique != nil { + _spec.Unique = *unique + } else if _q.path != nil { + _spec.Unique = true + } + if fields := _q.ctx.Fields; len(fields) > 0 { + _spec.Node.Columns = make([]string, 0, len(fields)) + _spec.Node.Columns = append(_spec.Node.Columns, batchimagejob.FieldID) + for i := range fields { + if fields[i] != batchimagejob.FieldID { + _spec.Node.Columns = append(_spec.Node.Columns, fields[i]) + } + } + } + if ps := _q.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if limit := _q.ctx.Limit; limit != nil { + _spec.Limit = *limit + } + if offset := _q.ctx.Offset; offset != nil { + _spec.Offset = *offset + } + if ps := _q.order; len(ps) > 0 { + _spec.Order = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + return _spec +} + +func (_q *BatchImageJobQuery) sqlQuery(ctx context.Context) *sql.Selector { + builder := sql.Dialect(_q.driver.Dialect()) + t1 := builder.Table(batchimagejob.Table) + columns := _q.ctx.Fields + if len(columns) == 0 { + columns = batchimagejob.Columns + } + selector := builder.Select(t1.Columns(columns...)...).From(t1) + if _q.sql != nil { + selector = _q.sql + selector.Select(selector.Columns(columns...)...) + } + if _q.ctx.Unique != nil && *_q.ctx.Unique { + selector.Distinct() + } + for _, m := range _q.modifiers { + m(selector) + } + for _, p := range _q.predicates { + p(selector) + } + for _, p := range _q.order { + p(selector) + } + if offset := _q.ctx.Offset; offset != nil { + // limit is mandatory for offset clause. We start + // with default value, and override it below if needed. + selector.Offset(*offset).Limit(math.MaxInt32) + } + if limit := _q.ctx.Limit; limit != nil { + selector.Limit(*limit) + } + return selector +} + +// ForUpdate locks the selected rows against concurrent updates, and prevent them from being +// updated, deleted or "selected ... for update" by other sessions, until the transaction is +// either committed or rolled-back. +func (_q *BatchImageJobQuery) ForUpdate(opts ...sql.LockOption) *BatchImageJobQuery { + if _q.driver.Dialect() == dialect.Postgres { + _q.Unique(false) + } + _q.modifiers = append(_q.modifiers, func(s *sql.Selector) { + s.ForUpdate(opts...) + }) + return _q +} + +// ForShare behaves similarly to ForUpdate, except that it acquires a shared mode lock +// on any rows that are read. Other sessions can read the rows, but cannot modify them +// until your transaction commits. +func (_q *BatchImageJobQuery) ForShare(opts ...sql.LockOption) *BatchImageJobQuery { + if _q.driver.Dialect() == dialect.Postgres { + _q.Unique(false) + } + _q.modifiers = append(_q.modifiers, func(s *sql.Selector) { + s.ForShare(opts...) + }) + return _q +} + +// BatchImageJobGroupBy is the group-by builder for BatchImageJob entities. +type BatchImageJobGroupBy struct { + selector + build *BatchImageJobQuery +} + +// Aggregate adds the given aggregation functions to the group-by query. +func (_g *BatchImageJobGroupBy) Aggregate(fns ...AggregateFunc) *BatchImageJobGroupBy { + _g.fns = append(_g.fns, fns...) + return _g +} + +// Scan applies the selector query and scans the result into the given value. +func (_g *BatchImageJobGroupBy) Scan(ctx context.Context, v any) error { + ctx = setContextOp(ctx, _g.build.ctx, ent.OpQueryGroupBy) + if err := _g.build.prepareQuery(ctx); err != nil { + return err + } + return scanWithInterceptors[*BatchImageJobQuery, *BatchImageJobGroupBy](ctx, _g.build, _g, _g.build.inters, v) +} + +func (_g *BatchImageJobGroupBy) sqlScan(ctx context.Context, root *BatchImageJobQuery, v any) error { + selector := root.sqlQuery(ctx).Select() + aggregation := make([]string, 0, len(_g.fns)) + for _, fn := range _g.fns { + aggregation = append(aggregation, fn(selector)) + } + if len(selector.SelectedColumns()) == 0 { + columns := make([]string, 0, len(*_g.flds)+len(_g.fns)) + for _, f := range *_g.flds { + columns = append(columns, selector.C(f)) + } + columns = append(columns, aggregation...) + selector.Select(columns...) + } + selector.GroupBy(selector.Columns(*_g.flds...)...) + if err := selector.Err(); err != nil { + return err + } + rows := &sql.Rows{} + query, args := selector.Query() + if err := _g.build.driver.Query(ctx, query, args, rows); err != nil { + return err + } + defer rows.Close() + return sql.ScanSlice(rows, v) +} + +// BatchImageJobSelect is the builder for selecting fields of BatchImageJob entities. +type BatchImageJobSelect struct { + *BatchImageJobQuery + selector +} + +// Aggregate adds the given aggregation functions to the selector query. +func (_s *BatchImageJobSelect) Aggregate(fns ...AggregateFunc) *BatchImageJobSelect { + _s.fns = append(_s.fns, fns...) + return _s +} + +// Scan applies the selector query and scans the result into the given value. +func (_s *BatchImageJobSelect) Scan(ctx context.Context, v any) error { + ctx = setContextOp(ctx, _s.ctx, ent.OpQuerySelect) + if err := _s.prepareQuery(ctx); err != nil { + return err + } + return scanWithInterceptors[*BatchImageJobQuery, *BatchImageJobSelect](ctx, _s.BatchImageJobQuery, _s, _s.inters, v) +} + +func (_s *BatchImageJobSelect) sqlScan(ctx context.Context, root *BatchImageJobQuery, v any) error { + selector := root.sqlQuery(ctx) + aggregation := make([]string, 0, len(_s.fns)) + for _, fn := range _s.fns { + aggregation = append(aggregation, fn(selector)) + } + switch n := len(*_s.selector.flds); { + case n == 0 && len(aggregation) > 0: + selector.Select(aggregation...) + case n != 0 && len(aggregation) > 0: + selector.AppendSelect(aggregation...) + } + rows := &sql.Rows{} + query, args := selector.Query() + if err := _s.driver.Query(ctx, query, args, rows); err != nil { + return err + } + defer rows.Close() + return sql.ScanSlice(rows, v) +} diff --git a/backend/ent/batchimagejob_update.go b/backend/ent/batchimagejob_update.go new file mode 100644 index 0000000000..8df7302500 --- /dev/null +++ b/backend/ent/batchimagejob_update.go @@ -0,0 +1,2308 @@ +// Code generated by ent, DO NOT EDIT. + +package ent + +import ( + "context" + "errors" + "fmt" + "time" + + "entgo.io/ent/dialect/sql" + "entgo.io/ent/dialect/sql/sqlgraph" + "entgo.io/ent/schema/field" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" + "github.com/Wei-Shaw/sub2api/ent/predicate" +) + +// BatchImageJobUpdate is the builder for updating BatchImageJob entities. +type BatchImageJobUpdate struct { + config + hooks []Hook + mutation *BatchImageJobMutation +} + +// Where appends a list predicates to the BatchImageJobUpdate builder. +func (_u *BatchImageJobUpdate) Where(ps ...predicate.BatchImageJob) *BatchImageJobUpdate { + _u.mutation.Where(ps...) + return _u +} + +// SetUserID sets the "user_id" field. +func (_u *BatchImageJobUpdate) SetUserID(v int64) *BatchImageJobUpdate { + _u.mutation.ResetUserID() + _u.mutation.SetUserID(v) + return _u +} + +// SetNillableUserID sets the "user_id" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableUserID(v *int64) *BatchImageJobUpdate { + if v != nil { + _u.SetUserID(*v) + } + return _u +} + +// AddUserID adds value to the "user_id" field. +func (_u *BatchImageJobUpdate) AddUserID(v int64) *BatchImageJobUpdate { + _u.mutation.AddUserID(v) + return _u +} + +// SetAPIKeyID sets the "api_key_id" field. +func (_u *BatchImageJobUpdate) SetAPIKeyID(v int64) *BatchImageJobUpdate { + _u.mutation.ResetAPIKeyID() + _u.mutation.SetAPIKeyID(v) + return _u +} + +// SetNillableAPIKeyID sets the "api_key_id" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableAPIKeyID(v *int64) *BatchImageJobUpdate { + if v != nil { + _u.SetAPIKeyID(*v) + } + return _u +} + +// AddAPIKeyID adds value to the "api_key_id" field. +func (_u *BatchImageJobUpdate) AddAPIKeyID(v int64) *BatchImageJobUpdate { + _u.mutation.AddAPIKeyID(v) + return _u +} + +// ClearAPIKeyID clears the value of the "api_key_id" field. +func (_u *BatchImageJobUpdate) ClearAPIKeyID() *BatchImageJobUpdate { + _u.mutation.ClearAPIKeyID() + return _u +} + +// SetAccountID sets the "account_id" field. +func (_u *BatchImageJobUpdate) SetAccountID(v int64) *BatchImageJobUpdate { + _u.mutation.ResetAccountID() + _u.mutation.SetAccountID(v) + return _u +} + +// SetNillableAccountID sets the "account_id" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableAccountID(v *int64) *BatchImageJobUpdate { + if v != nil { + _u.SetAccountID(*v) + } + return _u +} + +// AddAccountID adds value to the "account_id" field. +func (_u *BatchImageJobUpdate) AddAccountID(v int64) *BatchImageJobUpdate { + _u.mutation.AddAccountID(v) + return _u +} + +// ClearAccountID clears the value of the "account_id" field. +func (_u *BatchImageJobUpdate) ClearAccountID() *BatchImageJobUpdate { + _u.mutation.ClearAccountID() + return _u +} + +// SetProvider sets the "provider" field. +func (_u *BatchImageJobUpdate) SetProvider(v string) *BatchImageJobUpdate { + _u.mutation.SetProvider(v) + return _u +} + +// SetNillableProvider sets the "provider" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableProvider(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetProvider(*v) + } + return _u +} + +// SetModel sets the "model" field. +func (_u *BatchImageJobUpdate) SetModel(v string) *BatchImageJobUpdate { + _u.mutation.SetModel(v) + return _u +} + +// SetNillableModel sets the "model" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableModel(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetModel(*v) + } + return _u +} + +// SetTaskName sets the "task_name" field. +func (_u *BatchImageJobUpdate) SetTaskName(v string) *BatchImageJobUpdate { + _u.mutation.SetTaskName(v) + return _u +} + +// SetNillableTaskName sets the "task_name" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableTaskName(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetTaskName(*v) + } + return _u +} + +// SetStatus sets the "status" field. +func (_u *BatchImageJobUpdate) SetStatus(v string) *BatchImageJobUpdate { + _u.mutation.SetStatus(v) + return _u +} + +// SetNillableStatus sets the "status" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableStatus(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetStatus(*v) + } + return _u +} + +// SetProviderJobName sets the "provider_job_name" field. +func (_u *BatchImageJobUpdate) SetProviderJobName(v string) *BatchImageJobUpdate { + _u.mutation.SetProviderJobName(v) + return _u +} + +// SetNillableProviderJobName sets the "provider_job_name" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableProviderJobName(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetProviderJobName(*v) + } + return _u +} + +// ClearProviderJobName clears the value of the "provider_job_name" field. +func (_u *BatchImageJobUpdate) ClearProviderJobName() *BatchImageJobUpdate { + _u.mutation.ClearProviderJobName() + return _u +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (_u *BatchImageJobUpdate) SetProviderInputRef(v string) *BatchImageJobUpdate { + _u.mutation.SetProviderInputRef(v) + return _u +} + +// SetNillableProviderInputRef sets the "provider_input_ref" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableProviderInputRef(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetProviderInputRef(*v) + } + return _u +} + +// ClearProviderInputRef clears the value of the "provider_input_ref" field. +func (_u *BatchImageJobUpdate) ClearProviderInputRef() *BatchImageJobUpdate { + _u.mutation.ClearProviderInputRef() + return _u +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (_u *BatchImageJobUpdate) SetProviderOutputRef(v string) *BatchImageJobUpdate { + _u.mutation.SetProviderOutputRef(v) + return _u +} + +// SetNillableProviderOutputRef sets the "provider_output_ref" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableProviderOutputRef(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetProviderOutputRef(*v) + } + return _u +} + +// ClearProviderOutputRef clears the value of the "provider_output_ref" field. +func (_u *BatchImageJobUpdate) ClearProviderOutputRef() *BatchImageJobUpdate { + _u.mutation.ClearProviderOutputRef() + return _u +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (_u *BatchImageJobUpdate) SetGcsInputURI(v string) *BatchImageJobUpdate { + _u.mutation.SetGcsInputURI(v) + return _u +} + +// SetNillableGcsInputURI sets the "gcs_input_uri" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableGcsInputURI(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetGcsInputURI(*v) + } + return _u +} + +// ClearGcsInputURI clears the value of the "gcs_input_uri" field. +func (_u *BatchImageJobUpdate) ClearGcsInputURI() *BatchImageJobUpdate { + _u.mutation.ClearGcsInputURI() + return _u +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (_u *BatchImageJobUpdate) SetGcsOutputURI(v string) *BatchImageJobUpdate { + _u.mutation.SetGcsOutputURI(v) + return _u +} + +// SetNillableGcsOutputURI sets the "gcs_output_uri" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableGcsOutputURI(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetGcsOutputURI(*v) + } + return _u +} + +// ClearGcsOutputURI clears the value of the "gcs_output_uri" field. +func (_u *BatchImageJobUpdate) ClearGcsOutputURI() *BatchImageJobUpdate { + _u.mutation.ClearGcsOutputURI() + return _u +} + +// SetItemCount sets the "item_count" field. +func (_u *BatchImageJobUpdate) SetItemCount(v int) *BatchImageJobUpdate { + _u.mutation.ResetItemCount() + _u.mutation.SetItemCount(v) + return _u +} + +// SetNillableItemCount sets the "item_count" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableItemCount(v *int) *BatchImageJobUpdate { + if v != nil { + _u.SetItemCount(*v) + } + return _u +} + +// AddItemCount adds value to the "item_count" field. +func (_u *BatchImageJobUpdate) AddItemCount(v int) *BatchImageJobUpdate { + _u.mutation.AddItemCount(v) + return _u +} + +// SetSuccessCount sets the "success_count" field. +func (_u *BatchImageJobUpdate) SetSuccessCount(v int) *BatchImageJobUpdate { + _u.mutation.ResetSuccessCount() + _u.mutation.SetSuccessCount(v) + return _u +} + +// SetNillableSuccessCount sets the "success_count" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableSuccessCount(v *int) *BatchImageJobUpdate { + if v != nil { + _u.SetSuccessCount(*v) + } + return _u +} + +// AddSuccessCount adds value to the "success_count" field. +func (_u *BatchImageJobUpdate) AddSuccessCount(v int) *BatchImageJobUpdate { + _u.mutation.AddSuccessCount(v) + return _u +} + +// SetFailCount sets the "fail_count" field. +func (_u *BatchImageJobUpdate) SetFailCount(v int) *BatchImageJobUpdate { + _u.mutation.ResetFailCount() + _u.mutation.SetFailCount(v) + return _u +} + +// SetNillableFailCount sets the "fail_count" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableFailCount(v *int) *BatchImageJobUpdate { + if v != nil { + _u.SetFailCount(*v) + } + return _u +} + +// AddFailCount adds value to the "fail_count" field. +func (_u *BatchImageJobUpdate) AddFailCount(v int) *BatchImageJobUpdate { + _u.mutation.AddFailCount(v) + return _u +} + +// SetCancelledCount sets the "cancelled_count" field. +func (_u *BatchImageJobUpdate) SetCancelledCount(v int) *BatchImageJobUpdate { + _u.mutation.ResetCancelledCount() + _u.mutation.SetCancelledCount(v) + return _u +} + +// SetNillableCancelledCount sets the "cancelled_count" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableCancelledCount(v *int) *BatchImageJobUpdate { + if v != nil { + _u.SetCancelledCount(*v) + } + return _u +} + +// AddCancelledCount adds value to the "cancelled_count" field. +func (_u *BatchImageJobUpdate) AddCancelledCount(v int) *BatchImageJobUpdate { + _u.mutation.AddCancelledCount(v) + return _u +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (_u *BatchImageJobUpdate) SetEstimatedCost(v float64) *BatchImageJobUpdate { + _u.mutation.ResetEstimatedCost() + _u.mutation.SetEstimatedCost(v) + return _u +} + +// SetNillableEstimatedCost sets the "estimated_cost" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableEstimatedCost(v *float64) *BatchImageJobUpdate { + if v != nil { + _u.SetEstimatedCost(*v) + } + return _u +} + +// AddEstimatedCost adds value to the "estimated_cost" field. +func (_u *BatchImageJobUpdate) AddEstimatedCost(v float64) *BatchImageJobUpdate { + _u.mutation.AddEstimatedCost(v) + return _u +} + +// SetHoldAmount sets the "hold_amount" field. +func (_u *BatchImageJobUpdate) SetHoldAmount(v float64) *BatchImageJobUpdate { + _u.mutation.ResetHoldAmount() + _u.mutation.SetHoldAmount(v) + return _u +} + +// SetNillableHoldAmount sets the "hold_amount" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableHoldAmount(v *float64) *BatchImageJobUpdate { + if v != nil { + _u.SetHoldAmount(*v) + } + return _u +} + +// AddHoldAmount adds value to the "hold_amount" field. +func (_u *BatchImageJobUpdate) AddHoldAmount(v float64) *BatchImageJobUpdate { + _u.mutation.AddHoldAmount(v) + return _u +} + +// ClearHoldAmount clears the value of the "hold_amount" field. +func (_u *BatchImageJobUpdate) ClearHoldAmount() *BatchImageJobUpdate { + _u.mutation.ClearHoldAmount() + return _u +} + +// SetActualCost sets the "actual_cost" field. +func (_u *BatchImageJobUpdate) SetActualCost(v float64) *BatchImageJobUpdate { + _u.mutation.ResetActualCost() + _u.mutation.SetActualCost(v) + return _u +} + +// SetNillableActualCost sets the "actual_cost" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableActualCost(v *float64) *BatchImageJobUpdate { + if v != nil { + _u.SetActualCost(*v) + } + return _u +} + +// AddActualCost adds value to the "actual_cost" field. +func (_u *BatchImageJobUpdate) AddActualCost(v float64) *BatchImageJobUpdate { + _u.mutation.AddActualCost(v) + return _u +} + +// ClearActualCost clears the value of the "actual_cost" field. +func (_u *BatchImageJobUpdate) ClearActualCost() *BatchImageJobUpdate { + _u.mutation.ClearActualCost() + return _u +} + +// SetCurrency sets the "currency" field. +func (_u *BatchImageJobUpdate) SetCurrency(v string) *BatchImageJobUpdate { + _u.mutation.SetCurrency(v) + return _u +} + +// SetNillableCurrency sets the "currency" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableCurrency(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetCurrency(*v) + } + return _u +} + +// SetHoldID sets the "hold_id" field. +func (_u *BatchImageJobUpdate) SetHoldID(v string) *BatchImageJobUpdate { + _u.mutation.SetHoldID(v) + return _u +} + +// SetNillableHoldID sets the "hold_id" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableHoldID(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetHoldID(*v) + } + return _u +} + +// ClearHoldID clears the value of the "hold_id" field. +func (_u *BatchImageJobUpdate) ClearHoldID() *BatchImageJobUpdate { + _u.mutation.ClearHoldID() + return _u +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (_u *BatchImageJobUpdate) SetIdempotencyKey(v string) *BatchImageJobUpdate { + _u.mutation.SetIdempotencyKey(v) + return _u +} + +// SetNillableIdempotencyKey sets the "idempotency_key" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableIdempotencyKey(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetIdempotencyKey(*v) + } + return _u +} + +// ClearIdempotencyKey clears the value of the "idempotency_key" field. +func (_u *BatchImageJobUpdate) ClearIdempotencyKey() *BatchImageJobUpdate { + _u.mutation.ClearIdempotencyKey() + return _u +} + +// SetRequestHash sets the "request_hash" field. +func (_u *BatchImageJobUpdate) SetRequestHash(v string) *BatchImageJobUpdate { + _u.mutation.SetRequestHash(v) + return _u +} + +// SetNillableRequestHash sets the "request_hash" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableRequestHash(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetRequestHash(*v) + } + return _u +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (_u *BatchImageJobUpdate) ClearRequestHash() *BatchImageJobUpdate { + _u.mutation.ClearRequestHash() + return _u +} + +// SetManifestHash sets the "manifest_hash" field. +func (_u *BatchImageJobUpdate) SetManifestHash(v string) *BatchImageJobUpdate { + _u.mutation.SetManifestHash(v) + return _u +} + +// SetNillableManifestHash sets the "manifest_hash" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableManifestHash(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetManifestHash(*v) + } + return _u +} + +// ClearManifestHash clears the value of the "manifest_hash" field. +func (_u *BatchImageJobUpdate) ClearManifestHash() *BatchImageJobUpdate { + _u.mutation.ClearManifestHash() + return _u +} + +// SetRetryCount sets the "retry_count" field. +func (_u *BatchImageJobUpdate) SetRetryCount(v int) *BatchImageJobUpdate { + _u.mutation.ResetRetryCount() + _u.mutation.SetRetryCount(v) + return _u +} + +// SetNillableRetryCount sets the "retry_count" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableRetryCount(v *int) *BatchImageJobUpdate { + if v != nil { + _u.SetRetryCount(*v) + } + return _u +} + +// AddRetryCount adds value to the "retry_count" field. +func (_u *BatchImageJobUpdate) AddRetryCount(v int) *BatchImageJobUpdate { + _u.mutation.AddRetryCount(v) + return _u +} + +// SetVersion sets the "version" field. +func (_u *BatchImageJobUpdate) SetVersion(v int) *BatchImageJobUpdate { + _u.mutation.ResetVersion() + _u.mutation.SetVersion(v) + return _u +} + +// SetNillableVersion sets the "version" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableVersion(v *int) *BatchImageJobUpdate { + if v != nil { + _u.SetVersion(*v) + } + return _u +} + +// AddVersion adds value to the "version" field. +func (_u *BatchImageJobUpdate) AddVersion(v int) *BatchImageJobUpdate { + _u.mutation.AddVersion(v) + return _u +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (_u *BatchImageJobUpdate) SetOutputExpiresAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetOutputExpiresAt(v) + return _u +} + +// SetNillableOutputExpiresAt sets the "output_expires_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableOutputExpiresAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetOutputExpiresAt(*v) + } + return _u +} + +// ClearOutputExpiresAt clears the value of the "output_expires_at" field. +func (_u *BatchImageJobUpdate) ClearOutputExpiresAt() *BatchImageJobUpdate { + _u.mutation.ClearOutputExpiresAt() + return _u +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (_u *BatchImageJobUpdate) SetInputDeletedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetInputDeletedAt(v) + return _u +} + +// SetNillableInputDeletedAt sets the "input_deleted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableInputDeletedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetInputDeletedAt(*v) + } + return _u +} + +// ClearInputDeletedAt clears the value of the "input_deleted_at" field. +func (_u *BatchImageJobUpdate) ClearInputDeletedAt() *BatchImageJobUpdate { + _u.mutation.ClearInputDeletedAt() + return _u +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (_u *BatchImageJobUpdate) SetOutputDeletedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetOutputDeletedAt(v) + return _u +} + +// SetNillableOutputDeletedAt sets the "output_deleted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableOutputDeletedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetOutputDeletedAt(*v) + } + return _u +} + +// ClearOutputDeletedAt clears the value of the "output_deleted_at" field. +func (_u *BatchImageJobUpdate) ClearOutputDeletedAt() *BatchImageJobUpdate { + _u.mutation.ClearOutputDeletedAt() + return _u +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (_u *BatchImageJobUpdate) SetDownloadedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetDownloadedAt(v) + return _u +} + +// SetNillableDownloadedAt sets the "downloaded_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableDownloadedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetDownloadedAt(*v) + } + return _u +} + +// ClearDownloadedAt clears the value of the "downloaded_at" field. +func (_u *BatchImageJobUpdate) ClearDownloadedAt() *BatchImageJobUpdate { + _u.mutation.ClearDownloadedAt() + return _u +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (_u *BatchImageJobUpdate) SetUserDeletedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetUserDeletedAt(v) + return _u +} + +// SetNillableUserDeletedAt sets the "user_deleted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableUserDeletedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetUserDeletedAt(*v) + } + return _u +} + +// ClearUserDeletedAt clears the value of the "user_deleted_at" field. +func (_u *BatchImageJobUpdate) ClearUserDeletedAt() *BatchImageJobUpdate { + _u.mutation.ClearUserDeletedAt() + return _u +} + +// SetLastErrorCode sets the "last_error_code" field. +func (_u *BatchImageJobUpdate) SetLastErrorCode(v string) *BatchImageJobUpdate { + _u.mutation.SetLastErrorCode(v) + return _u +} + +// SetNillableLastErrorCode sets the "last_error_code" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableLastErrorCode(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetLastErrorCode(*v) + } + return _u +} + +// ClearLastErrorCode clears the value of the "last_error_code" field. +func (_u *BatchImageJobUpdate) ClearLastErrorCode() *BatchImageJobUpdate { + _u.mutation.ClearLastErrorCode() + return _u +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (_u *BatchImageJobUpdate) SetLastErrorMessage(v string) *BatchImageJobUpdate { + _u.mutation.SetLastErrorMessage(v) + return _u +} + +// SetNillableLastErrorMessage sets the "last_error_message" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableLastErrorMessage(v *string) *BatchImageJobUpdate { + if v != nil { + _u.SetLastErrorMessage(*v) + } + return _u +} + +// ClearLastErrorMessage clears the value of the "last_error_message" field. +func (_u *BatchImageJobUpdate) ClearLastErrorMessage() *BatchImageJobUpdate { + _u.mutation.ClearLastErrorMessage() + return _u +} + +// SetUpdatedAt sets the "updated_at" field. +func (_u *BatchImageJobUpdate) SetUpdatedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetUpdatedAt(v) + return _u +} + +// SetSubmittedAt sets the "submitted_at" field. +func (_u *BatchImageJobUpdate) SetSubmittedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetSubmittedAt(v) + return _u +} + +// SetNillableSubmittedAt sets the "submitted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableSubmittedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetSubmittedAt(*v) + } + return _u +} + +// ClearSubmittedAt clears the value of the "submitted_at" field. +func (_u *BatchImageJobUpdate) ClearSubmittedAt() *BatchImageJobUpdate { + _u.mutation.ClearSubmittedAt() + return _u +} + +// SetStartedAt sets the "started_at" field. +func (_u *BatchImageJobUpdate) SetStartedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetStartedAt(v) + return _u +} + +// SetNillableStartedAt sets the "started_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableStartedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetStartedAt(*v) + } + return _u +} + +// ClearStartedAt clears the value of the "started_at" field. +func (_u *BatchImageJobUpdate) ClearStartedAt() *BatchImageJobUpdate { + _u.mutation.ClearStartedAt() + return _u +} + +// SetFinishedAt sets the "finished_at" field. +func (_u *BatchImageJobUpdate) SetFinishedAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetFinishedAt(v) + return _u +} + +// SetNillableFinishedAt sets the "finished_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableFinishedAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetFinishedAt(*v) + } + return _u +} + +// ClearFinishedAt clears the value of the "finished_at" field. +func (_u *BatchImageJobUpdate) ClearFinishedAt() *BatchImageJobUpdate { + _u.mutation.ClearFinishedAt() + return _u +} + +// SetSettledAt sets the "settled_at" field. +func (_u *BatchImageJobUpdate) SetSettledAt(v time.Time) *BatchImageJobUpdate { + _u.mutation.SetSettledAt(v) + return _u +} + +// SetNillableSettledAt sets the "settled_at" field if the given value is not nil. +func (_u *BatchImageJobUpdate) SetNillableSettledAt(v *time.Time) *BatchImageJobUpdate { + if v != nil { + _u.SetSettledAt(*v) + } + return _u +} + +// ClearSettledAt clears the value of the "settled_at" field. +func (_u *BatchImageJobUpdate) ClearSettledAt() *BatchImageJobUpdate { + _u.mutation.ClearSettledAt() + return _u +} + +// Mutation returns the BatchImageJobMutation object of the builder. +func (_u *BatchImageJobUpdate) Mutation() *BatchImageJobMutation { + return _u.mutation +} + +// Save executes the query and returns the number of nodes affected by the update operation. +func (_u *BatchImageJobUpdate) Save(ctx context.Context) (int, error) { + _u.defaults() + return withHooks(ctx, _u.sqlSave, _u.mutation, _u.hooks) +} + +// SaveX is like Save, but panics if an error occurs. +func (_u *BatchImageJobUpdate) SaveX(ctx context.Context) int { + affected, err := _u.Save(ctx) + if err != nil { + panic(err) + } + return affected +} + +// Exec executes the query. +func (_u *BatchImageJobUpdate) Exec(ctx context.Context) error { + _, err := _u.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_u *BatchImageJobUpdate) ExecX(ctx context.Context) { + if err := _u.Exec(ctx); err != nil { + panic(err) + } +} + +// defaults sets the default values of the builder before save. +func (_u *BatchImageJobUpdate) defaults() { + if _, ok := _u.mutation.UpdatedAt(); !ok { + v := batchimagejob.UpdateDefaultUpdatedAt() + _u.mutation.SetUpdatedAt(v) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_u *BatchImageJobUpdate) check() error { + if v, ok := _u.mutation.Provider(); ok { + if err := batchimagejob.ProviderValidator(v); err != nil { + return &ValidationError{Name: "provider", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider": %w`, err)} + } + } + if v, ok := _u.mutation.Model(); ok { + if err := batchimagejob.ModelValidator(v); err != nil { + return &ValidationError{Name: "model", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.model": %w`, err)} + } + } + if v, ok := _u.mutation.TaskName(); ok { + if err := batchimagejob.TaskNameValidator(v); err != nil { + return &ValidationError{Name: "task_name", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.task_name": %w`, err)} + } + } + if v, ok := _u.mutation.Status(); ok { + if err := batchimagejob.StatusValidator(v); err != nil { + return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.status": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderJobName(); ok { + if err := batchimagejob.ProviderJobNameValidator(v); err != nil { + return &ValidationError{Name: "provider_job_name", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_job_name": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderInputRef(); ok { + if err := batchimagejob.ProviderInputRefValidator(v); err != nil { + return &ValidationError{Name: "provider_input_ref", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_input_ref": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderOutputRef(); ok { + if err := batchimagejob.ProviderOutputRefValidator(v); err != nil { + return &ValidationError{Name: "provider_output_ref", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_output_ref": %w`, err)} + } + } + if v, ok := _u.mutation.GcsInputURI(); ok { + if err := batchimagejob.GcsInputURIValidator(v); err != nil { + return &ValidationError{Name: "gcs_input_uri", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.gcs_input_uri": %w`, err)} + } + } + if v, ok := _u.mutation.GcsOutputURI(); ok { + if err := batchimagejob.GcsOutputURIValidator(v); err != nil { + return &ValidationError{Name: "gcs_output_uri", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.gcs_output_uri": %w`, err)} + } + } + if v, ok := _u.mutation.Currency(); ok { + if err := batchimagejob.CurrencyValidator(v); err != nil { + return &ValidationError{Name: "currency", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.currency": %w`, err)} + } + } + if v, ok := _u.mutation.HoldID(); ok { + if err := batchimagejob.HoldIDValidator(v); err != nil { + return &ValidationError{Name: "hold_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.hold_id": %w`, err)} + } + } + if v, ok := _u.mutation.IdempotencyKey(); ok { + if err := batchimagejob.IdempotencyKeyValidator(v); err != nil { + return &ValidationError{Name: "idempotency_key", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.idempotency_key": %w`, err)} + } + } + if v, ok := _u.mutation.RequestHash(); ok { + if err := batchimagejob.RequestHashValidator(v); err != nil { + return &ValidationError{Name: "request_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.request_hash": %w`, err)} + } + } + if v, ok := _u.mutation.ManifestHash(); ok { + if err := batchimagejob.ManifestHashValidator(v); err != nil { + return &ValidationError{Name: "manifest_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.manifest_hash": %w`, err)} + } + } + if v, ok := _u.mutation.LastErrorCode(); ok { + if err := batchimagejob.LastErrorCodeValidator(v); err != nil { + return &ValidationError{Name: "last_error_code", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.last_error_code": %w`, err)} + } + } + return nil +} + +func (_u *BatchImageJobUpdate) sqlSave(ctx context.Context) (_node int, err error) { + if err := _u.check(); err != nil { + return _node, err + } + _spec := sqlgraph.NewUpdateSpec(batchimagejob.Table, batchimagejob.Columns, sqlgraph.NewFieldSpec(batchimagejob.FieldID, field.TypeInt64)) + if ps := _u.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if value, ok := _u.mutation.UserID(); ok { + _spec.SetField(batchimagejob.FieldUserID, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedUserID(); ok { + _spec.AddField(batchimagejob.FieldUserID, field.TypeInt64, value) + } + if value, ok := _u.mutation.APIKeyID(); ok { + _spec.SetField(batchimagejob.FieldAPIKeyID, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedAPIKeyID(); ok { + _spec.AddField(batchimagejob.FieldAPIKeyID, field.TypeInt64, value) + } + if _u.mutation.APIKeyIDCleared() { + _spec.ClearField(batchimagejob.FieldAPIKeyID, field.TypeInt64) + } + if value, ok := _u.mutation.AccountID(); ok { + _spec.SetField(batchimagejob.FieldAccountID, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedAccountID(); ok { + _spec.AddField(batchimagejob.FieldAccountID, field.TypeInt64, value) + } + if _u.mutation.AccountIDCleared() { + _spec.ClearField(batchimagejob.FieldAccountID, field.TypeInt64) + } + if value, ok := _u.mutation.Provider(); ok { + _spec.SetField(batchimagejob.FieldProvider, field.TypeString, value) + } + if value, ok := _u.mutation.Model(); ok { + _spec.SetField(batchimagejob.FieldModel, field.TypeString, value) + } + if value, ok := _u.mutation.TaskName(); ok { + _spec.SetField(batchimagejob.FieldTaskName, field.TypeString, value) + } + if value, ok := _u.mutation.Status(); ok { + _spec.SetField(batchimagejob.FieldStatus, field.TypeString, value) + } + if value, ok := _u.mutation.ProviderJobName(); ok { + _spec.SetField(batchimagejob.FieldProviderJobName, field.TypeString, value) + } + if _u.mutation.ProviderJobNameCleared() { + _spec.ClearField(batchimagejob.FieldProviderJobName, field.TypeString) + } + if value, ok := _u.mutation.ProviderInputRef(); ok { + _spec.SetField(batchimagejob.FieldProviderInputRef, field.TypeString, value) + } + if _u.mutation.ProviderInputRefCleared() { + _spec.ClearField(batchimagejob.FieldProviderInputRef, field.TypeString) + } + if value, ok := _u.mutation.ProviderOutputRef(); ok { + _spec.SetField(batchimagejob.FieldProviderOutputRef, field.TypeString, value) + } + if _u.mutation.ProviderOutputRefCleared() { + _spec.ClearField(batchimagejob.FieldProviderOutputRef, field.TypeString) + } + if value, ok := _u.mutation.GcsInputURI(); ok { + _spec.SetField(batchimagejob.FieldGcsInputURI, field.TypeString, value) + } + if _u.mutation.GcsInputURICleared() { + _spec.ClearField(batchimagejob.FieldGcsInputURI, field.TypeString) + } + if value, ok := _u.mutation.GcsOutputURI(); ok { + _spec.SetField(batchimagejob.FieldGcsOutputURI, field.TypeString, value) + } + if _u.mutation.GcsOutputURICleared() { + _spec.ClearField(batchimagejob.FieldGcsOutputURI, field.TypeString) + } + if value, ok := _u.mutation.ItemCount(); ok { + _spec.SetField(batchimagejob.FieldItemCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedItemCount(); ok { + _spec.AddField(batchimagejob.FieldItemCount, field.TypeInt, value) + } + if value, ok := _u.mutation.SuccessCount(); ok { + _spec.SetField(batchimagejob.FieldSuccessCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedSuccessCount(); ok { + _spec.AddField(batchimagejob.FieldSuccessCount, field.TypeInt, value) + } + if value, ok := _u.mutation.FailCount(); ok { + _spec.SetField(batchimagejob.FieldFailCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedFailCount(); ok { + _spec.AddField(batchimagejob.FieldFailCount, field.TypeInt, value) + } + if value, ok := _u.mutation.CancelledCount(); ok { + _spec.SetField(batchimagejob.FieldCancelledCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedCancelledCount(); ok { + _spec.AddField(batchimagejob.FieldCancelledCount, field.TypeInt, value) + } + if value, ok := _u.mutation.EstimatedCost(); ok { + _spec.SetField(batchimagejob.FieldEstimatedCost, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedEstimatedCost(); ok { + _spec.AddField(batchimagejob.FieldEstimatedCost, field.TypeFloat64, value) + } + if value, ok := _u.mutation.HoldAmount(); ok { + _spec.SetField(batchimagejob.FieldHoldAmount, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedHoldAmount(); ok { + _spec.AddField(batchimagejob.FieldHoldAmount, field.TypeFloat64, value) + } + if _u.mutation.HoldAmountCleared() { + _spec.ClearField(batchimagejob.FieldHoldAmount, field.TypeFloat64) + } + if value, ok := _u.mutation.ActualCost(); ok { + _spec.SetField(batchimagejob.FieldActualCost, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedActualCost(); ok { + _spec.AddField(batchimagejob.FieldActualCost, field.TypeFloat64, value) + } + if _u.mutation.ActualCostCleared() { + _spec.ClearField(batchimagejob.FieldActualCost, field.TypeFloat64) + } + if value, ok := _u.mutation.Currency(); ok { + _spec.SetField(batchimagejob.FieldCurrency, field.TypeString, value) + } + if value, ok := _u.mutation.HoldID(); ok { + _spec.SetField(batchimagejob.FieldHoldID, field.TypeString, value) + } + if _u.mutation.HoldIDCleared() { + _spec.ClearField(batchimagejob.FieldHoldID, field.TypeString) + } + if value, ok := _u.mutation.IdempotencyKey(); ok { + _spec.SetField(batchimagejob.FieldIdempotencyKey, field.TypeString, value) + } + if _u.mutation.IdempotencyKeyCleared() { + _spec.ClearField(batchimagejob.FieldIdempotencyKey, field.TypeString) + } + if value, ok := _u.mutation.RequestHash(); ok { + _spec.SetField(batchimagejob.FieldRequestHash, field.TypeString, value) + } + if _u.mutation.RequestHashCleared() { + _spec.ClearField(batchimagejob.FieldRequestHash, field.TypeString) + } + if value, ok := _u.mutation.ManifestHash(); ok { + _spec.SetField(batchimagejob.FieldManifestHash, field.TypeString, value) + } + if _u.mutation.ManifestHashCleared() { + _spec.ClearField(batchimagejob.FieldManifestHash, field.TypeString) + } + if value, ok := _u.mutation.RetryCount(); ok { + _spec.SetField(batchimagejob.FieldRetryCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedRetryCount(); ok { + _spec.AddField(batchimagejob.FieldRetryCount, field.TypeInt, value) + } + if value, ok := _u.mutation.Version(); ok { + _spec.SetField(batchimagejob.FieldVersion, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedVersion(); ok { + _spec.AddField(batchimagejob.FieldVersion, field.TypeInt, value) + } + if value, ok := _u.mutation.OutputExpiresAt(); ok { + _spec.SetField(batchimagejob.FieldOutputExpiresAt, field.TypeTime, value) + } + if _u.mutation.OutputExpiresAtCleared() { + _spec.ClearField(batchimagejob.FieldOutputExpiresAt, field.TypeTime) + } + if value, ok := _u.mutation.InputDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldInputDeletedAt, field.TypeTime, value) + } + if _u.mutation.InputDeletedAtCleared() { + _spec.ClearField(batchimagejob.FieldInputDeletedAt, field.TypeTime) + } + if value, ok := _u.mutation.OutputDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldOutputDeletedAt, field.TypeTime, value) + } + if _u.mutation.OutputDeletedAtCleared() { + _spec.ClearField(batchimagejob.FieldOutputDeletedAt, field.TypeTime) + } + if value, ok := _u.mutation.DownloadedAt(); ok { + _spec.SetField(batchimagejob.FieldDownloadedAt, field.TypeTime, value) + } + if _u.mutation.DownloadedAtCleared() { + _spec.ClearField(batchimagejob.FieldDownloadedAt, field.TypeTime) + } + if value, ok := _u.mutation.UserDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldUserDeletedAt, field.TypeTime, value) + } + if _u.mutation.UserDeletedAtCleared() { + _spec.ClearField(batchimagejob.FieldUserDeletedAt, field.TypeTime) + } + if value, ok := _u.mutation.LastErrorCode(); ok { + _spec.SetField(batchimagejob.FieldLastErrorCode, field.TypeString, value) + } + if _u.mutation.LastErrorCodeCleared() { + _spec.ClearField(batchimagejob.FieldLastErrorCode, field.TypeString) + } + if value, ok := _u.mutation.LastErrorMessage(); ok { + _spec.SetField(batchimagejob.FieldLastErrorMessage, field.TypeString, value) + } + if _u.mutation.LastErrorMessageCleared() { + _spec.ClearField(batchimagejob.FieldLastErrorMessage, field.TypeString) + } + if value, ok := _u.mutation.UpdatedAt(); ok { + _spec.SetField(batchimagejob.FieldUpdatedAt, field.TypeTime, value) + } + if value, ok := _u.mutation.SubmittedAt(); ok { + _spec.SetField(batchimagejob.FieldSubmittedAt, field.TypeTime, value) + } + if _u.mutation.SubmittedAtCleared() { + _spec.ClearField(batchimagejob.FieldSubmittedAt, field.TypeTime) + } + if value, ok := _u.mutation.StartedAt(); ok { + _spec.SetField(batchimagejob.FieldStartedAt, field.TypeTime, value) + } + if _u.mutation.StartedAtCleared() { + _spec.ClearField(batchimagejob.FieldStartedAt, field.TypeTime) + } + if value, ok := _u.mutation.FinishedAt(); ok { + _spec.SetField(batchimagejob.FieldFinishedAt, field.TypeTime, value) + } + if _u.mutation.FinishedAtCleared() { + _spec.ClearField(batchimagejob.FieldFinishedAt, field.TypeTime) + } + if value, ok := _u.mutation.SettledAt(); ok { + _spec.SetField(batchimagejob.FieldSettledAt, field.TypeTime, value) + } + if _u.mutation.SettledAtCleared() { + _spec.ClearField(batchimagejob.FieldSettledAt, field.TypeTime) + } + if _node, err = sqlgraph.UpdateNodes(ctx, _u.driver, _spec); err != nil { + if _, ok := err.(*sqlgraph.NotFoundError); ok { + err = &NotFoundError{batchimagejob.Label} + } else if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return 0, err + } + _u.mutation.done = true + return _node, nil +} + +// BatchImageJobUpdateOne is the builder for updating a single BatchImageJob entity. +type BatchImageJobUpdateOne struct { + config + fields []string + hooks []Hook + mutation *BatchImageJobMutation +} + +// SetUserID sets the "user_id" field. +func (_u *BatchImageJobUpdateOne) SetUserID(v int64) *BatchImageJobUpdateOne { + _u.mutation.ResetUserID() + _u.mutation.SetUserID(v) + return _u +} + +// SetNillableUserID sets the "user_id" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableUserID(v *int64) *BatchImageJobUpdateOne { + if v != nil { + _u.SetUserID(*v) + } + return _u +} + +// AddUserID adds value to the "user_id" field. +func (_u *BatchImageJobUpdateOne) AddUserID(v int64) *BatchImageJobUpdateOne { + _u.mutation.AddUserID(v) + return _u +} + +// SetAPIKeyID sets the "api_key_id" field. +func (_u *BatchImageJobUpdateOne) SetAPIKeyID(v int64) *BatchImageJobUpdateOne { + _u.mutation.ResetAPIKeyID() + _u.mutation.SetAPIKeyID(v) + return _u +} + +// SetNillableAPIKeyID sets the "api_key_id" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableAPIKeyID(v *int64) *BatchImageJobUpdateOne { + if v != nil { + _u.SetAPIKeyID(*v) + } + return _u +} + +// AddAPIKeyID adds value to the "api_key_id" field. +func (_u *BatchImageJobUpdateOne) AddAPIKeyID(v int64) *BatchImageJobUpdateOne { + _u.mutation.AddAPIKeyID(v) + return _u +} + +// ClearAPIKeyID clears the value of the "api_key_id" field. +func (_u *BatchImageJobUpdateOne) ClearAPIKeyID() *BatchImageJobUpdateOne { + _u.mutation.ClearAPIKeyID() + return _u +} + +// SetAccountID sets the "account_id" field. +func (_u *BatchImageJobUpdateOne) SetAccountID(v int64) *BatchImageJobUpdateOne { + _u.mutation.ResetAccountID() + _u.mutation.SetAccountID(v) + return _u +} + +// SetNillableAccountID sets the "account_id" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableAccountID(v *int64) *BatchImageJobUpdateOne { + if v != nil { + _u.SetAccountID(*v) + } + return _u +} + +// AddAccountID adds value to the "account_id" field. +func (_u *BatchImageJobUpdateOne) AddAccountID(v int64) *BatchImageJobUpdateOne { + _u.mutation.AddAccountID(v) + return _u +} + +// ClearAccountID clears the value of the "account_id" field. +func (_u *BatchImageJobUpdateOne) ClearAccountID() *BatchImageJobUpdateOne { + _u.mutation.ClearAccountID() + return _u +} + +// SetProvider sets the "provider" field. +func (_u *BatchImageJobUpdateOne) SetProvider(v string) *BatchImageJobUpdateOne { + _u.mutation.SetProvider(v) + return _u +} + +// SetNillableProvider sets the "provider" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableProvider(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetProvider(*v) + } + return _u +} + +// SetModel sets the "model" field. +func (_u *BatchImageJobUpdateOne) SetModel(v string) *BatchImageJobUpdateOne { + _u.mutation.SetModel(v) + return _u +} + +// SetNillableModel sets the "model" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableModel(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetModel(*v) + } + return _u +} + +// SetTaskName sets the "task_name" field. +func (_u *BatchImageJobUpdateOne) SetTaskName(v string) *BatchImageJobUpdateOne { + _u.mutation.SetTaskName(v) + return _u +} + +// SetNillableTaskName sets the "task_name" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableTaskName(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetTaskName(*v) + } + return _u +} + +// SetStatus sets the "status" field. +func (_u *BatchImageJobUpdateOne) SetStatus(v string) *BatchImageJobUpdateOne { + _u.mutation.SetStatus(v) + return _u +} + +// SetNillableStatus sets the "status" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableStatus(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetStatus(*v) + } + return _u +} + +// SetProviderJobName sets the "provider_job_name" field. +func (_u *BatchImageJobUpdateOne) SetProviderJobName(v string) *BatchImageJobUpdateOne { + _u.mutation.SetProviderJobName(v) + return _u +} + +// SetNillableProviderJobName sets the "provider_job_name" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableProviderJobName(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetProviderJobName(*v) + } + return _u +} + +// ClearProviderJobName clears the value of the "provider_job_name" field. +func (_u *BatchImageJobUpdateOne) ClearProviderJobName() *BatchImageJobUpdateOne { + _u.mutation.ClearProviderJobName() + return _u +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (_u *BatchImageJobUpdateOne) SetProviderInputRef(v string) *BatchImageJobUpdateOne { + _u.mutation.SetProviderInputRef(v) + return _u +} + +// SetNillableProviderInputRef sets the "provider_input_ref" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableProviderInputRef(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetProviderInputRef(*v) + } + return _u +} + +// ClearProviderInputRef clears the value of the "provider_input_ref" field. +func (_u *BatchImageJobUpdateOne) ClearProviderInputRef() *BatchImageJobUpdateOne { + _u.mutation.ClearProviderInputRef() + return _u +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (_u *BatchImageJobUpdateOne) SetProviderOutputRef(v string) *BatchImageJobUpdateOne { + _u.mutation.SetProviderOutputRef(v) + return _u +} + +// SetNillableProviderOutputRef sets the "provider_output_ref" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableProviderOutputRef(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetProviderOutputRef(*v) + } + return _u +} + +// ClearProviderOutputRef clears the value of the "provider_output_ref" field. +func (_u *BatchImageJobUpdateOne) ClearProviderOutputRef() *BatchImageJobUpdateOne { + _u.mutation.ClearProviderOutputRef() + return _u +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (_u *BatchImageJobUpdateOne) SetGcsInputURI(v string) *BatchImageJobUpdateOne { + _u.mutation.SetGcsInputURI(v) + return _u +} + +// SetNillableGcsInputURI sets the "gcs_input_uri" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableGcsInputURI(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetGcsInputURI(*v) + } + return _u +} + +// ClearGcsInputURI clears the value of the "gcs_input_uri" field. +func (_u *BatchImageJobUpdateOne) ClearGcsInputURI() *BatchImageJobUpdateOne { + _u.mutation.ClearGcsInputURI() + return _u +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (_u *BatchImageJobUpdateOne) SetGcsOutputURI(v string) *BatchImageJobUpdateOne { + _u.mutation.SetGcsOutputURI(v) + return _u +} + +// SetNillableGcsOutputURI sets the "gcs_output_uri" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableGcsOutputURI(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetGcsOutputURI(*v) + } + return _u +} + +// ClearGcsOutputURI clears the value of the "gcs_output_uri" field. +func (_u *BatchImageJobUpdateOne) ClearGcsOutputURI() *BatchImageJobUpdateOne { + _u.mutation.ClearGcsOutputURI() + return _u +} + +// SetItemCount sets the "item_count" field. +func (_u *BatchImageJobUpdateOne) SetItemCount(v int) *BatchImageJobUpdateOne { + _u.mutation.ResetItemCount() + _u.mutation.SetItemCount(v) + return _u +} + +// SetNillableItemCount sets the "item_count" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableItemCount(v *int) *BatchImageJobUpdateOne { + if v != nil { + _u.SetItemCount(*v) + } + return _u +} + +// AddItemCount adds value to the "item_count" field. +func (_u *BatchImageJobUpdateOne) AddItemCount(v int) *BatchImageJobUpdateOne { + _u.mutation.AddItemCount(v) + return _u +} + +// SetSuccessCount sets the "success_count" field. +func (_u *BatchImageJobUpdateOne) SetSuccessCount(v int) *BatchImageJobUpdateOne { + _u.mutation.ResetSuccessCount() + _u.mutation.SetSuccessCount(v) + return _u +} + +// SetNillableSuccessCount sets the "success_count" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableSuccessCount(v *int) *BatchImageJobUpdateOne { + if v != nil { + _u.SetSuccessCount(*v) + } + return _u +} + +// AddSuccessCount adds value to the "success_count" field. +func (_u *BatchImageJobUpdateOne) AddSuccessCount(v int) *BatchImageJobUpdateOne { + _u.mutation.AddSuccessCount(v) + return _u +} + +// SetFailCount sets the "fail_count" field. +func (_u *BatchImageJobUpdateOne) SetFailCount(v int) *BatchImageJobUpdateOne { + _u.mutation.ResetFailCount() + _u.mutation.SetFailCount(v) + return _u +} + +// SetNillableFailCount sets the "fail_count" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableFailCount(v *int) *BatchImageJobUpdateOne { + if v != nil { + _u.SetFailCount(*v) + } + return _u +} + +// AddFailCount adds value to the "fail_count" field. +func (_u *BatchImageJobUpdateOne) AddFailCount(v int) *BatchImageJobUpdateOne { + _u.mutation.AddFailCount(v) + return _u +} + +// SetCancelledCount sets the "cancelled_count" field. +func (_u *BatchImageJobUpdateOne) SetCancelledCount(v int) *BatchImageJobUpdateOne { + _u.mutation.ResetCancelledCount() + _u.mutation.SetCancelledCount(v) + return _u +} + +// SetNillableCancelledCount sets the "cancelled_count" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableCancelledCount(v *int) *BatchImageJobUpdateOne { + if v != nil { + _u.SetCancelledCount(*v) + } + return _u +} + +// AddCancelledCount adds value to the "cancelled_count" field. +func (_u *BatchImageJobUpdateOne) AddCancelledCount(v int) *BatchImageJobUpdateOne { + _u.mutation.AddCancelledCount(v) + return _u +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (_u *BatchImageJobUpdateOne) SetEstimatedCost(v float64) *BatchImageJobUpdateOne { + _u.mutation.ResetEstimatedCost() + _u.mutation.SetEstimatedCost(v) + return _u +} + +// SetNillableEstimatedCost sets the "estimated_cost" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableEstimatedCost(v *float64) *BatchImageJobUpdateOne { + if v != nil { + _u.SetEstimatedCost(*v) + } + return _u +} + +// AddEstimatedCost adds value to the "estimated_cost" field. +func (_u *BatchImageJobUpdateOne) AddEstimatedCost(v float64) *BatchImageJobUpdateOne { + _u.mutation.AddEstimatedCost(v) + return _u +} + +// SetHoldAmount sets the "hold_amount" field. +func (_u *BatchImageJobUpdateOne) SetHoldAmount(v float64) *BatchImageJobUpdateOne { + _u.mutation.ResetHoldAmount() + _u.mutation.SetHoldAmount(v) + return _u +} + +// SetNillableHoldAmount sets the "hold_amount" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableHoldAmount(v *float64) *BatchImageJobUpdateOne { + if v != nil { + _u.SetHoldAmount(*v) + } + return _u +} + +// AddHoldAmount adds value to the "hold_amount" field. +func (_u *BatchImageJobUpdateOne) AddHoldAmount(v float64) *BatchImageJobUpdateOne { + _u.mutation.AddHoldAmount(v) + return _u +} + +// ClearHoldAmount clears the value of the "hold_amount" field. +func (_u *BatchImageJobUpdateOne) ClearHoldAmount() *BatchImageJobUpdateOne { + _u.mutation.ClearHoldAmount() + return _u +} + +// SetActualCost sets the "actual_cost" field. +func (_u *BatchImageJobUpdateOne) SetActualCost(v float64) *BatchImageJobUpdateOne { + _u.mutation.ResetActualCost() + _u.mutation.SetActualCost(v) + return _u +} + +// SetNillableActualCost sets the "actual_cost" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableActualCost(v *float64) *BatchImageJobUpdateOne { + if v != nil { + _u.SetActualCost(*v) + } + return _u +} + +// AddActualCost adds value to the "actual_cost" field. +func (_u *BatchImageJobUpdateOne) AddActualCost(v float64) *BatchImageJobUpdateOne { + _u.mutation.AddActualCost(v) + return _u +} + +// ClearActualCost clears the value of the "actual_cost" field. +func (_u *BatchImageJobUpdateOne) ClearActualCost() *BatchImageJobUpdateOne { + _u.mutation.ClearActualCost() + return _u +} + +// SetCurrency sets the "currency" field. +func (_u *BatchImageJobUpdateOne) SetCurrency(v string) *BatchImageJobUpdateOne { + _u.mutation.SetCurrency(v) + return _u +} + +// SetNillableCurrency sets the "currency" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableCurrency(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetCurrency(*v) + } + return _u +} + +// SetHoldID sets the "hold_id" field. +func (_u *BatchImageJobUpdateOne) SetHoldID(v string) *BatchImageJobUpdateOne { + _u.mutation.SetHoldID(v) + return _u +} + +// SetNillableHoldID sets the "hold_id" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableHoldID(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetHoldID(*v) + } + return _u +} + +// ClearHoldID clears the value of the "hold_id" field. +func (_u *BatchImageJobUpdateOne) ClearHoldID() *BatchImageJobUpdateOne { + _u.mutation.ClearHoldID() + return _u +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (_u *BatchImageJobUpdateOne) SetIdempotencyKey(v string) *BatchImageJobUpdateOne { + _u.mutation.SetIdempotencyKey(v) + return _u +} + +// SetNillableIdempotencyKey sets the "idempotency_key" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableIdempotencyKey(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetIdempotencyKey(*v) + } + return _u +} + +// ClearIdempotencyKey clears the value of the "idempotency_key" field. +func (_u *BatchImageJobUpdateOne) ClearIdempotencyKey() *BatchImageJobUpdateOne { + _u.mutation.ClearIdempotencyKey() + return _u +} + +// SetRequestHash sets the "request_hash" field. +func (_u *BatchImageJobUpdateOne) SetRequestHash(v string) *BatchImageJobUpdateOne { + _u.mutation.SetRequestHash(v) + return _u +} + +// SetNillableRequestHash sets the "request_hash" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableRequestHash(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetRequestHash(*v) + } + return _u +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (_u *BatchImageJobUpdateOne) ClearRequestHash() *BatchImageJobUpdateOne { + _u.mutation.ClearRequestHash() + return _u +} + +// SetManifestHash sets the "manifest_hash" field. +func (_u *BatchImageJobUpdateOne) SetManifestHash(v string) *BatchImageJobUpdateOne { + _u.mutation.SetManifestHash(v) + return _u +} + +// SetNillableManifestHash sets the "manifest_hash" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableManifestHash(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetManifestHash(*v) + } + return _u +} + +// ClearManifestHash clears the value of the "manifest_hash" field. +func (_u *BatchImageJobUpdateOne) ClearManifestHash() *BatchImageJobUpdateOne { + _u.mutation.ClearManifestHash() + return _u +} + +// SetRetryCount sets the "retry_count" field. +func (_u *BatchImageJobUpdateOne) SetRetryCount(v int) *BatchImageJobUpdateOne { + _u.mutation.ResetRetryCount() + _u.mutation.SetRetryCount(v) + return _u +} + +// SetNillableRetryCount sets the "retry_count" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableRetryCount(v *int) *BatchImageJobUpdateOne { + if v != nil { + _u.SetRetryCount(*v) + } + return _u +} + +// AddRetryCount adds value to the "retry_count" field. +func (_u *BatchImageJobUpdateOne) AddRetryCount(v int) *BatchImageJobUpdateOne { + _u.mutation.AddRetryCount(v) + return _u +} + +// SetVersion sets the "version" field. +func (_u *BatchImageJobUpdateOne) SetVersion(v int) *BatchImageJobUpdateOne { + _u.mutation.ResetVersion() + _u.mutation.SetVersion(v) + return _u +} + +// SetNillableVersion sets the "version" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableVersion(v *int) *BatchImageJobUpdateOne { + if v != nil { + _u.SetVersion(*v) + } + return _u +} + +// AddVersion adds value to the "version" field. +func (_u *BatchImageJobUpdateOne) AddVersion(v int) *BatchImageJobUpdateOne { + _u.mutation.AddVersion(v) + return _u +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (_u *BatchImageJobUpdateOne) SetOutputExpiresAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetOutputExpiresAt(v) + return _u +} + +// SetNillableOutputExpiresAt sets the "output_expires_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableOutputExpiresAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetOutputExpiresAt(*v) + } + return _u +} + +// ClearOutputExpiresAt clears the value of the "output_expires_at" field. +func (_u *BatchImageJobUpdateOne) ClearOutputExpiresAt() *BatchImageJobUpdateOne { + _u.mutation.ClearOutputExpiresAt() + return _u +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (_u *BatchImageJobUpdateOne) SetInputDeletedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetInputDeletedAt(v) + return _u +} + +// SetNillableInputDeletedAt sets the "input_deleted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableInputDeletedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetInputDeletedAt(*v) + } + return _u +} + +// ClearInputDeletedAt clears the value of the "input_deleted_at" field. +func (_u *BatchImageJobUpdateOne) ClearInputDeletedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearInputDeletedAt() + return _u +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (_u *BatchImageJobUpdateOne) SetOutputDeletedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetOutputDeletedAt(v) + return _u +} + +// SetNillableOutputDeletedAt sets the "output_deleted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableOutputDeletedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetOutputDeletedAt(*v) + } + return _u +} + +// ClearOutputDeletedAt clears the value of the "output_deleted_at" field. +func (_u *BatchImageJobUpdateOne) ClearOutputDeletedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearOutputDeletedAt() + return _u +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (_u *BatchImageJobUpdateOne) SetDownloadedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetDownloadedAt(v) + return _u +} + +// SetNillableDownloadedAt sets the "downloaded_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableDownloadedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetDownloadedAt(*v) + } + return _u +} + +// ClearDownloadedAt clears the value of the "downloaded_at" field. +func (_u *BatchImageJobUpdateOne) ClearDownloadedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearDownloadedAt() + return _u +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (_u *BatchImageJobUpdateOne) SetUserDeletedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetUserDeletedAt(v) + return _u +} + +// SetNillableUserDeletedAt sets the "user_deleted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableUserDeletedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetUserDeletedAt(*v) + } + return _u +} + +// ClearUserDeletedAt clears the value of the "user_deleted_at" field. +func (_u *BatchImageJobUpdateOne) ClearUserDeletedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearUserDeletedAt() + return _u +} + +// SetLastErrorCode sets the "last_error_code" field. +func (_u *BatchImageJobUpdateOne) SetLastErrorCode(v string) *BatchImageJobUpdateOne { + _u.mutation.SetLastErrorCode(v) + return _u +} + +// SetNillableLastErrorCode sets the "last_error_code" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableLastErrorCode(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetLastErrorCode(*v) + } + return _u +} + +// ClearLastErrorCode clears the value of the "last_error_code" field. +func (_u *BatchImageJobUpdateOne) ClearLastErrorCode() *BatchImageJobUpdateOne { + _u.mutation.ClearLastErrorCode() + return _u +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (_u *BatchImageJobUpdateOne) SetLastErrorMessage(v string) *BatchImageJobUpdateOne { + _u.mutation.SetLastErrorMessage(v) + return _u +} + +// SetNillableLastErrorMessage sets the "last_error_message" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableLastErrorMessage(v *string) *BatchImageJobUpdateOne { + if v != nil { + _u.SetLastErrorMessage(*v) + } + return _u +} + +// ClearLastErrorMessage clears the value of the "last_error_message" field. +func (_u *BatchImageJobUpdateOne) ClearLastErrorMessage() *BatchImageJobUpdateOne { + _u.mutation.ClearLastErrorMessage() + return _u +} + +// SetUpdatedAt sets the "updated_at" field. +func (_u *BatchImageJobUpdateOne) SetUpdatedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetUpdatedAt(v) + return _u +} + +// SetSubmittedAt sets the "submitted_at" field. +func (_u *BatchImageJobUpdateOne) SetSubmittedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetSubmittedAt(v) + return _u +} + +// SetNillableSubmittedAt sets the "submitted_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableSubmittedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetSubmittedAt(*v) + } + return _u +} + +// ClearSubmittedAt clears the value of the "submitted_at" field. +func (_u *BatchImageJobUpdateOne) ClearSubmittedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearSubmittedAt() + return _u +} + +// SetStartedAt sets the "started_at" field. +func (_u *BatchImageJobUpdateOne) SetStartedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetStartedAt(v) + return _u +} + +// SetNillableStartedAt sets the "started_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableStartedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetStartedAt(*v) + } + return _u +} + +// ClearStartedAt clears the value of the "started_at" field. +func (_u *BatchImageJobUpdateOne) ClearStartedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearStartedAt() + return _u +} + +// SetFinishedAt sets the "finished_at" field. +func (_u *BatchImageJobUpdateOne) SetFinishedAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetFinishedAt(v) + return _u +} + +// SetNillableFinishedAt sets the "finished_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableFinishedAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetFinishedAt(*v) + } + return _u +} + +// ClearFinishedAt clears the value of the "finished_at" field. +func (_u *BatchImageJobUpdateOne) ClearFinishedAt() *BatchImageJobUpdateOne { + _u.mutation.ClearFinishedAt() + return _u +} + +// SetSettledAt sets the "settled_at" field. +func (_u *BatchImageJobUpdateOne) SetSettledAt(v time.Time) *BatchImageJobUpdateOne { + _u.mutation.SetSettledAt(v) + return _u +} + +// SetNillableSettledAt sets the "settled_at" field if the given value is not nil. +func (_u *BatchImageJobUpdateOne) SetNillableSettledAt(v *time.Time) *BatchImageJobUpdateOne { + if v != nil { + _u.SetSettledAt(*v) + } + return _u +} + +// ClearSettledAt clears the value of the "settled_at" field. +func (_u *BatchImageJobUpdateOne) ClearSettledAt() *BatchImageJobUpdateOne { + _u.mutation.ClearSettledAt() + return _u +} + +// Mutation returns the BatchImageJobMutation object of the builder. +func (_u *BatchImageJobUpdateOne) Mutation() *BatchImageJobMutation { + return _u.mutation +} + +// Where appends a list predicates to the BatchImageJobUpdate builder. +func (_u *BatchImageJobUpdateOne) Where(ps ...predicate.BatchImageJob) *BatchImageJobUpdateOne { + _u.mutation.Where(ps...) + return _u +} + +// Select allows selecting one or more fields (columns) of the returned entity. +// The default is selecting all fields defined in the entity schema. +func (_u *BatchImageJobUpdateOne) Select(field string, fields ...string) *BatchImageJobUpdateOne { + _u.fields = append([]string{field}, fields...) + return _u +} + +// Save executes the query and returns the updated BatchImageJob entity. +func (_u *BatchImageJobUpdateOne) Save(ctx context.Context) (*BatchImageJob, error) { + _u.defaults() + return withHooks(ctx, _u.sqlSave, _u.mutation, _u.hooks) +} + +// SaveX is like Save, but panics if an error occurs. +func (_u *BatchImageJobUpdateOne) SaveX(ctx context.Context) *BatchImageJob { + node, err := _u.Save(ctx) + if err != nil { + panic(err) + } + return node +} + +// Exec executes the query on the entity. +func (_u *BatchImageJobUpdateOne) Exec(ctx context.Context) error { + _, err := _u.Save(ctx) + return err +} + +// ExecX is like Exec, but panics if an error occurs. +func (_u *BatchImageJobUpdateOne) ExecX(ctx context.Context) { + if err := _u.Exec(ctx); err != nil { + panic(err) + } +} + +// defaults sets the default values of the builder before save. +func (_u *BatchImageJobUpdateOne) defaults() { + if _, ok := _u.mutation.UpdatedAt(); !ok { + v := batchimagejob.UpdateDefaultUpdatedAt() + _u.mutation.SetUpdatedAt(v) + } +} + +// check runs all checks and user-defined validators on the builder. +func (_u *BatchImageJobUpdateOne) check() error { + if v, ok := _u.mutation.Provider(); ok { + if err := batchimagejob.ProviderValidator(v); err != nil { + return &ValidationError{Name: "provider", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider": %w`, err)} + } + } + if v, ok := _u.mutation.Model(); ok { + if err := batchimagejob.ModelValidator(v); err != nil { + return &ValidationError{Name: "model", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.model": %w`, err)} + } + } + if v, ok := _u.mutation.TaskName(); ok { + if err := batchimagejob.TaskNameValidator(v); err != nil { + return &ValidationError{Name: "task_name", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.task_name": %w`, err)} + } + } + if v, ok := _u.mutation.Status(); ok { + if err := batchimagejob.StatusValidator(v); err != nil { + return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.status": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderJobName(); ok { + if err := batchimagejob.ProviderJobNameValidator(v); err != nil { + return &ValidationError{Name: "provider_job_name", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_job_name": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderInputRef(); ok { + if err := batchimagejob.ProviderInputRefValidator(v); err != nil { + return &ValidationError{Name: "provider_input_ref", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_input_ref": %w`, err)} + } + } + if v, ok := _u.mutation.ProviderOutputRef(); ok { + if err := batchimagejob.ProviderOutputRefValidator(v); err != nil { + return &ValidationError{Name: "provider_output_ref", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.provider_output_ref": %w`, err)} + } + } + if v, ok := _u.mutation.GcsInputURI(); ok { + if err := batchimagejob.GcsInputURIValidator(v); err != nil { + return &ValidationError{Name: "gcs_input_uri", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.gcs_input_uri": %w`, err)} + } + } + if v, ok := _u.mutation.GcsOutputURI(); ok { + if err := batchimagejob.GcsOutputURIValidator(v); err != nil { + return &ValidationError{Name: "gcs_output_uri", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.gcs_output_uri": %w`, err)} + } + } + if v, ok := _u.mutation.Currency(); ok { + if err := batchimagejob.CurrencyValidator(v); err != nil { + return &ValidationError{Name: "currency", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.currency": %w`, err)} + } + } + if v, ok := _u.mutation.HoldID(); ok { + if err := batchimagejob.HoldIDValidator(v); err != nil { + return &ValidationError{Name: "hold_id", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.hold_id": %w`, err)} + } + } + if v, ok := _u.mutation.IdempotencyKey(); ok { + if err := batchimagejob.IdempotencyKeyValidator(v); err != nil { + return &ValidationError{Name: "idempotency_key", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.idempotency_key": %w`, err)} + } + } + if v, ok := _u.mutation.RequestHash(); ok { + if err := batchimagejob.RequestHashValidator(v); err != nil { + return &ValidationError{Name: "request_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.request_hash": %w`, err)} + } + } + if v, ok := _u.mutation.ManifestHash(); ok { + if err := batchimagejob.ManifestHashValidator(v); err != nil { + return &ValidationError{Name: "manifest_hash", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.manifest_hash": %w`, err)} + } + } + if v, ok := _u.mutation.LastErrorCode(); ok { + if err := batchimagejob.LastErrorCodeValidator(v); err != nil { + return &ValidationError{Name: "last_error_code", err: fmt.Errorf(`ent: validator failed for field "BatchImageJob.last_error_code": %w`, err)} + } + } + return nil +} + +func (_u *BatchImageJobUpdateOne) sqlSave(ctx context.Context) (_node *BatchImageJob, err error) { + if err := _u.check(); err != nil { + return _node, err + } + _spec := sqlgraph.NewUpdateSpec(batchimagejob.Table, batchimagejob.Columns, sqlgraph.NewFieldSpec(batchimagejob.FieldID, field.TypeInt64)) + id, ok := _u.mutation.ID() + if !ok { + return nil, &ValidationError{Name: "id", err: errors.New(`ent: missing "BatchImageJob.id" for update`)} + } + _spec.Node.ID.Value = id + if fields := _u.fields; len(fields) > 0 { + _spec.Node.Columns = make([]string, 0, len(fields)) + _spec.Node.Columns = append(_spec.Node.Columns, batchimagejob.FieldID) + for _, f := range fields { + if !batchimagejob.ValidColumn(f) { + return nil, &ValidationError{Name: f, err: fmt.Errorf("ent: invalid field %q for query", f)} + } + if f != batchimagejob.FieldID { + _spec.Node.Columns = append(_spec.Node.Columns, f) + } + } + } + if ps := _u.mutation.predicates; len(ps) > 0 { + _spec.Predicate = func(selector *sql.Selector) { + for i := range ps { + ps[i](selector) + } + } + } + if value, ok := _u.mutation.UserID(); ok { + _spec.SetField(batchimagejob.FieldUserID, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedUserID(); ok { + _spec.AddField(batchimagejob.FieldUserID, field.TypeInt64, value) + } + if value, ok := _u.mutation.APIKeyID(); ok { + _spec.SetField(batchimagejob.FieldAPIKeyID, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedAPIKeyID(); ok { + _spec.AddField(batchimagejob.FieldAPIKeyID, field.TypeInt64, value) + } + if _u.mutation.APIKeyIDCleared() { + _spec.ClearField(batchimagejob.FieldAPIKeyID, field.TypeInt64) + } + if value, ok := _u.mutation.AccountID(); ok { + _spec.SetField(batchimagejob.FieldAccountID, field.TypeInt64, value) + } + if value, ok := _u.mutation.AddedAccountID(); ok { + _spec.AddField(batchimagejob.FieldAccountID, field.TypeInt64, value) + } + if _u.mutation.AccountIDCleared() { + _spec.ClearField(batchimagejob.FieldAccountID, field.TypeInt64) + } + if value, ok := _u.mutation.Provider(); ok { + _spec.SetField(batchimagejob.FieldProvider, field.TypeString, value) + } + if value, ok := _u.mutation.Model(); ok { + _spec.SetField(batchimagejob.FieldModel, field.TypeString, value) + } + if value, ok := _u.mutation.TaskName(); ok { + _spec.SetField(batchimagejob.FieldTaskName, field.TypeString, value) + } + if value, ok := _u.mutation.Status(); ok { + _spec.SetField(batchimagejob.FieldStatus, field.TypeString, value) + } + if value, ok := _u.mutation.ProviderJobName(); ok { + _spec.SetField(batchimagejob.FieldProviderJobName, field.TypeString, value) + } + if _u.mutation.ProviderJobNameCleared() { + _spec.ClearField(batchimagejob.FieldProviderJobName, field.TypeString) + } + if value, ok := _u.mutation.ProviderInputRef(); ok { + _spec.SetField(batchimagejob.FieldProviderInputRef, field.TypeString, value) + } + if _u.mutation.ProviderInputRefCleared() { + _spec.ClearField(batchimagejob.FieldProviderInputRef, field.TypeString) + } + if value, ok := _u.mutation.ProviderOutputRef(); ok { + _spec.SetField(batchimagejob.FieldProviderOutputRef, field.TypeString, value) + } + if _u.mutation.ProviderOutputRefCleared() { + _spec.ClearField(batchimagejob.FieldProviderOutputRef, field.TypeString) + } + if value, ok := _u.mutation.GcsInputURI(); ok { + _spec.SetField(batchimagejob.FieldGcsInputURI, field.TypeString, value) + } + if _u.mutation.GcsInputURICleared() { + _spec.ClearField(batchimagejob.FieldGcsInputURI, field.TypeString) + } + if value, ok := _u.mutation.GcsOutputURI(); ok { + _spec.SetField(batchimagejob.FieldGcsOutputURI, field.TypeString, value) + } + if _u.mutation.GcsOutputURICleared() { + _spec.ClearField(batchimagejob.FieldGcsOutputURI, field.TypeString) + } + if value, ok := _u.mutation.ItemCount(); ok { + _spec.SetField(batchimagejob.FieldItemCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedItemCount(); ok { + _spec.AddField(batchimagejob.FieldItemCount, field.TypeInt, value) + } + if value, ok := _u.mutation.SuccessCount(); ok { + _spec.SetField(batchimagejob.FieldSuccessCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedSuccessCount(); ok { + _spec.AddField(batchimagejob.FieldSuccessCount, field.TypeInt, value) + } + if value, ok := _u.mutation.FailCount(); ok { + _spec.SetField(batchimagejob.FieldFailCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedFailCount(); ok { + _spec.AddField(batchimagejob.FieldFailCount, field.TypeInt, value) + } + if value, ok := _u.mutation.CancelledCount(); ok { + _spec.SetField(batchimagejob.FieldCancelledCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedCancelledCount(); ok { + _spec.AddField(batchimagejob.FieldCancelledCount, field.TypeInt, value) + } + if value, ok := _u.mutation.EstimatedCost(); ok { + _spec.SetField(batchimagejob.FieldEstimatedCost, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedEstimatedCost(); ok { + _spec.AddField(batchimagejob.FieldEstimatedCost, field.TypeFloat64, value) + } + if value, ok := _u.mutation.HoldAmount(); ok { + _spec.SetField(batchimagejob.FieldHoldAmount, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedHoldAmount(); ok { + _spec.AddField(batchimagejob.FieldHoldAmount, field.TypeFloat64, value) + } + if _u.mutation.HoldAmountCleared() { + _spec.ClearField(batchimagejob.FieldHoldAmount, field.TypeFloat64) + } + if value, ok := _u.mutation.ActualCost(); ok { + _spec.SetField(batchimagejob.FieldActualCost, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedActualCost(); ok { + _spec.AddField(batchimagejob.FieldActualCost, field.TypeFloat64, value) + } + if _u.mutation.ActualCostCleared() { + _spec.ClearField(batchimagejob.FieldActualCost, field.TypeFloat64) + } + if value, ok := _u.mutation.Currency(); ok { + _spec.SetField(batchimagejob.FieldCurrency, field.TypeString, value) + } + if value, ok := _u.mutation.HoldID(); ok { + _spec.SetField(batchimagejob.FieldHoldID, field.TypeString, value) + } + if _u.mutation.HoldIDCleared() { + _spec.ClearField(batchimagejob.FieldHoldID, field.TypeString) + } + if value, ok := _u.mutation.IdempotencyKey(); ok { + _spec.SetField(batchimagejob.FieldIdempotencyKey, field.TypeString, value) + } + if _u.mutation.IdempotencyKeyCleared() { + _spec.ClearField(batchimagejob.FieldIdempotencyKey, field.TypeString) + } + if value, ok := _u.mutation.RequestHash(); ok { + _spec.SetField(batchimagejob.FieldRequestHash, field.TypeString, value) + } + if _u.mutation.RequestHashCleared() { + _spec.ClearField(batchimagejob.FieldRequestHash, field.TypeString) + } + if value, ok := _u.mutation.ManifestHash(); ok { + _spec.SetField(batchimagejob.FieldManifestHash, field.TypeString, value) + } + if _u.mutation.ManifestHashCleared() { + _spec.ClearField(batchimagejob.FieldManifestHash, field.TypeString) + } + if value, ok := _u.mutation.RetryCount(); ok { + _spec.SetField(batchimagejob.FieldRetryCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedRetryCount(); ok { + _spec.AddField(batchimagejob.FieldRetryCount, field.TypeInt, value) + } + if value, ok := _u.mutation.Version(); ok { + _spec.SetField(batchimagejob.FieldVersion, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedVersion(); ok { + _spec.AddField(batchimagejob.FieldVersion, field.TypeInt, value) + } + if value, ok := _u.mutation.OutputExpiresAt(); ok { + _spec.SetField(batchimagejob.FieldOutputExpiresAt, field.TypeTime, value) + } + if _u.mutation.OutputExpiresAtCleared() { + _spec.ClearField(batchimagejob.FieldOutputExpiresAt, field.TypeTime) + } + if value, ok := _u.mutation.InputDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldInputDeletedAt, field.TypeTime, value) + } + if _u.mutation.InputDeletedAtCleared() { + _spec.ClearField(batchimagejob.FieldInputDeletedAt, field.TypeTime) + } + if value, ok := _u.mutation.OutputDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldOutputDeletedAt, field.TypeTime, value) + } + if _u.mutation.OutputDeletedAtCleared() { + _spec.ClearField(batchimagejob.FieldOutputDeletedAt, field.TypeTime) + } + if value, ok := _u.mutation.DownloadedAt(); ok { + _spec.SetField(batchimagejob.FieldDownloadedAt, field.TypeTime, value) + } + if _u.mutation.DownloadedAtCleared() { + _spec.ClearField(batchimagejob.FieldDownloadedAt, field.TypeTime) + } + if value, ok := _u.mutation.UserDeletedAt(); ok { + _spec.SetField(batchimagejob.FieldUserDeletedAt, field.TypeTime, value) + } + if _u.mutation.UserDeletedAtCleared() { + _spec.ClearField(batchimagejob.FieldUserDeletedAt, field.TypeTime) + } + if value, ok := _u.mutation.LastErrorCode(); ok { + _spec.SetField(batchimagejob.FieldLastErrorCode, field.TypeString, value) + } + if _u.mutation.LastErrorCodeCleared() { + _spec.ClearField(batchimagejob.FieldLastErrorCode, field.TypeString) + } + if value, ok := _u.mutation.LastErrorMessage(); ok { + _spec.SetField(batchimagejob.FieldLastErrorMessage, field.TypeString, value) + } + if _u.mutation.LastErrorMessageCleared() { + _spec.ClearField(batchimagejob.FieldLastErrorMessage, field.TypeString) + } + if value, ok := _u.mutation.UpdatedAt(); ok { + _spec.SetField(batchimagejob.FieldUpdatedAt, field.TypeTime, value) + } + if value, ok := _u.mutation.SubmittedAt(); ok { + _spec.SetField(batchimagejob.FieldSubmittedAt, field.TypeTime, value) + } + if _u.mutation.SubmittedAtCleared() { + _spec.ClearField(batchimagejob.FieldSubmittedAt, field.TypeTime) + } + if value, ok := _u.mutation.StartedAt(); ok { + _spec.SetField(batchimagejob.FieldStartedAt, field.TypeTime, value) + } + if _u.mutation.StartedAtCleared() { + _spec.ClearField(batchimagejob.FieldStartedAt, field.TypeTime) + } + if value, ok := _u.mutation.FinishedAt(); ok { + _spec.SetField(batchimagejob.FieldFinishedAt, field.TypeTime, value) + } + if _u.mutation.FinishedAtCleared() { + _spec.ClearField(batchimagejob.FieldFinishedAt, field.TypeTime) + } + if value, ok := _u.mutation.SettledAt(); ok { + _spec.SetField(batchimagejob.FieldSettledAt, field.TypeTime, value) + } + if _u.mutation.SettledAtCleared() { + _spec.ClearField(batchimagejob.FieldSettledAt, field.TypeTime) + } + _node = &BatchImageJob{config: _u.config} + _spec.Assign = _node.assignValues + _spec.ScanValues = _node.scanValues + if err = sqlgraph.UpdateNode(ctx, _u.driver, _spec); err != nil { + if _, ok := err.(*sqlgraph.NotFoundError); ok { + err = &NotFoundError{batchimagejob.Label} + } else if sqlgraph.IsConstraintError(err) { + err = &ConstraintError{msg: err.Error(), wrap: err} + } + return nil, err + } + _u.mutation.done = true + return _node, nil +} diff --git a/backend/ent/client.go b/backend/ent/client.go index 0b4edcf203..e9b74fcfef 100644 --- a/backend/ent/client.go +++ b/backend/ent/client.go @@ -22,6 +22,9 @@ import ( "github.com/Wei-Shaw/sub2api/ent/apikey" "github.com/Wei-Shaw/sub2api/ent/authidentity" "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" "github.com/Wei-Shaw/sub2api/ent/channelmonitor" "github.com/Wei-Shaw/sub2api/ent/channelmonitordailyrollup" "github.com/Wei-Shaw/sub2api/ent/channelmonitorhistory" @@ -73,6 +76,12 @@ type Client struct { AuthIdentity *AuthIdentityClient // AuthIdentityChannel is the client for interacting with the AuthIdentityChannel builders. AuthIdentityChannel *AuthIdentityChannelClient + // BatchImageEvent is the client for interacting with the BatchImageEvent builders. + BatchImageEvent *BatchImageEventClient + // BatchImageItem is the client for interacting with the BatchImageItem builders. + BatchImageItem *BatchImageItemClient + // BatchImageJob is the client for interacting with the BatchImageJob builders. + BatchImageJob *BatchImageJobClient // ChannelMonitor is the client for interacting with the ChannelMonitor builders. ChannelMonitor *ChannelMonitorClient // ChannelMonitorDailyRollup is the client for interacting with the ChannelMonitorDailyRollup builders. @@ -147,6 +156,9 @@ func (c *Client) init() { c.AnnouncementRead = NewAnnouncementReadClient(c.config) c.AuthIdentity = NewAuthIdentityClient(c.config) c.AuthIdentityChannel = NewAuthIdentityChannelClient(c.config) + c.BatchImageEvent = NewBatchImageEventClient(c.config) + c.BatchImageItem = NewBatchImageItemClient(c.config) + c.BatchImageJob = NewBatchImageJobClient(c.config) c.ChannelMonitor = NewChannelMonitorClient(c.config) c.ChannelMonitorDailyRollup = NewChannelMonitorDailyRollupClient(c.config) c.ChannelMonitorHistory = NewChannelMonitorHistoryClient(c.config) @@ -274,6 +286,9 @@ func (c *Client) Tx(ctx context.Context) (*Tx, error) { AnnouncementRead: NewAnnouncementReadClient(cfg), AuthIdentity: NewAuthIdentityClient(cfg), AuthIdentityChannel: NewAuthIdentityChannelClient(cfg), + BatchImageEvent: NewBatchImageEventClient(cfg), + BatchImageItem: NewBatchImageItemClient(cfg), + BatchImageJob: NewBatchImageJobClient(cfg), ChannelMonitor: NewChannelMonitorClient(cfg), ChannelMonitorDailyRollup: NewChannelMonitorDailyRollupClient(cfg), ChannelMonitorHistory: NewChannelMonitorHistoryClient(cfg), @@ -328,6 +343,9 @@ func (c *Client) BeginTx(ctx context.Context, opts *sql.TxOptions) (*Tx, error) AnnouncementRead: NewAnnouncementReadClient(cfg), AuthIdentity: NewAuthIdentityClient(cfg), AuthIdentityChannel: NewAuthIdentityChannelClient(cfg), + BatchImageEvent: NewBatchImageEventClient(cfg), + BatchImageItem: NewBatchImageItemClient(cfg), + BatchImageJob: NewBatchImageJobClient(cfg), ChannelMonitor: NewChannelMonitorClient(cfg), ChannelMonitorDailyRollup: NewChannelMonitorDailyRollupClient(cfg), ChannelMonitorHistory: NewChannelMonitorHistoryClient(cfg), @@ -386,14 +404,15 @@ func (c *Client) Close() error { func (c *Client) Use(hooks ...Hook) { for _, n := range []interface{ Use(...Hook) }{ c.APIKey, c.Account, c.AccountGroup, c.Announcement, c.AnnouncementRead, - c.AuthIdentity, c.AuthIdentityChannel, c.ChannelMonitor, - c.ChannelMonitorDailyRollup, c.ChannelMonitorHistory, - c.ChannelMonitorRequestTemplate, c.ErrorPassthroughRule, c.Group, - c.IdempotencyRecord, c.IdentityAdoptionDecision, c.PaymentAuditLog, - c.PaymentOrder, c.PaymentProviderInstance, c.PendingAuthSession, c.PromoCode, - c.PromoCodeUsage, c.Proxy, c.RedeemCode, c.SecuritySecret, c.Setting, - c.SubscriptionPlan, c.TLSFingerprintProfile, c.UsageCleanupTask, c.UsageLog, - c.User, c.UserAllowedGroup, c.UserAttributeDefinition, c.UserAttributeValue, + c.AuthIdentity, c.AuthIdentityChannel, c.BatchImageEvent, c.BatchImageItem, + c.BatchImageJob, c.ChannelMonitor, c.ChannelMonitorDailyRollup, + c.ChannelMonitorHistory, c.ChannelMonitorRequestTemplate, + c.ErrorPassthroughRule, c.Group, c.IdempotencyRecord, + c.IdentityAdoptionDecision, c.PaymentAuditLog, c.PaymentOrder, + c.PaymentProviderInstance, c.PendingAuthSession, c.PromoCode, c.PromoCodeUsage, + c.Proxy, c.RedeemCode, c.SecuritySecret, c.Setting, c.SubscriptionPlan, + c.TLSFingerprintProfile, c.UsageCleanupTask, c.UsageLog, c.User, + c.UserAllowedGroup, c.UserAttributeDefinition, c.UserAttributeValue, c.UserPlatformQuota, c.UserSubscription, } { n.Use(hooks...) @@ -405,14 +424,15 @@ func (c *Client) Use(hooks ...Hook) { func (c *Client) Intercept(interceptors ...Interceptor) { for _, n := range []interface{ Intercept(...Interceptor) }{ c.APIKey, c.Account, c.AccountGroup, c.Announcement, c.AnnouncementRead, - c.AuthIdentity, c.AuthIdentityChannel, c.ChannelMonitor, - c.ChannelMonitorDailyRollup, c.ChannelMonitorHistory, - c.ChannelMonitorRequestTemplate, c.ErrorPassthroughRule, c.Group, - c.IdempotencyRecord, c.IdentityAdoptionDecision, c.PaymentAuditLog, - c.PaymentOrder, c.PaymentProviderInstance, c.PendingAuthSession, c.PromoCode, - c.PromoCodeUsage, c.Proxy, c.RedeemCode, c.SecuritySecret, c.Setting, - c.SubscriptionPlan, c.TLSFingerprintProfile, c.UsageCleanupTask, c.UsageLog, - c.User, c.UserAllowedGroup, c.UserAttributeDefinition, c.UserAttributeValue, + c.AuthIdentity, c.AuthIdentityChannel, c.BatchImageEvent, c.BatchImageItem, + c.BatchImageJob, c.ChannelMonitor, c.ChannelMonitorDailyRollup, + c.ChannelMonitorHistory, c.ChannelMonitorRequestTemplate, + c.ErrorPassthroughRule, c.Group, c.IdempotencyRecord, + c.IdentityAdoptionDecision, c.PaymentAuditLog, c.PaymentOrder, + c.PaymentProviderInstance, c.PendingAuthSession, c.PromoCode, c.PromoCodeUsage, + c.Proxy, c.RedeemCode, c.SecuritySecret, c.Setting, c.SubscriptionPlan, + c.TLSFingerprintProfile, c.UsageCleanupTask, c.UsageLog, c.User, + c.UserAllowedGroup, c.UserAttributeDefinition, c.UserAttributeValue, c.UserPlatformQuota, c.UserSubscription, } { n.Intercept(interceptors...) @@ -436,6 +456,12 @@ func (c *Client) Mutate(ctx context.Context, m Mutation) (Value, error) { return c.AuthIdentity.mutate(ctx, m) case *AuthIdentityChannelMutation: return c.AuthIdentityChannel.mutate(ctx, m) + case *BatchImageEventMutation: + return c.BatchImageEvent.mutate(ctx, m) + case *BatchImageItemMutation: + return c.BatchImageItem.mutate(ctx, m) + case *BatchImageJobMutation: + return c.BatchImageJob.mutate(ctx, m) case *ChannelMonitorMutation: return c.ChannelMonitor.mutate(ctx, m) case *ChannelMonitorDailyRollupMutation: @@ -1671,6 +1697,405 @@ func (c *AuthIdentityChannelClient) mutate(ctx context.Context, m *AuthIdentityC } } +// BatchImageEventClient is a client for the BatchImageEvent schema. +type BatchImageEventClient struct { + config +} + +// NewBatchImageEventClient returns a client for the BatchImageEvent from the given config. +func NewBatchImageEventClient(c config) *BatchImageEventClient { + return &BatchImageEventClient{config: c} +} + +// Use adds a list of mutation hooks to the hooks stack. +// A call to `Use(f, g, h)` equals to `batchimageevent.Hooks(f(g(h())))`. +func (c *BatchImageEventClient) Use(hooks ...Hook) { + c.hooks.BatchImageEvent = append(c.hooks.BatchImageEvent, hooks...) +} + +// Intercept adds a list of query interceptors to the interceptors stack. +// A call to `Intercept(f, g, h)` equals to `batchimageevent.Intercept(f(g(h())))`. +func (c *BatchImageEventClient) Intercept(interceptors ...Interceptor) { + c.inters.BatchImageEvent = append(c.inters.BatchImageEvent, interceptors...) +} + +// Create returns a builder for creating a BatchImageEvent entity. +func (c *BatchImageEventClient) Create() *BatchImageEventCreate { + mutation := newBatchImageEventMutation(c.config, OpCreate) + return &BatchImageEventCreate{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// CreateBulk returns a builder for creating a bulk of BatchImageEvent entities. +func (c *BatchImageEventClient) CreateBulk(builders ...*BatchImageEventCreate) *BatchImageEventCreateBulk { + return &BatchImageEventCreateBulk{config: c.config, builders: builders} +} + +// MapCreateBulk creates a bulk creation builder from the given slice. For each item in the slice, the function creates +// a builder and applies setFunc on it. +func (c *BatchImageEventClient) MapCreateBulk(slice any, setFunc func(*BatchImageEventCreate, int)) *BatchImageEventCreateBulk { + rv := reflect.ValueOf(slice) + if rv.Kind() != reflect.Slice { + return &BatchImageEventCreateBulk{err: fmt.Errorf("calling to BatchImageEventClient.MapCreateBulk with wrong type %T, need slice", slice)} + } + builders := make([]*BatchImageEventCreate, rv.Len()) + for i := 0; i < rv.Len(); i++ { + builders[i] = c.Create() + setFunc(builders[i], i) + } + return &BatchImageEventCreateBulk{config: c.config, builders: builders} +} + +// Update returns an update builder for BatchImageEvent. +func (c *BatchImageEventClient) Update() *BatchImageEventUpdate { + mutation := newBatchImageEventMutation(c.config, OpUpdate) + return &BatchImageEventUpdate{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// UpdateOne returns an update builder for the given entity. +func (c *BatchImageEventClient) UpdateOne(_m *BatchImageEvent) *BatchImageEventUpdateOne { + mutation := newBatchImageEventMutation(c.config, OpUpdateOne, withBatchImageEvent(_m)) + return &BatchImageEventUpdateOne{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// UpdateOneID returns an update builder for the given id. +func (c *BatchImageEventClient) UpdateOneID(id int64) *BatchImageEventUpdateOne { + mutation := newBatchImageEventMutation(c.config, OpUpdateOne, withBatchImageEventID(id)) + return &BatchImageEventUpdateOne{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// Delete returns a delete builder for BatchImageEvent. +func (c *BatchImageEventClient) Delete() *BatchImageEventDelete { + mutation := newBatchImageEventMutation(c.config, OpDelete) + return &BatchImageEventDelete{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// DeleteOne returns a builder for deleting the given entity. +func (c *BatchImageEventClient) DeleteOne(_m *BatchImageEvent) *BatchImageEventDeleteOne { + return c.DeleteOneID(_m.ID) +} + +// DeleteOneID returns a builder for deleting the given entity by its id. +func (c *BatchImageEventClient) DeleteOneID(id int64) *BatchImageEventDeleteOne { + builder := c.Delete().Where(batchimageevent.ID(id)) + builder.mutation.id = &id + builder.mutation.op = OpDeleteOne + return &BatchImageEventDeleteOne{builder} +} + +// Query returns a query builder for BatchImageEvent. +func (c *BatchImageEventClient) Query() *BatchImageEventQuery { + return &BatchImageEventQuery{ + config: c.config, + ctx: &QueryContext{Type: TypeBatchImageEvent}, + inters: c.Interceptors(), + } +} + +// Get returns a BatchImageEvent entity by its id. +func (c *BatchImageEventClient) Get(ctx context.Context, id int64) (*BatchImageEvent, error) { + return c.Query().Where(batchimageevent.ID(id)).Only(ctx) +} + +// GetX is like Get, but panics if an error occurs. +func (c *BatchImageEventClient) GetX(ctx context.Context, id int64) *BatchImageEvent { + obj, err := c.Get(ctx, id) + if err != nil { + panic(err) + } + return obj +} + +// Hooks returns the client hooks. +func (c *BatchImageEventClient) Hooks() []Hook { + return c.hooks.BatchImageEvent +} + +// Interceptors returns the client interceptors. +func (c *BatchImageEventClient) Interceptors() []Interceptor { + return c.inters.BatchImageEvent +} + +func (c *BatchImageEventClient) mutate(ctx context.Context, m *BatchImageEventMutation) (Value, error) { + switch m.Op() { + case OpCreate: + return (&BatchImageEventCreate{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpUpdate: + return (&BatchImageEventUpdate{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpUpdateOne: + return (&BatchImageEventUpdateOne{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpDelete, OpDeleteOne: + return (&BatchImageEventDelete{config: c.config, hooks: c.Hooks(), mutation: m}).Exec(ctx) + default: + return nil, fmt.Errorf("ent: unknown BatchImageEvent mutation op: %q", m.Op()) + } +} + +// BatchImageItemClient is a client for the BatchImageItem schema. +type BatchImageItemClient struct { + config +} + +// NewBatchImageItemClient returns a client for the BatchImageItem from the given config. +func NewBatchImageItemClient(c config) *BatchImageItemClient { + return &BatchImageItemClient{config: c} +} + +// Use adds a list of mutation hooks to the hooks stack. +// A call to `Use(f, g, h)` equals to `batchimageitem.Hooks(f(g(h())))`. +func (c *BatchImageItemClient) Use(hooks ...Hook) { + c.hooks.BatchImageItem = append(c.hooks.BatchImageItem, hooks...) +} + +// Intercept adds a list of query interceptors to the interceptors stack. +// A call to `Intercept(f, g, h)` equals to `batchimageitem.Intercept(f(g(h())))`. +func (c *BatchImageItemClient) Intercept(interceptors ...Interceptor) { + c.inters.BatchImageItem = append(c.inters.BatchImageItem, interceptors...) +} + +// Create returns a builder for creating a BatchImageItem entity. +func (c *BatchImageItemClient) Create() *BatchImageItemCreate { + mutation := newBatchImageItemMutation(c.config, OpCreate) + return &BatchImageItemCreate{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// CreateBulk returns a builder for creating a bulk of BatchImageItem entities. +func (c *BatchImageItemClient) CreateBulk(builders ...*BatchImageItemCreate) *BatchImageItemCreateBulk { + return &BatchImageItemCreateBulk{config: c.config, builders: builders} +} + +// MapCreateBulk creates a bulk creation builder from the given slice. For each item in the slice, the function creates +// a builder and applies setFunc on it. +func (c *BatchImageItemClient) MapCreateBulk(slice any, setFunc func(*BatchImageItemCreate, int)) *BatchImageItemCreateBulk { + rv := reflect.ValueOf(slice) + if rv.Kind() != reflect.Slice { + return &BatchImageItemCreateBulk{err: fmt.Errorf("calling to BatchImageItemClient.MapCreateBulk with wrong type %T, need slice", slice)} + } + builders := make([]*BatchImageItemCreate, rv.Len()) + for i := 0; i < rv.Len(); i++ { + builders[i] = c.Create() + setFunc(builders[i], i) + } + return &BatchImageItemCreateBulk{config: c.config, builders: builders} +} + +// Update returns an update builder for BatchImageItem. +func (c *BatchImageItemClient) Update() *BatchImageItemUpdate { + mutation := newBatchImageItemMutation(c.config, OpUpdate) + return &BatchImageItemUpdate{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// UpdateOne returns an update builder for the given entity. +func (c *BatchImageItemClient) UpdateOne(_m *BatchImageItem) *BatchImageItemUpdateOne { + mutation := newBatchImageItemMutation(c.config, OpUpdateOne, withBatchImageItem(_m)) + return &BatchImageItemUpdateOne{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// UpdateOneID returns an update builder for the given id. +func (c *BatchImageItemClient) UpdateOneID(id int64) *BatchImageItemUpdateOne { + mutation := newBatchImageItemMutation(c.config, OpUpdateOne, withBatchImageItemID(id)) + return &BatchImageItemUpdateOne{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// Delete returns a delete builder for BatchImageItem. +func (c *BatchImageItemClient) Delete() *BatchImageItemDelete { + mutation := newBatchImageItemMutation(c.config, OpDelete) + return &BatchImageItemDelete{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// DeleteOne returns a builder for deleting the given entity. +func (c *BatchImageItemClient) DeleteOne(_m *BatchImageItem) *BatchImageItemDeleteOne { + return c.DeleteOneID(_m.ID) +} + +// DeleteOneID returns a builder for deleting the given entity by its id. +func (c *BatchImageItemClient) DeleteOneID(id int64) *BatchImageItemDeleteOne { + builder := c.Delete().Where(batchimageitem.ID(id)) + builder.mutation.id = &id + builder.mutation.op = OpDeleteOne + return &BatchImageItemDeleteOne{builder} +} + +// Query returns a query builder for BatchImageItem. +func (c *BatchImageItemClient) Query() *BatchImageItemQuery { + return &BatchImageItemQuery{ + config: c.config, + ctx: &QueryContext{Type: TypeBatchImageItem}, + inters: c.Interceptors(), + } +} + +// Get returns a BatchImageItem entity by its id. +func (c *BatchImageItemClient) Get(ctx context.Context, id int64) (*BatchImageItem, error) { + return c.Query().Where(batchimageitem.ID(id)).Only(ctx) +} + +// GetX is like Get, but panics if an error occurs. +func (c *BatchImageItemClient) GetX(ctx context.Context, id int64) *BatchImageItem { + obj, err := c.Get(ctx, id) + if err != nil { + panic(err) + } + return obj +} + +// Hooks returns the client hooks. +func (c *BatchImageItemClient) Hooks() []Hook { + return c.hooks.BatchImageItem +} + +// Interceptors returns the client interceptors. +func (c *BatchImageItemClient) Interceptors() []Interceptor { + return c.inters.BatchImageItem +} + +func (c *BatchImageItemClient) mutate(ctx context.Context, m *BatchImageItemMutation) (Value, error) { + switch m.Op() { + case OpCreate: + return (&BatchImageItemCreate{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpUpdate: + return (&BatchImageItemUpdate{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpUpdateOne: + return (&BatchImageItemUpdateOne{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpDelete, OpDeleteOne: + return (&BatchImageItemDelete{config: c.config, hooks: c.Hooks(), mutation: m}).Exec(ctx) + default: + return nil, fmt.Errorf("ent: unknown BatchImageItem mutation op: %q", m.Op()) + } +} + +// BatchImageJobClient is a client for the BatchImageJob schema. +type BatchImageJobClient struct { + config +} + +// NewBatchImageJobClient returns a client for the BatchImageJob from the given config. +func NewBatchImageJobClient(c config) *BatchImageJobClient { + return &BatchImageJobClient{config: c} +} + +// Use adds a list of mutation hooks to the hooks stack. +// A call to `Use(f, g, h)` equals to `batchimagejob.Hooks(f(g(h())))`. +func (c *BatchImageJobClient) Use(hooks ...Hook) { + c.hooks.BatchImageJob = append(c.hooks.BatchImageJob, hooks...) +} + +// Intercept adds a list of query interceptors to the interceptors stack. +// A call to `Intercept(f, g, h)` equals to `batchimagejob.Intercept(f(g(h())))`. +func (c *BatchImageJobClient) Intercept(interceptors ...Interceptor) { + c.inters.BatchImageJob = append(c.inters.BatchImageJob, interceptors...) +} + +// Create returns a builder for creating a BatchImageJob entity. +func (c *BatchImageJobClient) Create() *BatchImageJobCreate { + mutation := newBatchImageJobMutation(c.config, OpCreate) + return &BatchImageJobCreate{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// CreateBulk returns a builder for creating a bulk of BatchImageJob entities. +func (c *BatchImageJobClient) CreateBulk(builders ...*BatchImageJobCreate) *BatchImageJobCreateBulk { + return &BatchImageJobCreateBulk{config: c.config, builders: builders} +} + +// MapCreateBulk creates a bulk creation builder from the given slice. For each item in the slice, the function creates +// a builder and applies setFunc on it. +func (c *BatchImageJobClient) MapCreateBulk(slice any, setFunc func(*BatchImageJobCreate, int)) *BatchImageJobCreateBulk { + rv := reflect.ValueOf(slice) + if rv.Kind() != reflect.Slice { + return &BatchImageJobCreateBulk{err: fmt.Errorf("calling to BatchImageJobClient.MapCreateBulk with wrong type %T, need slice", slice)} + } + builders := make([]*BatchImageJobCreate, rv.Len()) + for i := 0; i < rv.Len(); i++ { + builders[i] = c.Create() + setFunc(builders[i], i) + } + return &BatchImageJobCreateBulk{config: c.config, builders: builders} +} + +// Update returns an update builder for BatchImageJob. +func (c *BatchImageJobClient) Update() *BatchImageJobUpdate { + mutation := newBatchImageJobMutation(c.config, OpUpdate) + return &BatchImageJobUpdate{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// UpdateOne returns an update builder for the given entity. +func (c *BatchImageJobClient) UpdateOne(_m *BatchImageJob) *BatchImageJobUpdateOne { + mutation := newBatchImageJobMutation(c.config, OpUpdateOne, withBatchImageJob(_m)) + return &BatchImageJobUpdateOne{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// UpdateOneID returns an update builder for the given id. +func (c *BatchImageJobClient) UpdateOneID(id int64) *BatchImageJobUpdateOne { + mutation := newBatchImageJobMutation(c.config, OpUpdateOne, withBatchImageJobID(id)) + return &BatchImageJobUpdateOne{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// Delete returns a delete builder for BatchImageJob. +func (c *BatchImageJobClient) Delete() *BatchImageJobDelete { + mutation := newBatchImageJobMutation(c.config, OpDelete) + return &BatchImageJobDelete{config: c.config, hooks: c.Hooks(), mutation: mutation} +} + +// DeleteOne returns a builder for deleting the given entity. +func (c *BatchImageJobClient) DeleteOne(_m *BatchImageJob) *BatchImageJobDeleteOne { + return c.DeleteOneID(_m.ID) +} + +// DeleteOneID returns a builder for deleting the given entity by its id. +func (c *BatchImageJobClient) DeleteOneID(id int64) *BatchImageJobDeleteOne { + builder := c.Delete().Where(batchimagejob.ID(id)) + builder.mutation.id = &id + builder.mutation.op = OpDeleteOne + return &BatchImageJobDeleteOne{builder} +} + +// Query returns a query builder for BatchImageJob. +func (c *BatchImageJobClient) Query() *BatchImageJobQuery { + return &BatchImageJobQuery{ + config: c.config, + ctx: &QueryContext{Type: TypeBatchImageJob}, + inters: c.Interceptors(), + } +} + +// Get returns a BatchImageJob entity by its id. +func (c *BatchImageJobClient) Get(ctx context.Context, id int64) (*BatchImageJob, error) { + return c.Query().Where(batchimagejob.ID(id)).Only(ctx) +} + +// GetX is like Get, but panics if an error occurs. +func (c *BatchImageJobClient) GetX(ctx context.Context, id int64) *BatchImageJob { + obj, err := c.Get(ctx, id) + if err != nil { + panic(err) + } + return obj +} + +// Hooks returns the client hooks. +func (c *BatchImageJobClient) Hooks() []Hook { + return c.hooks.BatchImageJob +} + +// Interceptors returns the client interceptors. +func (c *BatchImageJobClient) Interceptors() []Interceptor { + return c.inters.BatchImageJob +} + +func (c *BatchImageJobClient) mutate(ctx context.Context, m *BatchImageJobMutation) (Value, error) { + switch m.Op() { + case OpCreate: + return (&BatchImageJobCreate{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpUpdate: + return (&BatchImageJobUpdate{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpUpdateOne: + return (&BatchImageJobUpdateOne{config: c.config, hooks: c.Hooks(), mutation: m}).Save(ctx) + case OpDelete, OpDeleteOne: + return (&BatchImageJobDelete{config: c.config, hooks: c.Hooks(), mutation: m}).Exec(ctx) + default: + return nil, fmt.Errorf("ent: unknown BatchImageJob mutation op: %q", m.Op()) + } +} + // ChannelMonitorClient is a client for the ChannelMonitor schema. type ChannelMonitorClient struct { config @@ -6242,25 +6667,25 @@ func (c *UserSubscriptionClient) mutate(ctx context.Context, m *UserSubscription type ( hooks struct { APIKey, Account, AccountGroup, Announcement, AnnouncementRead, AuthIdentity, - AuthIdentityChannel, ChannelMonitor, ChannelMonitorDailyRollup, - ChannelMonitorHistory, ChannelMonitorRequestTemplate, ErrorPassthroughRule, - Group, IdempotencyRecord, IdentityAdoptionDecision, PaymentAuditLog, - PaymentOrder, PaymentProviderInstance, PendingAuthSession, PromoCode, - PromoCodeUsage, Proxy, RedeemCode, SecuritySecret, Setting, SubscriptionPlan, - TLSFingerprintProfile, UsageCleanupTask, UsageLog, User, UserAllowedGroup, - UserAttributeDefinition, UserAttributeValue, UserPlatformQuota, - UserSubscription []ent.Hook + AuthIdentityChannel, BatchImageEvent, BatchImageItem, BatchImageJob, + ChannelMonitor, ChannelMonitorDailyRollup, ChannelMonitorHistory, + ChannelMonitorRequestTemplate, ErrorPassthroughRule, Group, IdempotencyRecord, + IdentityAdoptionDecision, PaymentAuditLog, PaymentOrder, + PaymentProviderInstance, PendingAuthSession, PromoCode, PromoCodeUsage, Proxy, + RedeemCode, SecuritySecret, Setting, SubscriptionPlan, TLSFingerprintProfile, + UsageCleanupTask, UsageLog, User, UserAllowedGroup, UserAttributeDefinition, + UserAttributeValue, UserPlatformQuota, UserSubscription []ent.Hook } inters struct { APIKey, Account, AccountGroup, Announcement, AnnouncementRead, AuthIdentity, - AuthIdentityChannel, ChannelMonitor, ChannelMonitorDailyRollup, - ChannelMonitorHistory, ChannelMonitorRequestTemplate, ErrorPassthroughRule, - Group, IdempotencyRecord, IdentityAdoptionDecision, PaymentAuditLog, - PaymentOrder, PaymentProviderInstance, PendingAuthSession, PromoCode, - PromoCodeUsage, Proxy, RedeemCode, SecuritySecret, Setting, SubscriptionPlan, - TLSFingerprintProfile, UsageCleanupTask, UsageLog, User, UserAllowedGroup, - UserAttributeDefinition, UserAttributeValue, UserPlatformQuota, - UserSubscription []ent.Interceptor + AuthIdentityChannel, BatchImageEvent, BatchImageItem, BatchImageJob, + ChannelMonitor, ChannelMonitorDailyRollup, ChannelMonitorHistory, + ChannelMonitorRequestTemplate, ErrorPassthroughRule, Group, IdempotencyRecord, + IdentityAdoptionDecision, PaymentAuditLog, PaymentOrder, + PaymentProviderInstance, PendingAuthSession, PromoCode, PromoCodeUsage, Proxy, + RedeemCode, SecuritySecret, Setting, SubscriptionPlan, TLSFingerprintProfile, + UsageCleanupTask, UsageLog, User, UserAllowedGroup, UserAttributeDefinition, + UserAttributeValue, UserPlatformQuota, UserSubscription []ent.Interceptor } ) diff --git a/backend/ent/ent.go b/backend/ent/ent.go index 33d36e70ee..d23f61327f 100644 --- a/backend/ent/ent.go +++ b/backend/ent/ent.go @@ -19,6 +19,9 @@ import ( "github.com/Wei-Shaw/sub2api/ent/apikey" "github.com/Wei-Shaw/sub2api/ent/authidentity" "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" "github.com/Wei-Shaw/sub2api/ent/channelmonitor" "github.com/Wei-Shaw/sub2api/ent/channelmonitordailyrollup" "github.com/Wei-Shaw/sub2api/ent/channelmonitorhistory" @@ -114,6 +117,9 @@ func checkColumn(t, c string) error { announcementread.Table: announcementread.ValidColumn, authidentity.Table: authidentity.ValidColumn, authidentitychannel.Table: authidentitychannel.ValidColumn, + batchimageevent.Table: batchimageevent.ValidColumn, + batchimageitem.Table: batchimageitem.ValidColumn, + batchimagejob.Table: batchimagejob.ValidColumn, channelmonitor.Table: channelmonitor.ValidColumn, channelmonitordailyrollup.Table: channelmonitordailyrollup.ValidColumn, channelmonitorhistory.Table: channelmonitorhistory.ValidColumn, diff --git a/backend/ent/group.go b/backend/ent/group.go index 172b67777e..5bec594977 100644 --- a/backend/ent/group.go +++ b/backend/ent/group.go @@ -33,9 +33,9 @@ type Group struct { RateMultiplier float64 `json:"rate_multiplier,omitempty"` // 是否启用高峰时段倍率 PeakRateEnabled bool `json:"peak_rate_enabled,omitempty"` - // 高峰开始时间 HH:MM(含),如 14:00;空表示未配置 + // 高峰开始时间 HH:MM(含),如 14:00;空表示未配置;不支持跨天 PeakStart string `json:"peak_start,omitempty"` - // 高峰结束时间 HH:MM(不含),如 18:00 + // 高峰结束时间 HH:MM(不含),必须大于 peak_start;不支持跨天,如 22:00-02:00 PeakEnd string `json:"peak_end,omitempty"` // 高峰时段叠加倍率,仅在 peak_rate_enabled 且处于 [peak_start, peak_end) 时乘入文本倍率 PeakRateMultiplier float64 `json:"peak_rate_multiplier,omitempty"` @@ -57,6 +57,8 @@ type Group struct { DefaultValidityDays int `json:"default_validity_days,omitempty"` // 是否允许该分组使用图片生成能力 AllowImageGeneration bool `json:"allow_image_generation,omitempty"` + // 是否允许该分组使用批量图片生成能力 + AllowBatchImageGeneration bool `json:"allow_batch_image_generation,omitempty"` // 图片生成是否使用独立倍率;false 表示共享分组有效倍率 ImageRateIndependent bool `json:"image_rate_independent,omitempty"` // 图片生成独立倍率,仅 image_rate_independent=true 时生效 @@ -67,6 +69,20 @@ type Group struct { ImagePrice2k *float64 `json:"image_price_2k,omitempty"` // ImagePrice4k holds the value of the "image_price_4k" field. ImagePrice4k *float64 `json:"image_price_4k,omitempty"` + // 批量图片生成折扣倍率,最终单价会乘以该值;0 表示免费 + BatchImageDiscountMultiplier float64 `json:"batch_image_discount_multiplier,omitempty"` + // 批量图片生成冻结价格比例,按普通生图原价乘以该比例冻结,结算后释放差额 + BatchImageHoldMultiplier float64 `json:"batch_image_hold_multiplier,omitempty"` + // 视频生成是否使用独立倍率;false 表示共享分组有效倍率 + VideoRateIndependent bool `json:"video_rate_independent,omitempty"` + // 视频生成独立倍率,仅 video_rate_independent=true 时生效 + VideoRateMultiplier float64 `json:"video_rate_multiplier,omitempty"` + // VideoPrice480p holds the value of the "video_price_480p" field. + VideoPrice480p *float64 `json:"video_price_480p,omitempty"` + // VideoPrice720p holds the value of the "video_price_720p" field. + VideoPrice720p *float64 `json:"video_price_720p,omitempty"` + // VideoPrice1080p holds the value of the "video_price_1080p" field. + VideoPrice1080p *float64 `json:"video_price_1080p,omitempty"` // 是否仅允许 Claude Code 客户端 ClaudeCodeOnly bool `json:"claude_code_only,omitempty"` // 非 Claude Code 请求降级使用的分组 ID @@ -205,9 +221,9 @@ func (*Group) scanValues(columns []string) ([]any, error) { switch columns[i] { case group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig: values[i] = new([]byte) - case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldImageRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet: + case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet: values[i] = new(sql.NullBool) - case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k: + case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldBatchImageDiscountMultiplier, group.FieldBatchImageHoldMultiplier, group.FieldVideoRateMultiplier, group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p: values[i] = new(sql.NullFloat64) case group.FieldID, group.FieldDefaultValidityDays, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest, group.FieldSortOrder, group.FieldRpmLimit: values[i] = new(sql.NullInt64) @@ -355,6 +371,12 @@ func (_m *Group) assignValues(columns []string, values []any) error { } else if value.Valid { _m.AllowImageGeneration = value.Bool } + case group.FieldAllowBatchImageGeneration: + if value, ok := values[i].(*sql.NullBool); !ok { + return fmt.Errorf("unexpected type %T for field allow_batch_image_generation", values[i]) + } else if value.Valid { + _m.AllowBatchImageGeneration = value.Bool + } case group.FieldImageRateIndependent: if value, ok := values[i].(*sql.NullBool); !ok { return fmt.Errorf("unexpected type %T for field image_rate_independent", values[i]) @@ -388,6 +410,51 @@ func (_m *Group) assignValues(columns []string, values []any) error { _m.ImagePrice4k = new(float64) *_m.ImagePrice4k = value.Float64 } + case group.FieldBatchImageDiscountMultiplier: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field batch_image_discount_multiplier", values[i]) + } else if value.Valid { + _m.BatchImageDiscountMultiplier = value.Float64 + } + case group.FieldBatchImageHoldMultiplier: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field batch_image_hold_multiplier", values[i]) + } else if value.Valid { + _m.BatchImageHoldMultiplier = value.Float64 + } + case group.FieldVideoRateIndependent: + if value, ok := values[i].(*sql.NullBool); !ok { + return fmt.Errorf("unexpected type %T for field video_rate_independent", values[i]) + } else if value.Valid { + _m.VideoRateIndependent = value.Bool + } + case group.FieldVideoRateMultiplier: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field video_rate_multiplier", values[i]) + } else if value.Valid { + _m.VideoRateMultiplier = value.Float64 + } + case group.FieldVideoPrice480p: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field video_price_480p", values[i]) + } else if value.Valid { + _m.VideoPrice480p = new(float64) + *_m.VideoPrice480p = value.Float64 + } + case group.FieldVideoPrice720p: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field video_price_720p", values[i]) + } else if value.Valid { + _m.VideoPrice720p = new(float64) + *_m.VideoPrice720p = value.Float64 + } + case group.FieldVideoPrice1080p: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field video_price_1080p", values[i]) + } else if value.Valid { + _m.VideoPrice1080p = new(float64) + *_m.VideoPrice1080p = value.Float64 + } case group.FieldClaudeCodeOnly: if value, ok := values[i].(*sql.NullBool); !ok { return fmt.Errorf("unexpected type %T for field claude_code_only", values[i]) @@ -631,6 +698,9 @@ func (_m *Group) String() string { builder.WriteString("allow_image_generation=") builder.WriteString(fmt.Sprintf("%v", _m.AllowImageGeneration)) builder.WriteString(", ") + builder.WriteString("allow_batch_image_generation=") + builder.WriteString(fmt.Sprintf("%v", _m.AllowBatchImageGeneration)) + builder.WriteString(", ") builder.WriteString("image_rate_independent=") builder.WriteString(fmt.Sprintf("%v", _m.ImageRateIndependent)) builder.WriteString(", ") @@ -652,6 +722,33 @@ func (_m *Group) String() string { builder.WriteString(fmt.Sprintf("%v", *v)) } builder.WriteString(", ") + builder.WriteString("batch_image_discount_multiplier=") + builder.WriteString(fmt.Sprintf("%v", _m.BatchImageDiscountMultiplier)) + builder.WriteString(", ") + builder.WriteString("batch_image_hold_multiplier=") + builder.WriteString(fmt.Sprintf("%v", _m.BatchImageHoldMultiplier)) + builder.WriteString(", ") + builder.WriteString("video_rate_independent=") + builder.WriteString(fmt.Sprintf("%v", _m.VideoRateIndependent)) + builder.WriteString(", ") + builder.WriteString("video_rate_multiplier=") + builder.WriteString(fmt.Sprintf("%v", _m.VideoRateMultiplier)) + builder.WriteString(", ") + if v := _m.VideoPrice480p; v != nil { + builder.WriteString("video_price_480p=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.VideoPrice720p; v != nil { + builder.WriteString("video_price_720p=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") + if v := _m.VideoPrice1080p; v != nil { + builder.WriteString("video_price_1080p=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") builder.WriteString("claude_code_only=") builder.WriteString(fmt.Sprintf("%v", _m.ClaudeCodeOnly)) builder.WriteString(", ") diff --git a/backend/ent/group/group.go b/backend/ent/group/group.go index bc95af71b6..769c63e6b1 100644 --- a/backend/ent/group/group.go +++ b/backend/ent/group/group.go @@ -54,6 +54,8 @@ const ( FieldDefaultValidityDays = "default_validity_days" // FieldAllowImageGeneration holds the string denoting the allow_image_generation field in the database. FieldAllowImageGeneration = "allow_image_generation" + // FieldAllowBatchImageGeneration holds the string denoting the allow_batch_image_generation field in the database. + FieldAllowBatchImageGeneration = "allow_batch_image_generation" // FieldImageRateIndependent holds the string denoting the image_rate_independent field in the database. FieldImageRateIndependent = "image_rate_independent" // FieldImageRateMultiplier holds the string denoting the image_rate_multiplier field in the database. @@ -64,6 +66,20 @@ const ( FieldImagePrice2k = "image_price_2k" // FieldImagePrice4k holds the string denoting the image_price_4k field in the database. FieldImagePrice4k = "image_price_4k" + // FieldBatchImageDiscountMultiplier holds the string denoting the batch_image_discount_multiplier field in the database. + FieldBatchImageDiscountMultiplier = "batch_image_discount_multiplier" + // FieldBatchImageHoldMultiplier holds the string denoting the batch_image_hold_multiplier field in the database. + FieldBatchImageHoldMultiplier = "batch_image_hold_multiplier" + // FieldVideoRateIndependent holds the string denoting the video_rate_independent field in the database. + FieldVideoRateIndependent = "video_rate_independent" + // FieldVideoRateMultiplier holds the string denoting the video_rate_multiplier field in the database. + FieldVideoRateMultiplier = "video_rate_multiplier" + // FieldVideoPrice480p holds the string denoting the video_price_480p field in the database. + FieldVideoPrice480p = "video_price_480p" + // FieldVideoPrice720p holds the string denoting the video_price_720p field in the database. + FieldVideoPrice720p = "video_price_720p" + // FieldVideoPrice1080p holds the string denoting the video_price_1080p field in the database. + FieldVideoPrice1080p = "video_price_1080p" // FieldClaudeCodeOnly holds the string denoting the claude_code_only field in the database. FieldClaudeCodeOnly = "claude_code_only" // FieldFallbackGroupID holds the string denoting the fallback_group_id field in the database. @@ -188,11 +204,19 @@ var Columns = []string{ FieldMonthlyLimitUsd, FieldDefaultValidityDays, FieldAllowImageGeneration, + FieldAllowBatchImageGeneration, FieldImageRateIndependent, FieldImageRateMultiplier, FieldImagePrice1k, FieldImagePrice2k, FieldImagePrice4k, + FieldBatchImageDiscountMultiplier, + FieldBatchImageHoldMultiplier, + FieldVideoRateIndependent, + FieldVideoRateMultiplier, + FieldVideoPrice480p, + FieldVideoPrice720p, + FieldVideoPrice1080p, FieldClaudeCodeOnly, FieldFallbackGroupID, FieldFallbackGroupIDOnInvalidRequest, @@ -277,10 +301,20 @@ var ( DefaultDefaultValidityDays int // DefaultAllowImageGeneration holds the default value on creation for the "allow_image_generation" field. DefaultAllowImageGeneration bool + // DefaultAllowBatchImageGeneration holds the default value on creation for the "allow_batch_image_generation" field. + DefaultAllowBatchImageGeneration bool // DefaultImageRateIndependent holds the default value on creation for the "image_rate_independent" field. DefaultImageRateIndependent bool // DefaultImageRateMultiplier holds the default value on creation for the "image_rate_multiplier" field. DefaultImageRateMultiplier float64 + // DefaultBatchImageDiscountMultiplier holds the default value on creation for the "batch_image_discount_multiplier" field. + DefaultBatchImageDiscountMultiplier float64 + // DefaultBatchImageHoldMultiplier holds the default value on creation for the "batch_image_hold_multiplier" field. + DefaultBatchImageHoldMultiplier float64 + // DefaultVideoRateIndependent holds the default value on creation for the "video_rate_independent" field. + DefaultVideoRateIndependent bool + // DefaultVideoRateMultiplier holds the default value on creation for the "video_rate_multiplier" field. + DefaultVideoRateMultiplier float64 // DefaultClaudeCodeOnly holds the default value on creation for the "claude_code_only" field. DefaultClaudeCodeOnly bool // DefaultModelRoutingEnabled holds the default value on creation for the "model_routing_enabled" field. @@ -412,6 +446,11 @@ func ByAllowImageGeneration(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldAllowImageGeneration, opts...).ToFunc() } +// ByAllowBatchImageGeneration orders the results by the allow_batch_image_generation field. +func ByAllowBatchImageGeneration(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldAllowBatchImageGeneration, opts...).ToFunc() +} + // ByImageRateIndependent orders the results by the image_rate_independent field. func ByImageRateIndependent(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldImageRateIndependent, opts...).ToFunc() @@ -437,6 +476,41 @@ func ByImagePrice4k(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldImagePrice4k, opts...).ToFunc() } +// ByBatchImageDiscountMultiplier orders the results by the batch_image_discount_multiplier field. +func ByBatchImageDiscountMultiplier(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldBatchImageDiscountMultiplier, opts...).ToFunc() +} + +// ByBatchImageHoldMultiplier orders the results by the batch_image_hold_multiplier field. +func ByBatchImageHoldMultiplier(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldBatchImageHoldMultiplier, opts...).ToFunc() +} + +// ByVideoRateIndependent orders the results by the video_rate_independent field. +func ByVideoRateIndependent(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoRateIndependent, opts...).ToFunc() +} + +// ByVideoRateMultiplier orders the results by the video_rate_multiplier field. +func ByVideoRateMultiplier(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoRateMultiplier, opts...).ToFunc() +} + +// ByVideoPrice480p orders the results by the video_price_480p field. +func ByVideoPrice480p(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoPrice480p, opts...).ToFunc() +} + +// ByVideoPrice720p orders the results by the video_price_720p field. +func ByVideoPrice720p(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoPrice720p, opts...).ToFunc() +} + +// ByVideoPrice1080p orders the results by the video_price_1080p field. +func ByVideoPrice1080p(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoPrice1080p, opts...).ToFunc() +} + // ByClaudeCodeOnly orders the results by the claude_code_only field. func ByClaudeCodeOnly(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldClaudeCodeOnly, opts...).ToFunc() diff --git a/backend/ent/group/where.go b/backend/ent/group/where.go index 4a7fc01991..5a9d92d0f4 100644 --- a/backend/ent/group/where.go +++ b/backend/ent/group/where.go @@ -150,6 +150,11 @@ func AllowImageGeneration(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldAllowImageGeneration, v)) } +// AllowBatchImageGeneration applies equality check predicate on the "allow_batch_image_generation" field. It's identical to AllowBatchImageGenerationEQ. +func AllowBatchImageGeneration(v bool) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldAllowBatchImageGeneration, v)) +} + // ImageRateIndependent applies equality check predicate on the "image_rate_independent" field. It's identical to ImageRateIndependentEQ. func ImageRateIndependent(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldImageRateIndependent, v)) @@ -175,6 +180,41 @@ func ImagePrice4k(v float64) predicate.Group { return predicate.Group(sql.FieldEQ(FieldImagePrice4k, v)) } +// BatchImageDiscountMultiplier applies equality check predicate on the "batch_image_discount_multiplier" field. It's identical to BatchImageDiscountMultiplierEQ. +func BatchImageDiscountMultiplier(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageHoldMultiplier applies equality check predicate on the "batch_image_hold_multiplier" field. It's identical to BatchImageHoldMultiplierEQ. +func BatchImageHoldMultiplier(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldBatchImageHoldMultiplier, v)) +} + +// VideoRateIndependent applies equality check predicate on the "video_rate_independent" field. It's identical to VideoRateIndependentEQ. +func VideoRateIndependent(v bool) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoRateIndependent, v)) +} + +// VideoRateMultiplier applies equality check predicate on the "video_rate_multiplier" field. It's identical to VideoRateMultiplierEQ. +func VideoRateMultiplier(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoRateMultiplier, v)) +} + +// VideoPrice480p applies equality check predicate on the "video_price_480p" field. It's identical to VideoPrice480pEQ. +func VideoPrice480p(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoPrice480p, v)) +} + +// VideoPrice720p applies equality check predicate on the "video_price_720p" field. It's identical to VideoPrice720pEQ. +func VideoPrice720p(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoPrice720p, v)) +} + +// VideoPrice1080p applies equality check predicate on the "video_price_1080p" field. It's identical to VideoPrice1080pEQ. +func VideoPrice1080p(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoPrice1080p, v)) +} + // ClaudeCodeOnly applies equality check predicate on the "claude_code_only" field. It's identical to ClaudeCodeOnlyEQ. func ClaudeCodeOnly(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v)) @@ -1125,6 +1165,16 @@ func AllowImageGenerationNEQ(v bool) predicate.Group { return predicate.Group(sql.FieldNEQ(FieldAllowImageGeneration, v)) } +// AllowBatchImageGenerationEQ applies the EQ predicate on the "allow_batch_image_generation" field. +func AllowBatchImageGenerationEQ(v bool) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldAllowBatchImageGeneration, v)) +} + +// AllowBatchImageGenerationNEQ applies the NEQ predicate on the "allow_batch_image_generation" field. +func AllowBatchImageGenerationNEQ(v bool) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldAllowBatchImageGeneration, v)) +} + // ImageRateIndependentEQ applies the EQ predicate on the "image_rate_independent" field. func ImageRateIndependentEQ(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldImageRateIndependent, v)) @@ -1325,6 +1375,286 @@ func ImagePrice4kNotNil() predicate.Group { return predicate.Group(sql.FieldNotNull(FieldImagePrice4k)) } +// BatchImageDiscountMultiplierEQ applies the EQ predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageDiscountMultiplierNEQ applies the NEQ predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageDiscountMultiplierIn applies the In predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldBatchImageDiscountMultiplier, vs...)) +} + +// BatchImageDiscountMultiplierNotIn applies the NotIn predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldBatchImageDiscountMultiplier, vs...)) +} + +// BatchImageDiscountMultiplierGT applies the GT predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageDiscountMultiplierGTE applies the GTE predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageDiscountMultiplierLT applies the LT predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageDiscountMultiplierLTE applies the LTE predicate on the "batch_image_discount_multiplier" field. +func BatchImageDiscountMultiplierLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldBatchImageDiscountMultiplier, v)) +} + +// BatchImageHoldMultiplierEQ applies the EQ predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldBatchImageHoldMultiplier, v)) +} + +// BatchImageHoldMultiplierNEQ applies the NEQ predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldBatchImageHoldMultiplier, v)) +} + +// BatchImageHoldMultiplierIn applies the In predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldBatchImageHoldMultiplier, vs...)) +} + +// BatchImageHoldMultiplierNotIn applies the NotIn predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldBatchImageHoldMultiplier, vs...)) +} + +// BatchImageHoldMultiplierGT applies the GT predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldBatchImageHoldMultiplier, v)) +} + +// BatchImageHoldMultiplierGTE applies the GTE predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldBatchImageHoldMultiplier, v)) +} + +// BatchImageHoldMultiplierLT applies the LT predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldBatchImageHoldMultiplier, v)) +} + +// BatchImageHoldMultiplierLTE applies the LTE predicate on the "batch_image_hold_multiplier" field. +func BatchImageHoldMultiplierLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldBatchImageHoldMultiplier, v)) +} + +// VideoRateIndependentEQ applies the EQ predicate on the "video_rate_independent" field. +func VideoRateIndependentEQ(v bool) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoRateIndependent, v)) +} + +// VideoRateIndependentNEQ applies the NEQ predicate on the "video_rate_independent" field. +func VideoRateIndependentNEQ(v bool) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldVideoRateIndependent, v)) +} + +// VideoRateMultiplierEQ applies the EQ predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoRateMultiplier, v)) +} + +// VideoRateMultiplierNEQ applies the NEQ predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldVideoRateMultiplier, v)) +} + +// VideoRateMultiplierIn applies the In predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldVideoRateMultiplier, vs...)) +} + +// VideoRateMultiplierNotIn applies the NotIn predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldVideoRateMultiplier, vs...)) +} + +// VideoRateMultiplierGT applies the GT predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldVideoRateMultiplier, v)) +} + +// VideoRateMultiplierGTE applies the GTE predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldVideoRateMultiplier, v)) +} + +// VideoRateMultiplierLT applies the LT predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldVideoRateMultiplier, v)) +} + +// VideoRateMultiplierLTE applies the LTE predicate on the "video_rate_multiplier" field. +func VideoRateMultiplierLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldVideoRateMultiplier, v)) +} + +// VideoPrice480pEQ applies the EQ predicate on the "video_price_480p" field. +func VideoPrice480pEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoPrice480p, v)) +} + +// VideoPrice480pNEQ applies the NEQ predicate on the "video_price_480p" field. +func VideoPrice480pNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldVideoPrice480p, v)) +} + +// VideoPrice480pIn applies the In predicate on the "video_price_480p" field. +func VideoPrice480pIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldVideoPrice480p, vs...)) +} + +// VideoPrice480pNotIn applies the NotIn predicate on the "video_price_480p" field. +func VideoPrice480pNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldVideoPrice480p, vs...)) +} + +// VideoPrice480pGT applies the GT predicate on the "video_price_480p" field. +func VideoPrice480pGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldVideoPrice480p, v)) +} + +// VideoPrice480pGTE applies the GTE predicate on the "video_price_480p" field. +func VideoPrice480pGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldVideoPrice480p, v)) +} + +// VideoPrice480pLT applies the LT predicate on the "video_price_480p" field. +func VideoPrice480pLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldVideoPrice480p, v)) +} + +// VideoPrice480pLTE applies the LTE predicate on the "video_price_480p" field. +func VideoPrice480pLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldVideoPrice480p, v)) +} + +// VideoPrice480pIsNil applies the IsNil predicate on the "video_price_480p" field. +func VideoPrice480pIsNil() predicate.Group { + return predicate.Group(sql.FieldIsNull(FieldVideoPrice480p)) +} + +// VideoPrice480pNotNil applies the NotNil predicate on the "video_price_480p" field. +func VideoPrice480pNotNil() predicate.Group { + return predicate.Group(sql.FieldNotNull(FieldVideoPrice480p)) +} + +// VideoPrice720pEQ applies the EQ predicate on the "video_price_720p" field. +func VideoPrice720pEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoPrice720p, v)) +} + +// VideoPrice720pNEQ applies the NEQ predicate on the "video_price_720p" field. +func VideoPrice720pNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldVideoPrice720p, v)) +} + +// VideoPrice720pIn applies the In predicate on the "video_price_720p" field. +func VideoPrice720pIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldVideoPrice720p, vs...)) +} + +// VideoPrice720pNotIn applies the NotIn predicate on the "video_price_720p" field. +func VideoPrice720pNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldVideoPrice720p, vs...)) +} + +// VideoPrice720pGT applies the GT predicate on the "video_price_720p" field. +func VideoPrice720pGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldVideoPrice720p, v)) +} + +// VideoPrice720pGTE applies the GTE predicate on the "video_price_720p" field. +func VideoPrice720pGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldVideoPrice720p, v)) +} + +// VideoPrice720pLT applies the LT predicate on the "video_price_720p" field. +func VideoPrice720pLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldVideoPrice720p, v)) +} + +// VideoPrice720pLTE applies the LTE predicate on the "video_price_720p" field. +func VideoPrice720pLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldVideoPrice720p, v)) +} + +// VideoPrice720pIsNil applies the IsNil predicate on the "video_price_720p" field. +func VideoPrice720pIsNil() predicate.Group { + return predicate.Group(sql.FieldIsNull(FieldVideoPrice720p)) +} + +// VideoPrice720pNotNil applies the NotNil predicate on the "video_price_720p" field. +func VideoPrice720pNotNil() predicate.Group { + return predicate.Group(sql.FieldNotNull(FieldVideoPrice720p)) +} + +// VideoPrice1080pEQ applies the EQ predicate on the "video_price_1080p" field. +func VideoPrice1080pEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldVideoPrice1080p, v)) +} + +// VideoPrice1080pNEQ applies the NEQ predicate on the "video_price_1080p" field. +func VideoPrice1080pNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldVideoPrice1080p, v)) +} + +// VideoPrice1080pIn applies the In predicate on the "video_price_1080p" field. +func VideoPrice1080pIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldVideoPrice1080p, vs...)) +} + +// VideoPrice1080pNotIn applies the NotIn predicate on the "video_price_1080p" field. +func VideoPrice1080pNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldVideoPrice1080p, vs...)) +} + +// VideoPrice1080pGT applies the GT predicate on the "video_price_1080p" field. +func VideoPrice1080pGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldVideoPrice1080p, v)) +} + +// VideoPrice1080pGTE applies the GTE predicate on the "video_price_1080p" field. +func VideoPrice1080pGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldVideoPrice1080p, v)) +} + +// VideoPrice1080pLT applies the LT predicate on the "video_price_1080p" field. +func VideoPrice1080pLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldVideoPrice1080p, v)) +} + +// VideoPrice1080pLTE applies the LTE predicate on the "video_price_1080p" field. +func VideoPrice1080pLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldVideoPrice1080p, v)) +} + +// VideoPrice1080pIsNil applies the IsNil predicate on the "video_price_1080p" field. +func VideoPrice1080pIsNil() predicate.Group { + return predicate.Group(sql.FieldIsNull(FieldVideoPrice1080p)) +} + +// VideoPrice1080pNotNil applies the NotNil predicate on the "video_price_1080p" field. +func VideoPrice1080pNotNil() predicate.Group { + return predicate.Group(sql.FieldNotNull(FieldVideoPrice1080p)) +} + // ClaudeCodeOnlyEQ applies the EQ predicate on the "claude_code_only" field. func ClaudeCodeOnlyEQ(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v)) diff --git a/backend/ent/group_create.go b/backend/ent/group_create.go index 0f35f070ee..2a6c18e67d 100644 --- a/backend/ent/group_create.go +++ b/backend/ent/group_create.go @@ -287,6 +287,20 @@ func (_c *GroupCreate) SetNillableAllowImageGeneration(v *bool) *GroupCreate { return _c } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (_c *GroupCreate) SetAllowBatchImageGeneration(v bool) *GroupCreate { + _c.mutation.SetAllowBatchImageGeneration(v) + return _c +} + +// SetNillableAllowBatchImageGeneration sets the "allow_batch_image_generation" field if the given value is not nil. +func (_c *GroupCreate) SetNillableAllowBatchImageGeneration(v *bool) *GroupCreate { + if v != nil { + _c.SetAllowBatchImageGeneration(*v) + } + return _c +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (_c *GroupCreate) SetImageRateIndependent(v bool) *GroupCreate { _c.mutation.SetImageRateIndependent(v) @@ -357,6 +371,104 @@ func (_c *GroupCreate) SetNillableImagePrice4k(v *float64) *GroupCreate { return _c } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (_c *GroupCreate) SetBatchImageDiscountMultiplier(v float64) *GroupCreate { + _c.mutation.SetBatchImageDiscountMultiplier(v) + return _c +} + +// SetNillableBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field if the given value is not nil. +func (_c *GroupCreate) SetNillableBatchImageDiscountMultiplier(v *float64) *GroupCreate { + if v != nil { + _c.SetBatchImageDiscountMultiplier(*v) + } + return _c +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (_c *GroupCreate) SetBatchImageHoldMultiplier(v float64) *GroupCreate { + _c.mutation.SetBatchImageHoldMultiplier(v) + return _c +} + +// SetNillableBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field if the given value is not nil. +func (_c *GroupCreate) SetNillableBatchImageHoldMultiplier(v *float64) *GroupCreate { + if v != nil { + _c.SetBatchImageHoldMultiplier(*v) + } + return _c +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (_c *GroupCreate) SetVideoRateIndependent(v bool) *GroupCreate { + _c.mutation.SetVideoRateIndependent(v) + return _c +} + +// SetNillableVideoRateIndependent sets the "video_rate_independent" field if the given value is not nil. +func (_c *GroupCreate) SetNillableVideoRateIndependent(v *bool) *GroupCreate { + if v != nil { + _c.SetVideoRateIndependent(*v) + } + return _c +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (_c *GroupCreate) SetVideoRateMultiplier(v float64) *GroupCreate { + _c.mutation.SetVideoRateMultiplier(v) + return _c +} + +// SetNillableVideoRateMultiplier sets the "video_rate_multiplier" field if the given value is not nil. +func (_c *GroupCreate) SetNillableVideoRateMultiplier(v *float64) *GroupCreate { + if v != nil { + _c.SetVideoRateMultiplier(*v) + } + return _c +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (_c *GroupCreate) SetVideoPrice480p(v float64) *GroupCreate { + _c.mutation.SetVideoPrice480p(v) + return _c +} + +// SetNillableVideoPrice480p sets the "video_price_480p" field if the given value is not nil. +func (_c *GroupCreate) SetNillableVideoPrice480p(v *float64) *GroupCreate { + if v != nil { + _c.SetVideoPrice480p(*v) + } + return _c +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (_c *GroupCreate) SetVideoPrice720p(v float64) *GroupCreate { + _c.mutation.SetVideoPrice720p(v) + return _c +} + +// SetNillableVideoPrice720p sets the "video_price_720p" field if the given value is not nil. +func (_c *GroupCreate) SetNillableVideoPrice720p(v *float64) *GroupCreate { + if v != nil { + _c.SetVideoPrice720p(*v) + } + return _c +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (_c *GroupCreate) SetVideoPrice1080p(v float64) *GroupCreate { + _c.mutation.SetVideoPrice1080p(v) + return _c +} + +// SetNillableVideoPrice1080p sets the "video_price_1080p" field if the given value is not nil. +func (_c *GroupCreate) SetNillableVideoPrice1080p(v *float64) *GroupCreate { + if v != nil { + _c.SetVideoPrice1080p(*v) + } + return _c +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_c *GroupCreate) SetClaudeCodeOnly(v bool) *GroupCreate { _c.mutation.SetClaudeCodeOnly(v) @@ -736,6 +848,10 @@ func (_c *GroupCreate) defaults() error { v := group.DefaultAllowImageGeneration _c.mutation.SetAllowImageGeneration(v) } + if _, ok := _c.mutation.AllowBatchImageGeneration(); !ok { + v := group.DefaultAllowBatchImageGeneration + _c.mutation.SetAllowBatchImageGeneration(v) + } if _, ok := _c.mutation.ImageRateIndependent(); !ok { v := group.DefaultImageRateIndependent _c.mutation.SetImageRateIndependent(v) @@ -744,6 +860,22 @@ func (_c *GroupCreate) defaults() error { v := group.DefaultImageRateMultiplier _c.mutation.SetImageRateMultiplier(v) } + if _, ok := _c.mutation.BatchImageDiscountMultiplier(); !ok { + v := group.DefaultBatchImageDiscountMultiplier + _c.mutation.SetBatchImageDiscountMultiplier(v) + } + if _, ok := _c.mutation.BatchImageHoldMultiplier(); !ok { + v := group.DefaultBatchImageHoldMultiplier + _c.mutation.SetBatchImageHoldMultiplier(v) + } + if _, ok := _c.mutation.VideoRateIndependent(); !ok { + v := group.DefaultVideoRateIndependent + _c.mutation.SetVideoRateIndependent(v) + } + if _, ok := _c.mutation.VideoRateMultiplier(); !ok { + v := group.DefaultVideoRateMultiplier + _c.mutation.SetVideoRateMultiplier(v) + } if _, ok := _c.mutation.ClaudeCodeOnly(); !ok { v := group.DefaultClaudeCodeOnly _c.mutation.SetClaudeCodeOnly(v) @@ -869,12 +1001,27 @@ func (_c *GroupCreate) check() error { if _, ok := _c.mutation.AllowImageGeneration(); !ok { return &ValidationError{Name: "allow_image_generation", err: errors.New(`ent: missing required field "Group.allow_image_generation"`)} } + if _, ok := _c.mutation.AllowBatchImageGeneration(); !ok { + return &ValidationError{Name: "allow_batch_image_generation", err: errors.New(`ent: missing required field "Group.allow_batch_image_generation"`)} + } if _, ok := _c.mutation.ImageRateIndependent(); !ok { return &ValidationError{Name: "image_rate_independent", err: errors.New(`ent: missing required field "Group.image_rate_independent"`)} } if _, ok := _c.mutation.ImageRateMultiplier(); !ok { return &ValidationError{Name: "image_rate_multiplier", err: errors.New(`ent: missing required field "Group.image_rate_multiplier"`)} } + if _, ok := _c.mutation.BatchImageDiscountMultiplier(); !ok { + return &ValidationError{Name: "batch_image_discount_multiplier", err: errors.New(`ent: missing required field "Group.batch_image_discount_multiplier"`)} + } + if _, ok := _c.mutation.BatchImageHoldMultiplier(); !ok { + return &ValidationError{Name: "batch_image_hold_multiplier", err: errors.New(`ent: missing required field "Group.batch_image_hold_multiplier"`)} + } + if _, ok := _c.mutation.VideoRateIndependent(); !ok { + return &ValidationError{Name: "video_rate_independent", err: errors.New(`ent: missing required field "Group.video_rate_independent"`)} + } + if _, ok := _c.mutation.VideoRateMultiplier(); !ok { + return &ValidationError{Name: "video_rate_multiplier", err: errors.New(`ent: missing required field "Group.video_rate_multiplier"`)} + } if _, ok := _c.mutation.ClaudeCodeOnly(); !ok { return &ValidationError{Name: "claude_code_only", err: errors.New(`ent: missing required field "Group.claude_code_only"`)} } @@ -1019,6 +1166,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) { _spec.SetField(group.FieldAllowImageGeneration, field.TypeBool, value) _node.AllowImageGeneration = value } + if value, ok := _c.mutation.AllowBatchImageGeneration(); ok { + _spec.SetField(group.FieldAllowBatchImageGeneration, field.TypeBool, value) + _node.AllowBatchImageGeneration = value + } if value, ok := _c.mutation.ImageRateIndependent(); ok { _spec.SetField(group.FieldImageRateIndependent, field.TypeBool, value) _node.ImageRateIndependent = value @@ -1039,6 +1190,34 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) { _spec.SetField(group.FieldImagePrice4k, field.TypeFloat64, value) _node.ImagePrice4k = &value } + if value, ok := _c.mutation.BatchImageDiscountMultiplier(); ok { + _spec.SetField(group.FieldBatchImageDiscountMultiplier, field.TypeFloat64, value) + _node.BatchImageDiscountMultiplier = value + } + if value, ok := _c.mutation.BatchImageHoldMultiplier(); ok { + _spec.SetField(group.FieldBatchImageHoldMultiplier, field.TypeFloat64, value) + _node.BatchImageHoldMultiplier = value + } + if value, ok := _c.mutation.VideoRateIndependent(); ok { + _spec.SetField(group.FieldVideoRateIndependent, field.TypeBool, value) + _node.VideoRateIndependent = value + } + if value, ok := _c.mutation.VideoRateMultiplier(); ok { + _spec.SetField(group.FieldVideoRateMultiplier, field.TypeFloat64, value) + _node.VideoRateMultiplier = value + } + if value, ok := _c.mutation.VideoPrice480p(); ok { + _spec.SetField(group.FieldVideoPrice480p, field.TypeFloat64, value) + _node.VideoPrice480p = &value + } + if value, ok := _c.mutation.VideoPrice720p(); ok { + _spec.SetField(group.FieldVideoPrice720p, field.TypeFloat64, value) + _node.VideoPrice720p = &value + } + if value, ok := _c.mutation.VideoPrice1080p(); ok { + _spec.SetField(group.FieldVideoPrice1080p, field.TypeFloat64, value) + _node.VideoPrice1080p = &value + } if value, ok := _c.mutation.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) _node.ClaudeCodeOnly = value @@ -1537,6 +1716,18 @@ func (u *GroupUpsert) UpdateAllowImageGeneration() *GroupUpsert { return u } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (u *GroupUpsert) SetAllowBatchImageGeneration(v bool) *GroupUpsert { + u.Set(group.FieldAllowBatchImageGeneration, v) + return u +} + +// UpdateAllowBatchImageGeneration sets the "allow_batch_image_generation" field to the value that was provided on create. +func (u *GroupUpsert) UpdateAllowBatchImageGeneration() *GroupUpsert { + u.SetExcluded(group.FieldAllowBatchImageGeneration) + return u +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (u *GroupUpsert) SetImageRateIndependent(v bool) *GroupUpsert { u.Set(group.FieldImageRateIndependent, v) @@ -1639,6 +1830,144 @@ func (u *GroupUpsert) ClearImagePrice4k() *GroupUpsert { return u } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (u *GroupUpsert) SetBatchImageDiscountMultiplier(v float64) *GroupUpsert { + u.Set(group.FieldBatchImageDiscountMultiplier, v) + return u +} + +// UpdateBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field to the value that was provided on create. +func (u *GroupUpsert) UpdateBatchImageDiscountMultiplier() *GroupUpsert { + u.SetExcluded(group.FieldBatchImageDiscountMultiplier) + return u +} + +// AddBatchImageDiscountMultiplier adds v to the "batch_image_discount_multiplier" field. +func (u *GroupUpsert) AddBatchImageDiscountMultiplier(v float64) *GroupUpsert { + u.Add(group.FieldBatchImageDiscountMultiplier, v) + return u +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (u *GroupUpsert) SetBatchImageHoldMultiplier(v float64) *GroupUpsert { + u.Set(group.FieldBatchImageHoldMultiplier, v) + return u +} + +// UpdateBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field to the value that was provided on create. +func (u *GroupUpsert) UpdateBatchImageHoldMultiplier() *GroupUpsert { + u.SetExcluded(group.FieldBatchImageHoldMultiplier) + return u +} + +// AddBatchImageHoldMultiplier adds v to the "batch_image_hold_multiplier" field. +func (u *GroupUpsert) AddBatchImageHoldMultiplier(v float64) *GroupUpsert { + u.Add(group.FieldBatchImageHoldMultiplier, v) + return u +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (u *GroupUpsert) SetVideoRateIndependent(v bool) *GroupUpsert { + u.Set(group.FieldVideoRateIndependent, v) + return u +} + +// UpdateVideoRateIndependent sets the "video_rate_independent" field to the value that was provided on create. +func (u *GroupUpsert) UpdateVideoRateIndependent() *GroupUpsert { + u.SetExcluded(group.FieldVideoRateIndependent) + return u +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (u *GroupUpsert) SetVideoRateMultiplier(v float64) *GroupUpsert { + u.Set(group.FieldVideoRateMultiplier, v) + return u +} + +// UpdateVideoRateMultiplier sets the "video_rate_multiplier" field to the value that was provided on create. +func (u *GroupUpsert) UpdateVideoRateMultiplier() *GroupUpsert { + u.SetExcluded(group.FieldVideoRateMultiplier) + return u +} + +// AddVideoRateMultiplier adds v to the "video_rate_multiplier" field. +func (u *GroupUpsert) AddVideoRateMultiplier(v float64) *GroupUpsert { + u.Add(group.FieldVideoRateMultiplier, v) + return u +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (u *GroupUpsert) SetVideoPrice480p(v float64) *GroupUpsert { + u.Set(group.FieldVideoPrice480p, v) + return u +} + +// UpdateVideoPrice480p sets the "video_price_480p" field to the value that was provided on create. +func (u *GroupUpsert) UpdateVideoPrice480p() *GroupUpsert { + u.SetExcluded(group.FieldVideoPrice480p) + return u +} + +// AddVideoPrice480p adds v to the "video_price_480p" field. +func (u *GroupUpsert) AddVideoPrice480p(v float64) *GroupUpsert { + u.Add(group.FieldVideoPrice480p, v) + return u +} + +// ClearVideoPrice480p clears the value of the "video_price_480p" field. +func (u *GroupUpsert) ClearVideoPrice480p() *GroupUpsert { + u.SetNull(group.FieldVideoPrice480p) + return u +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (u *GroupUpsert) SetVideoPrice720p(v float64) *GroupUpsert { + u.Set(group.FieldVideoPrice720p, v) + return u +} + +// UpdateVideoPrice720p sets the "video_price_720p" field to the value that was provided on create. +func (u *GroupUpsert) UpdateVideoPrice720p() *GroupUpsert { + u.SetExcluded(group.FieldVideoPrice720p) + return u +} + +// AddVideoPrice720p adds v to the "video_price_720p" field. +func (u *GroupUpsert) AddVideoPrice720p(v float64) *GroupUpsert { + u.Add(group.FieldVideoPrice720p, v) + return u +} + +// ClearVideoPrice720p clears the value of the "video_price_720p" field. +func (u *GroupUpsert) ClearVideoPrice720p() *GroupUpsert { + u.SetNull(group.FieldVideoPrice720p) + return u +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (u *GroupUpsert) SetVideoPrice1080p(v float64) *GroupUpsert { + u.Set(group.FieldVideoPrice1080p, v) + return u +} + +// UpdateVideoPrice1080p sets the "video_price_1080p" field to the value that was provided on create. +func (u *GroupUpsert) UpdateVideoPrice1080p() *GroupUpsert { + u.SetExcluded(group.FieldVideoPrice1080p) + return u +} + +// AddVideoPrice1080p adds v to the "video_price_1080p" field. +func (u *GroupUpsert) AddVideoPrice1080p(v float64) *GroupUpsert { + u.Add(group.FieldVideoPrice1080p, v) + return u +} + +// ClearVideoPrice1080p clears the value of the "video_price_1080p" field. +func (u *GroupUpsert) ClearVideoPrice1080p() *GroupUpsert { + u.SetNull(group.FieldVideoPrice1080p) + return u +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsert) SetClaudeCodeOnly(v bool) *GroupUpsert { u.Set(group.FieldClaudeCodeOnly, v) @@ -2235,6 +2564,20 @@ func (u *GroupUpsertOne) UpdateAllowImageGeneration() *GroupUpsertOne { }) } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (u *GroupUpsertOne) SetAllowBatchImageGeneration(v bool) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetAllowBatchImageGeneration(v) + }) +} + +// UpdateAllowBatchImageGeneration sets the "allow_batch_image_generation" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateAllowBatchImageGeneration() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateAllowBatchImageGeneration() + }) +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (u *GroupUpsertOne) SetImageRateIndependent(v bool) *GroupUpsertOne { return u.Update(func(s *GroupUpsert) { @@ -2354,6 +2697,167 @@ func (u *GroupUpsertOne) ClearImagePrice4k() *GroupUpsertOne { }) } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (u *GroupUpsertOne) SetBatchImageDiscountMultiplier(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetBatchImageDiscountMultiplier(v) + }) +} + +// AddBatchImageDiscountMultiplier adds v to the "batch_image_discount_multiplier" field. +func (u *GroupUpsertOne) AddBatchImageDiscountMultiplier(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddBatchImageDiscountMultiplier(v) + }) +} + +// UpdateBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateBatchImageDiscountMultiplier() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateBatchImageDiscountMultiplier() + }) +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (u *GroupUpsertOne) SetBatchImageHoldMultiplier(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetBatchImageHoldMultiplier(v) + }) +} + +// AddBatchImageHoldMultiplier adds v to the "batch_image_hold_multiplier" field. +func (u *GroupUpsertOne) AddBatchImageHoldMultiplier(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddBatchImageHoldMultiplier(v) + }) +} + +// UpdateBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateBatchImageHoldMultiplier() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateBatchImageHoldMultiplier() + }) +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (u *GroupUpsertOne) SetVideoRateIndependent(v bool) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetVideoRateIndependent(v) + }) +} + +// UpdateVideoRateIndependent sets the "video_rate_independent" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateVideoRateIndependent() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoRateIndependent() + }) +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (u *GroupUpsertOne) SetVideoRateMultiplier(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetVideoRateMultiplier(v) + }) +} + +// AddVideoRateMultiplier adds v to the "video_rate_multiplier" field. +func (u *GroupUpsertOne) AddVideoRateMultiplier(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddVideoRateMultiplier(v) + }) +} + +// UpdateVideoRateMultiplier sets the "video_rate_multiplier" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateVideoRateMultiplier() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoRateMultiplier() + }) +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (u *GroupUpsertOne) SetVideoPrice480p(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetVideoPrice480p(v) + }) +} + +// AddVideoPrice480p adds v to the "video_price_480p" field. +func (u *GroupUpsertOne) AddVideoPrice480p(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddVideoPrice480p(v) + }) +} + +// UpdateVideoPrice480p sets the "video_price_480p" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateVideoPrice480p() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoPrice480p() + }) +} + +// ClearVideoPrice480p clears the value of the "video_price_480p" field. +func (u *GroupUpsertOne) ClearVideoPrice480p() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.ClearVideoPrice480p() + }) +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (u *GroupUpsertOne) SetVideoPrice720p(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetVideoPrice720p(v) + }) +} + +// AddVideoPrice720p adds v to the "video_price_720p" field. +func (u *GroupUpsertOne) AddVideoPrice720p(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddVideoPrice720p(v) + }) +} + +// UpdateVideoPrice720p sets the "video_price_720p" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateVideoPrice720p() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoPrice720p() + }) +} + +// ClearVideoPrice720p clears the value of the "video_price_720p" field. +func (u *GroupUpsertOne) ClearVideoPrice720p() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.ClearVideoPrice720p() + }) +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (u *GroupUpsertOne) SetVideoPrice1080p(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetVideoPrice1080p(v) + }) +} + +// AddVideoPrice1080p adds v to the "video_price_1080p" field. +func (u *GroupUpsertOne) AddVideoPrice1080p(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddVideoPrice1080p(v) + }) +} + +// UpdateVideoPrice1080p sets the "video_price_1080p" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateVideoPrice1080p() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoPrice1080p() + }) +} + +// ClearVideoPrice1080p clears the value of the "video_price_1080p" field. +func (u *GroupUpsertOne) ClearVideoPrice1080p() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.ClearVideoPrice1080p() + }) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsertOne) SetClaudeCodeOnly(v bool) *GroupUpsertOne { return u.Update(func(s *GroupUpsert) { @@ -3153,6 +3657,20 @@ func (u *GroupUpsertBulk) UpdateAllowImageGeneration() *GroupUpsertBulk { }) } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (u *GroupUpsertBulk) SetAllowBatchImageGeneration(v bool) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetAllowBatchImageGeneration(v) + }) +} + +// UpdateAllowBatchImageGeneration sets the "allow_batch_image_generation" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateAllowBatchImageGeneration() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateAllowBatchImageGeneration() + }) +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (u *GroupUpsertBulk) SetImageRateIndependent(v bool) *GroupUpsertBulk { return u.Update(func(s *GroupUpsert) { @@ -3272,6 +3790,167 @@ func (u *GroupUpsertBulk) ClearImagePrice4k() *GroupUpsertBulk { }) } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (u *GroupUpsertBulk) SetBatchImageDiscountMultiplier(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetBatchImageDiscountMultiplier(v) + }) +} + +// AddBatchImageDiscountMultiplier adds v to the "batch_image_discount_multiplier" field. +func (u *GroupUpsertBulk) AddBatchImageDiscountMultiplier(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddBatchImageDiscountMultiplier(v) + }) +} + +// UpdateBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateBatchImageDiscountMultiplier() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateBatchImageDiscountMultiplier() + }) +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (u *GroupUpsertBulk) SetBatchImageHoldMultiplier(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetBatchImageHoldMultiplier(v) + }) +} + +// AddBatchImageHoldMultiplier adds v to the "batch_image_hold_multiplier" field. +func (u *GroupUpsertBulk) AddBatchImageHoldMultiplier(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddBatchImageHoldMultiplier(v) + }) +} + +// UpdateBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateBatchImageHoldMultiplier() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateBatchImageHoldMultiplier() + }) +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (u *GroupUpsertBulk) SetVideoRateIndependent(v bool) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetVideoRateIndependent(v) + }) +} + +// UpdateVideoRateIndependent sets the "video_rate_independent" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateVideoRateIndependent() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoRateIndependent() + }) +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (u *GroupUpsertBulk) SetVideoRateMultiplier(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetVideoRateMultiplier(v) + }) +} + +// AddVideoRateMultiplier adds v to the "video_rate_multiplier" field. +func (u *GroupUpsertBulk) AddVideoRateMultiplier(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddVideoRateMultiplier(v) + }) +} + +// UpdateVideoRateMultiplier sets the "video_rate_multiplier" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateVideoRateMultiplier() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoRateMultiplier() + }) +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (u *GroupUpsertBulk) SetVideoPrice480p(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetVideoPrice480p(v) + }) +} + +// AddVideoPrice480p adds v to the "video_price_480p" field. +func (u *GroupUpsertBulk) AddVideoPrice480p(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddVideoPrice480p(v) + }) +} + +// UpdateVideoPrice480p sets the "video_price_480p" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateVideoPrice480p() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoPrice480p() + }) +} + +// ClearVideoPrice480p clears the value of the "video_price_480p" field. +func (u *GroupUpsertBulk) ClearVideoPrice480p() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.ClearVideoPrice480p() + }) +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (u *GroupUpsertBulk) SetVideoPrice720p(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetVideoPrice720p(v) + }) +} + +// AddVideoPrice720p adds v to the "video_price_720p" field. +func (u *GroupUpsertBulk) AddVideoPrice720p(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddVideoPrice720p(v) + }) +} + +// UpdateVideoPrice720p sets the "video_price_720p" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateVideoPrice720p() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoPrice720p() + }) +} + +// ClearVideoPrice720p clears the value of the "video_price_720p" field. +func (u *GroupUpsertBulk) ClearVideoPrice720p() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.ClearVideoPrice720p() + }) +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (u *GroupUpsertBulk) SetVideoPrice1080p(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetVideoPrice1080p(v) + }) +} + +// AddVideoPrice1080p adds v to the "video_price_1080p" field. +func (u *GroupUpsertBulk) AddVideoPrice1080p(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddVideoPrice1080p(v) + }) +} + +// UpdateVideoPrice1080p sets the "video_price_1080p" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateVideoPrice1080p() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateVideoPrice1080p() + }) +} + +// ClearVideoPrice1080p clears the value of the "video_price_1080p" field. +func (u *GroupUpsertBulk) ClearVideoPrice1080p() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.ClearVideoPrice1080p() + }) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsertBulk) SetClaudeCodeOnly(v bool) *GroupUpsertBulk { return u.Update(func(s *GroupUpsert) { diff --git a/backend/ent/group_update.go b/backend/ent/group_update.go index 55555f323c..3bb18d3e1a 100644 --- a/backend/ent/group_update.go +++ b/backend/ent/group_update.go @@ -352,6 +352,20 @@ func (_u *GroupUpdate) SetNillableAllowImageGeneration(v *bool) *GroupUpdate { return _u } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (_u *GroupUpdate) SetAllowBatchImageGeneration(v bool) *GroupUpdate { + _u.mutation.SetAllowBatchImageGeneration(v) + return _u +} + +// SetNillableAllowBatchImageGeneration sets the "allow_batch_image_generation" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableAllowBatchImageGeneration(v *bool) *GroupUpdate { + if v != nil { + _u.SetAllowBatchImageGeneration(*v) + } + return _u +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (_u *GroupUpdate) SetImageRateIndependent(v bool) *GroupUpdate { _u.mutation.SetImageRateIndependent(v) @@ -468,6 +482,164 @@ func (_u *GroupUpdate) ClearImagePrice4k() *GroupUpdate { return _u } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (_u *GroupUpdate) SetBatchImageDiscountMultiplier(v float64) *GroupUpdate { + _u.mutation.ResetBatchImageDiscountMultiplier() + _u.mutation.SetBatchImageDiscountMultiplier(v) + return _u +} + +// SetNillableBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableBatchImageDiscountMultiplier(v *float64) *GroupUpdate { + if v != nil { + _u.SetBatchImageDiscountMultiplier(*v) + } + return _u +} + +// AddBatchImageDiscountMultiplier adds value to the "batch_image_discount_multiplier" field. +func (_u *GroupUpdate) AddBatchImageDiscountMultiplier(v float64) *GroupUpdate { + _u.mutation.AddBatchImageDiscountMultiplier(v) + return _u +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (_u *GroupUpdate) SetBatchImageHoldMultiplier(v float64) *GroupUpdate { + _u.mutation.ResetBatchImageHoldMultiplier() + _u.mutation.SetBatchImageHoldMultiplier(v) + return _u +} + +// SetNillableBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableBatchImageHoldMultiplier(v *float64) *GroupUpdate { + if v != nil { + _u.SetBatchImageHoldMultiplier(*v) + } + return _u +} + +// AddBatchImageHoldMultiplier adds value to the "batch_image_hold_multiplier" field. +func (_u *GroupUpdate) AddBatchImageHoldMultiplier(v float64) *GroupUpdate { + _u.mutation.AddBatchImageHoldMultiplier(v) + return _u +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (_u *GroupUpdate) SetVideoRateIndependent(v bool) *GroupUpdate { + _u.mutation.SetVideoRateIndependent(v) + return _u +} + +// SetNillableVideoRateIndependent sets the "video_rate_independent" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableVideoRateIndependent(v *bool) *GroupUpdate { + if v != nil { + _u.SetVideoRateIndependent(*v) + } + return _u +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (_u *GroupUpdate) SetVideoRateMultiplier(v float64) *GroupUpdate { + _u.mutation.ResetVideoRateMultiplier() + _u.mutation.SetVideoRateMultiplier(v) + return _u +} + +// SetNillableVideoRateMultiplier sets the "video_rate_multiplier" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableVideoRateMultiplier(v *float64) *GroupUpdate { + if v != nil { + _u.SetVideoRateMultiplier(*v) + } + return _u +} + +// AddVideoRateMultiplier adds value to the "video_rate_multiplier" field. +func (_u *GroupUpdate) AddVideoRateMultiplier(v float64) *GroupUpdate { + _u.mutation.AddVideoRateMultiplier(v) + return _u +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (_u *GroupUpdate) SetVideoPrice480p(v float64) *GroupUpdate { + _u.mutation.ResetVideoPrice480p() + _u.mutation.SetVideoPrice480p(v) + return _u +} + +// SetNillableVideoPrice480p sets the "video_price_480p" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableVideoPrice480p(v *float64) *GroupUpdate { + if v != nil { + _u.SetVideoPrice480p(*v) + } + return _u +} + +// AddVideoPrice480p adds value to the "video_price_480p" field. +func (_u *GroupUpdate) AddVideoPrice480p(v float64) *GroupUpdate { + _u.mutation.AddVideoPrice480p(v) + return _u +} + +// ClearVideoPrice480p clears the value of the "video_price_480p" field. +func (_u *GroupUpdate) ClearVideoPrice480p() *GroupUpdate { + _u.mutation.ClearVideoPrice480p() + return _u +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (_u *GroupUpdate) SetVideoPrice720p(v float64) *GroupUpdate { + _u.mutation.ResetVideoPrice720p() + _u.mutation.SetVideoPrice720p(v) + return _u +} + +// SetNillableVideoPrice720p sets the "video_price_720p" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableVideoPrice720p(v *float64) *GroupUpdate { + if v != nil { + _u.SetVideoPrice720p(*v) + } + return _u +} + +// AddVideoPrice720p adds value to the "video_price_720p" field. +func (_u *GroupUpdate) AddVideoPrice720p(v float64) *GroupUpdate { + _u.mutation.AddVideoPrice720p(v) + return _u +} + +// ClearVideoPrice720p clears the value of the "video_price_720p" field. +func (_u *GroupUpdate) ClearVideoPrice720p() *GroupUpdate { + _u.mutation.ClearVideoPrice720p() + return _u +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (_u *GroupUpdate) SetVideoPrice1080p(v float64) *GroupUpdate { + _u.mutation.ResetVideoPrice1080p() + _u.mutation.SetVideoPrice1080p(v) + return _u +} + +// SetNillableVideoPrice1080p sets the "video_price_1080p" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableVideoPrice1080p(v *float64) *GroupUpdate { + if v != nil { + _u.SetVideoPrice1080p(*v) + } + return _u +} + +// AddVideoPrice1080p adds value to the "video_price_1080p" field. +func (_u *GroupUpdate) AddVideoPrice1080p(v float64) *GroupUpdate { + _u.mutation.AddVideoPrice1080p(v) + return _u +} + +// ClearVideoPrice1080p clears the value of the "video_price_1080p" field. +func (_u *GroupUpdate) ClearVideoPrice1080p() *GroupUpdate { + _u.mutation.ClearVideoPrice1080p() + return _u +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_u *GroupUpdate) SetClaudeCodeOnly(v bool) *GroupUpdate { _u.mutation.SetClaudeCodeOnly(v) @@ -1116,6 +1288,9 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) { if value, ok := _u.mutation.AllowImageGeneration(); ok { _spec.SetField(group.FieldAllowImageGeneration, field.TypeBool, value) } + if value, ok := _u.mutation.AllowBatchImageGeneration(); ok { + _spec.SetField(group.FieldAllowBatchImageGeneration, field.TypeBool, value) + } if value, ok := _u.mutation.ImageRateIndependent(); ok { _spec.SetField(group.FieldImageRateIndependent, field.TypeBool, value) } @@ -1152,6 +1327,54 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) { if _u.mutation.ImagePrice4kCleared() { _spec.ClearField(group.FieldImagePrice4k, field.TypeFloat64) } + if value, ok := _u.mutation.BatchImageDiscountMultiplier(); ok { + _spec.SetField(group.FieldBatchImageDiscountMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedBatchImageDiscountMultiplier(); ok { + _spec.AddField(group.FieldBatchImageDiscountMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.BatchImageHoldMultiplier(); ok { + _spec.SetField(group.FieldBatchImageHoldMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedBatchImageHoldMultiplier(); ok { + _spec.AddField(group.FieldBatchImageHoldMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.VideoRateIndependent(); ok { + _spec.SetField(group.FieldVideoRateIndependent, field.TypeBool, value) + } + if value, ok := _u.mutation.VideoRateMultiplier(); ok { + _spec.SetField(group.FieldVideoRateMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoRateMultiplier(); ok { + _spec.AddField(group.FieldVideoRateMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.VideoPrice480p(); ok { + _spec.SetField(group.FieldVideoPrice480p, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoPrice480p(); ok { + _spec.AddField(group.FieldVideoPrice480p, field.TypeFloat64, value) + } + if _u.mutation.VideoPrice480pCleared() { + _spec.ClearField(group.FieldVideoPrice480p, field.TypeFloat64) + } + if value, ok := _u.mutation.VideoPrice720p(); ok { + _spec.SetField(group.FieldVideoPrice720p, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoPrice720p(); ok { + _spec.AddField(group.FieldVideoPrice720p, field.TypeFloat64, value) + } + if _u.mutation.VideoPrice720pCleared() { + _spec.ClearField(group.FieldVideoPrice720p, field.TypeFloat64) + } + if value, ok := _u.mutation.VideoPrice1080p(); ok { + _spec.SetField(group.FieldVideoPrice1080p, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoPrice1080p(); ok { + _spec.AddField(group.FieldVideoPrice1080p, field.TypeFloat64, value) + } + if _u.mutation.VideoPrice1080pCleared() { + _spec.ClearField(group.FieldVideoPrice1080p, field.TypeFloat64) + } if value, ok := _u.mutation.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) } @@ -1853,6 +2076,20 @@ func (_u *GroupUpdateOne) SetNillableAllowImageGeneration(v *bool) *GroupUpdateO return _u } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (_u *GroupUpdateOne) SetAllowBatchImageGeneration(v bool) *GroupUpdateOne { + _u.mutation.SetAllowBatchImageGeneration(v) + return _u +} + +// SetNillableAllowBatchImageGeneration sets the "allow_batch_image_generation" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableAllowBatchImageGeneration(v *bool) *GroupUpdateOne { + if v != nil { + _u.SetAllowBatchImageGeneration(*v) + } + return _u +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (_u *GroupUpdateOne) SetImageRateIndependent(v bool) *GroupUpdateOne { _u.mutation.SetImageRateIndependent(v) @@ -1969,6 +2206,164 @@ func (_u *GroupUpdateOne) ClearImagePrice4k() *GroupUpdateOne { return _u } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (_u *GroupUpdateOne) SetBatchImageDiscountMultiplier(v float64) *GroupUpdateOne { + _u.mutation.ResetBatchImageDiscountMultiplier() + _u.mutation.SetBatchImageDiscountMultiplier(v) + return _u +} + +// SetNillableBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableBatchImageDiscountMultiplier(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetBatchImageDiscountMultiplier(*v) + } + return _u +} + +// AddBatchImageDiscountMultiplier adds value to the "batch_image_discount_multiplier" field. +func (_u *GroupUpdateOne) AddBatchImageDiscountMultiplier(v float64) *GroupUpdateOne { + _u.mutation.AddBatchImageDiscountMultiplier(v) + return _u +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (_u *GroupUpdateOne) SetBatchImageHoldMultiplier(v float64) *GroupUpdateOne { + _u.mutation.ResetBatchImageHoldMultiplier() + _u.mutation.SetBatchImageHoldMultiplier(v) + return _u +} + +// SetNillableBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableBatchImageHoldMultiplier(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetBatchImageHoldMultiplier(*v) + } + return _u +} + +// AddBatchImageHoldMultiplier adds value to the "batch_image_hold_multiplier" field. +func (_u *GroupUpdateOne) AddBatchImageHoldMultiplier(v float64) *GroupUpdateOne { + _u.mutation.AddBatchImageHoldMultiplier(v) + return _u +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (_u *GroupUpdateOne) SetVideoRateIndependent(v bool) *GroupUpdateOne { + _u.mutation.SetVideoRateIndependent(v) + return _u +} + +// SetNillableVideoRateIndependent sets the "video_rate_independent" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableVideoRateIndependent(v *bool) *GroupUpdateOne { + if v != nil { + _u.SetVideoRateIndependent(*v) + } + return _u +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (_u *GroupUpdateOne) SetVideoRateMultiplier(v float64) *GroupUpdateOne { + _u.mutation.ResetVideoRateMultiplier() + _u.mutation.SetVideoRateMultiplier(v) + return _u +} + +// SetNillableVideoRateMultiplier sets the "video_rate_multiplier" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableVideoRateMultiplier(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetVideoRateMultiplier(*v) + } + return _u +} + +// AddVideoRateMultiplier adds value to the "video_rate_multiplier" field. +func (_u *GroupUpdateOne) AddVideoRateMultiplier(v float64) *GroupUpdateOne { + _u.mutation.AddVideoRateMultiplier(v) + return _u +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (_u *GroupUpdateOne) SetVideoPrice480p(v float64) *GroupUpdateOne { + _u.mutation.ResetVideoPrice480p() + _u.mutation.SetVideoPrice480p(v) + return _u +} + +// SetNillableVideoPrice480p sets the "video_price_480p" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableVideoPrice480p(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetVideoPrice480p(*v) + } + return _u +} + +// AddVideoPrice480p adds value to the "video_price_480p" field. +func (_u *GroupUpdateOne) AddVideoPrice480p(v float64) *GroupUpdateOne { + _u.mutation.AddVideoPrice480p(v) + return _u +} + +// ClearVideoPrice480p clears the value of the "video_price_480p" field. +func (_u *GroupUpdateOne) ClearVideoPrice480p() *GroupUpdateOne { + _u.mutation.ClearVideoPrice480p() + return _u +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (_u *GroupUpdateOne) SetVideoPrice720p(v float64) *GroupUpdateOne { + _u.mutation.ResetVideoPrice720p() + _u.mutation.SetVideoPrice720p(v) + return _u +} + +// SetNillableVideoPrice720p sets the "video_price_720p" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableVideoPrice720p(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetVideoPrice720p(*v) + } + return _u +} + +// AddVideoPrice720p adds value to the "video_price_720p" field. +func (_u *GroupUpdateOne) AddVideoPrice720p(v float64) *GroupUpdateOne { + _u.mutation.AddVideoPrice720p(v) + return _u +} + +// ClearVideoPrice720p clears the value of the "video_price_720p" field. +func (_u *GroupUpdateOne) ClearVideoPrice720p() *GroupUpdateOne { + _u.mutation.ClearVideoPrice720p() + return _u +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (_u *GroupUpdateOne) SetVideoPrice1080p(v float64) *GroupUpdateOne { + _u.mutation.ResetVideoPrice1080p() + _u.mutation.SetVideoPrice1080p(v) + return _u +} + +// SetNillableVideoPrice1080p sets the "video_price_1080p" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableVideoPrice1080p(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetVideoPrice1080p(*v) + } + return _u +} + +// AddVideoPrice1080p adds value to the "video_price_1080p" field. +func (_u *GroupUpdateOne) AddVideoPrice1080p(v float64) *GroupUpdateOne { + _u.mutation.AddVideoPrice1080p(v) + return _u +} + +// ClearVideoPrice1080p clears the value of the "video_price_1080p" field. +func (_u *GroupUpdateOne) ClearVideoPrice1080p() *GroupUpdateOne { + _u.mutation.ClearVideoPrice1080p() + return _u +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_u *GroupUpdateOne) SetClaudeCodeOnly(v bool) *GroupUpdateOne { _u.mutation.SetClaudeCodeOnly(v) @@ -2647,6 +3042,9 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error) if value, ok := _u.mutation.AllowImageGeneration(); ok { _spec.SetField(group.FieldAllowImageGeneration, field.TypeBool, value) } + if value, ok := _u.mutation.AllowBatchImageGeneration(); ok { + _spec.SetField(group.FieldAllowBatchImageGeneration, field.TypeBool, value) + } if value, ok := _u.mutation.ImageRateIndependent(); ok { _spec.SetField(group.FieldImageRateIndependent, field.TypeBool, value) } @@ -2683,6 +3081,54 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error) if _u.mutation.ImagePrice4kCleared() { _spec.ClearField(group.FieldImagePrice4k, field.TypeFloat64) } + if value, ok := _u.mutation.BatchImageDiscountMultiplier(); ok { + _spec.SetField(group.FieldBatchImageDiscountMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedBatchImageDiscountMultiplier(); ok { + _spec.AddField(group.FieldBatchImageDiscountMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.BatchImageHoldMultiplier(); ok { + _spec.SetField(group.FieldBatchImageHoldMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedBatchImageHoldMultiplier(); ok { + _spec.AddField(group.FieldBatchImageHoldMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.VideoRateIndependent(); ok { + _spec.SetField(group.FieldVideoRateIndependent, field.TypeBool, value) + } + if value, ok := _u.mutation.VideoRateMultiplier(); ok { + _spec.SetField(group.FieldVideoRateMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoRateMultiplier(); ok { + _spec.AddField(group.FieldVideoRateMultiplier, field.TypeFloat64, value) + } + if value, ok := _u.mutation.VideoPrice480p(); ok { + _spec.SetField(group.FieldVideoPrice480p, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoPrice480p(); ok { + _spec.AddField(group.FieldVideoPrice480p, field.TypeFloat64, value) + } + if _u.mutation.VideoPrice480pCleared() { + _spec.ClearField(group.FieldVideoPrice480p, field.TypeFloat64) + } + if value, ok := _u.mutation.VideoPrice720p(); ok { + _spec.SetField(group.FieldVideoPrice720p, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoPrice720p(); ok { + _spec.AddField(group.FieldVideoPrice720p, field.TypeFloat64, value) + } + if _u.mutation.VideoPrice720pCleared() { + _spec.ClearField(group.FieldVideoPrice720p, field.TypeFloat64) + } + if value, ok := _u.mutation.VideoPrice1080p(); ok { + _spec.SetField(group.FieldVideoPrice1080p, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedVideoPrice1080p(); ok { + _spec.AddField(group.FieldVideoPrice1080p, field.TypeFloat64, value) + } + if _u.mutation.VideoPrice1080pCleared() { + _spec.ClearField(group.FieldVideoPrice1080p, field.TypeFloat64) + } if value, ok := _u.mutation.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) } diff --git a/backend/ent/hook/hook.go b/backend/ent/hook/hook.go index 71bfd3b88e..181f2f99db 100644 --- a/backend/ent/hook/hook.go +++ b/backend/ent/hook/hook.go @@ -93,6 +93,42 @@ func (f AuthIdentityChannelFunc) Mutate(ctx context.Context, m ent.Mutation) (en return nil, fmt.Errorf("unexpected mutation type %T. expect *ent.AuthIdentityChannelMutation", m) } +// The BatchImageEventFunc type is an adapter to allow the use of ordinary +// function as BatchImageEvent mutator. +type BatchImageEventFunc func(context.Context, *ent.BatchImageEventMutation) (ent.Value, error) + +// Mutate calls f(ctx, m). +func (f BatchImageEventFunc) Mutate(ctx context.Context, m ent.Mutation) (ent.Value, error) { + if mv, ok := m.(*ent.BatchImageEventMutation); ok { + return f(ctx, mv) + } + return nil, fmt.Errorf("unexpected mutation type %T. expect *ent.BatchImageEventMutation", m) +} + +// The BatchImageItemFunc type is an adapter to allow the use of ordinary +// function as BatchImageItem mutator. +type BatchImageItemFunc func(context.Context, *ent.BatchImageItemMutation) (ent.Value, error) + +// Mutate calls f(ctx, m). +func (f BatchImageItemFunc) Mutate(ctx context.Context, m ent.Mutation) (ent.Value, error) { + if mv, ok := m.(*ent.BatchImageItemMutation); ok { + return f(ctx, mv) + } + return nil, fmt.Errorf("unexpected mutation type %T. expect *ent.BatchImageItemMutation", m) +} + +// The BatchImageJobFunc type is an adapter to allow the use of ordinary +// function as BatchImageJob mutator. +type BatchImageJobFunc func(context.Context, *ent.BatchImageJobMutation) (ent.Value, error) + +// Mutate calls f(ctx, m). +func (f BatchImageJobFunc) Mutate(ctx context.Context, m ent.Mutation) (ent.Value, error) { + if mv, ok := m.(*ent.BatchImageJobMutation); ok { + return f(ctx, mv) + } + return nil, fmt.Errorf("unexpected mutation type %T. expect *ent.BatchImageJobMutation", m) +} + // The ChannelMonitorFunc type is an adapter to allow the use of ordinary // function as ChannelMonitor mutator. type ChannelMonitorFunc func(context.Context, *ent.ChannelMonitorMutation) (ent.Value, error) diff --git a/backend/ent/intercept/intercept.go b/backend/ent/intercept/intercept.go index 5d86e25bd5..7aeb07692d 100644 --- a/backend/ent/intercept/intercept.go +++ b/backend/ent/intercept/intercept.go @@ -15,6 +15,9 @@ import ( "github.com/Wei-Shaw/sub2api/ent/apikey" "github.com/Wei-Shaw/sub2api/ent/authidentity" "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" "github.com/Wei-Shaw/sub2api/ent/channelmonitor" "github.com/Wei-Shaw/sub2api/ent/channelmonitordailyrollup" "github.com/Wei-Shaw/sub2api/ent/channelmonitorhistory" @@ -291,6 +294,87 @@ func (f TraverseAuthIdentityChannel) Traverse(ctx context.Context, q ent.Query) return fmt.Errorf("unexpected query type %T. expect *ent.AuthIdentityChannelQuery", q) } +// The BatchImageEventFunc type is an adapter to allow the use of ordinary function as a Querier. +type BatchImageEventFunc func(context.Context, *ent.BatchImageEventQuery) (ent.Value, error) + +// Query calls f(ctx, q). +func (f BatchImageEventFunc) Query(ctx context.Context, q ent.Query) (ent.Value, error) { + if q, ok := q.(*ent.BatchImageEventQuery); ok { + return f(ctx, q) + } + return nil, fmt.Errorf("unexpected query type %T. expect *ent.BatchImageEventQuery", q) +} + +// The TraverseBatchImageEvent type is an adapter to allow the use of ordinary function as Traverser. +type TraverseBatchImageEvent func(context.Context, *ent.BatchImageEventQuery) error + +// Intercept is a dummy implementation of Intercept that returns the next Querier in the pipeline. +func (f TraverseBatchImageEvent) Intercept(next ent.Querier) ent.Querier { + return next +} + +// Traverse calls f(ctx, q). +func (f TraverseBatchImageEvent) Traverse(ctx context.Context, q ent.Query) error { + if q, ok := q.(*ent.BatchImageEventQuery); ok { + return f(ctx, q) + } + return fmt.Errorf("unexpected query type %T. expect *ent.BatchImageEventQuery", q) +} + +// The BatchImageItemFunc type is an adapter to allow the use of ordinary function as a Querier. +type BatchImageItemFunc func(context.Context, *ent.BatchImageItemQuery) (ent.Value, error) + +// Query calls f(ctx, q). +func (f BatchImageItemFunc) Query(ctx context.Context, q ent.Query) (ent.Value, error) { + if q, ok := q.(*ent.BatchImageItemQuery); ok { + return f(ctx, q) + } + return nil, fmt.Errorf("unexpected query type %T. expect *ent.BatchImageItemQuery", q) +} + +// The TraverseBatchImageItem type is an adapter to allow the use of ordinary function as Traverser. +type TraverseBatchImageItem func(context.Context, *ent.BatchImageItemQuery) error + +// Intercept is a dummy implementation of Intercept that returns the next Querier in the pipeline. +func (f TraverseBatchImageItem) Intercept(next ent.Querier) ent.Querier { + return next +} + +// Traverse calls f(ctx, q). +func (f TraverseBatchImageItem) Traverse(ctx context.Context, q ent.Query) error { + if q, ok := q.(*ent.BatchImageItemQuery); ok { + return f(ctx, q) + } + return fmt.Errorf("unexpected query type %T. expect *ent.BatchImageItemQuery", q) +} + +// The BatchImageJobFunc type is an adapter to allow the use of ordinary function as a Querier. +type BatchImageJobFunc func(context.Context, *ent.BatchImageJobQuery) (ent.Value, error) + +// Query calls f(ctx, q). +func (f BatchImageJobFunc) Query(ctx context.Context, q ent.Query) (ent.Value, error) { + if q, ok := q.(*ent.BatchImageJobQuery); ok { + return f(ctx, q) + } + return nil, fmt.Errorf("unexpected query type %T. expect *ent.BatchImageJobQuery", q) +} + +// The TraverseBatchImageJob type is an adapter to allow the use of ordinary function as Traverser. +type TraverseBatchImageJob func(context.Context, *ent.BatchImageJobQuery) error + +// Intercept is a dummy implementation of Intercept that returns the next Querier in the pipeline. +func (f TraverseBatchImageJob) Intercept(next ent.Querier) ent.Querier { + return next +} + +// Traverse calls f(ctx, q). +func (f TraverseBatchImageJob) Traverse(ctx context.Context, q ent.Query) error { + if q, ok := q.(*ent.BatchImageJobQuery); ok { + return f(ctx, q) + } + return fmt.Errorf("unexpected query type %T. expect *ent.BatchImageJobQuery", q) +} + // The ChannelMonitorFunc type is an adapter to allow the use of ordinary function as a Querier. type ChannelMonitorFunc func(context.Context, *ent.ChannelMonitorQuery) (ent.Value, error) @@ -1064,6 +1148,12 @@ func NewQuery(q ent.Query) (Query, error) { return &query[*ent.AuthIdentityQuery, predicate.AuthIdentity, authidentity.OrderOption]{typ: ent.TypeAuthIdentity, tq: q}, nil case *ent.AuthIdentityChannelQuery: return &query[*ent.AuthIdentityChannelQuery, predicate.AuthIdentityChannel, authidentitychannel.OrderOption]{typ: ent.TypeAuthIdentityChannel, tq: q}, nil + case *ent.BatchImageEventQuery: + return &query[*ent.BatchImageEventQuery, predicate.BatchImageEvent, batchimageevent.OrderOption]{typ: ent.TypeBatchImageEvent, tq: q}, nil + case *ent.BatchImageItemQuery: + return &query[*ent.BatchImageItemQuery, predicate.BatchImageItem, batchimageitem.OrderOption]{typ: ent.TypeBatchImageItem, tq: q}, nil + case *ent.BatchImageJobQuery: + return &query[*ent.BatchImageJobQuery, predicate.BatchImageJob, batchimagejob.OrderOption]{typ: ent.TypeBatchImageJob, tq: q}, nil case *ent.ChannelMonitorQuery: return &query[*ent.ChannelMonitorQuery, predicate.ChannelMonitor, channelmonitor.OrderOption]{typ: ent.TypeChannelMonitor, tq: q}, nil case *ent.ChannelMonitorDailyRollupQuery: diff --git a/backend/ent/migrate/schema.go b/backend/ent/migrate/schema.go index c771f9572d..d3e8bc5448 100644 --- a/backend/ent/migrate/schema.go +++ b/backend/ent/migrate/schema.go @@ -435,6 +435,188 @@ var ( }, }, } + // BatchImageEventsColumns holds the columns for the "batch_image_events" table. + BatchImageEventsColumns = []*schema.Column{ + {Name: "id", Type: field.TypeInt64, Increment: true}, + {Name: "job_id", Type: field.TypeString, Size: 64}, + {Name: "event_type", Type: field.TypeString, Size: 64}, + {Name: "payload", Type: field.TypeJSON, Nullable: true, SchemaType: map[string]string{"postgres": "jsonb"}}, + {Name: "event_hash", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}}, + } + // BatchImageEventsTable holds the schema information for the "batch_image_events" table. + BatchImageEventsTable = &schema.Table{ + Name: "batch_image_events", + Columns: BatchImageEventsColumns, + PrimaryKey: []*schema.Column{BatchImageEventsColumns[0]}, + Indexes: []*schema.Index{ + { + Name: "batchimageevent_job_id_created_at", + Unique: false, + Columns: []*schema.Column{BatchImageEventsColumns[1], BatchImageEventsColumns[5]}, + }, + { + Name: "batchimageevent_event_type", + Unique: false, + Columns: []*schema.Column{BatchImageEventsColumns[2]}, + }, + { + Name: "batchimageevent_job_id_event_hash", + Unique: true, + Columns: []*schema.Column{BatchImageEventsColumns[1], BatchImageEventsColumns[4]}, + Annotation: &entsql.IndexAnnotation{ + Where: "event_hash IS NOT NULL AND event_hash <> ''", + }, + }, + }, + } + // BatchImageItemsColumns holds the columns for the "batch_image_items" table. + BatchImageItemsColumns = []*schema.Column{ + {Name: "id", Type: field.TypeInt64, Increment: true}, + {Name: "job_id", Type: field.TypeString, Size: 64}, + {Name: "custom_id", Type: field.TypeString, Size: 255}, + {Name: "status", Type: field.TypeString, Size: 32}, + {Name: "request_hash", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "prompt_preview", Type: field.TypeString, Nullable: true, SchemaType: map[string]string{"postgres": "text"}}, + {Name: "provider_source_object", Type: field.TypeString, Nullable: true, Size: 1024}, + {Name: "source_line_number", Type: field.TypeInt, Nullable: true}, + {Name: "source_byte_offset", Type: field.TypeInt64, Nullable: true}, + {Name: "source_byte_length", Type: field.TypeInt64, Nullable: true}, + {Name: "mime_type", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "file_extension", Type: field.TypeString, Nullable: true, Size: 32}, + {Name: "image_count", Type: field.TypeInt, Default: 0}, + {Name: "error_code", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "error_message", Type: field.TypeString, Nullable: true, SchemaType: map[string]string{"postgres": "text"}}, + {Name: "billed_amount", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,10)"}}, + {Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "indexed_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + } + // BatchImageItemsTable holds the schema information for the "batch_image_items" table. + BatchImageItemsTable = &schema.Table{ + Name: "batch_image_items", + Columns: BatchImageItemsColumns, + PrimaryKey: []*schema.Column{BatchImageItemsColumns[0]}, + Indexes: []*schema.Index{ + { + Name: "batchimageitem_job_id_custom_id", + Unique: true, + Columns: []*schema.Column{BatchImageItemsColumns[1], BatchImageItemsColumns[2]}, + }, + { + Name: "batchimageitem_job_id_status", + Unique: false, + Columns: []*schema.Column{BatchImageItemsColumns[1], BatchImageItemsColumns[3]}, + }, + { + Name: "batchimageitem_provider_source_object", + Unique: false, + Columns: []*schema.Column{BatchImageItemsColumns[6]}, + }, + }, + } + // BatchImageJobsColumns holds the columns for the "batch_image_jobs" table. + BatchImageJobsColumns = []*schema.Column{ + {Name: "id", Type: field.TypeInt64, Increment: true}, + {Name: "batch_id", Type: field.TypeString, Size: 64}, + {Name: "user_id", Type: field.TypeInt64}, + {Name: "api_key_id", Type: field.TypeInt64, Nullable: true}, + {Name: "account_id", Type: field.TypeInt64, Nullable: true}, + {Name: "provider", Type: field.TypeString, Size: 32}, + {Name: "model", Type: field.TypeString, Size: 128}, + {Name: "task_name", Type: field.TypeString, Size: 255, Default: ""}, + {Name: "status", Type: field.TypeString, Size: 32, Default: "created"}, + {Name: "provider_job_name", Type: field.TypeString, Nullable: true, Size: 512}, + {Name: "provider_input_ref", Type: field.TypeString, Nullable: true, Size: 1024}, + {Name: "provider_output_ref", Type: field.TypeString, Nullable: true, Size: 1024}, + {Name: "gcs_input_uri", Type: field.TypeString, Nullable: true, Size: 1024}, + {Name: "gcs_output_uri", Type: field.TypeString, Nullable: true, Size: 1024}, + {Name: "item_count", Type: field.TypeInt}, + {Name: "success_count", Type: field.TypeInt, Default: 0}, + {Name: "fail_count", Type: field.TypeInt, Default: 0}, + {Name: "cancelled_count", Type: field.TypeInt, Default: 0}, + {Name: "estimated_cost", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,10)"}}, + {Name: "hold_amount", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,10)"}}, + {Name: "actual_cost", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,10)"}}, + {Name: "currency", Type: field.TypeString, Size: 16, Default: "USD"}, + {Name: "hold_id", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "idempotency_key", Type: field.TypeString, Nullable: true, Size: 255}, + {Name: "request_hash", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "manifest_hash", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "retry_count", Type: field.TypeInt, Default: 0}, + {Name: "version", Type: field.TypeInt, Default: 0}, + {Name: "output_expires_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "input_deleted_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "output_deleted_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "downloaded_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "user_deleted_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "last_error_code", Type: field.TypeString, Nullable: true, Size: 128}, + {Name: "last_error_message", Type: field.TypeString, Nullable: true, SchemaType: map[string]string{"postgres": "text"}}, + {Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "updated_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "submitted_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "started_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "finished_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + {Name: "settled_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}}, + } + // BatchImageJobsTable holds the schema information for the "batch_image_jobs" table. + BatchImageJobsTable = &schema.Table{ + Name: "batch_image_jobs", + Columns: BatchImageJobsColumns, + PrimaryKey: []*schema.Column{BatchImageJobsColumns[0]}, + Indexes: []*schema.Index{ + { + Name: "batchimagejob_batch_id", + Unique: true, + Columns: []*schema.Column{BatchImageJobsColumns[1]}, + }, + { + Name: "batchimagejob_user_id_created_at", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[2], BatchImageJobsColumns[35]}, + }, + { + Name: "batchimagejob_status", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[8]}, + }, + { + Name: "batchimagejob_provider_status", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[5], BatchImageJobsColumns[8]}, + }, + { + Name: "batchimagejob_idempotency_key", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[23]}, + Annotation: &entsql.IndexAnnotation{ + Where: "idempotency_key IS NOT NULL AND idempotency_key <> ''", + }, + }, + { + Name: "batchimagejob_manifest_hash", + Unique: true, + Columns: []*schema.Column{BatchImageJobsColumns[25]}, + Annotation: &entsql.IndexAnnotation{ + Where: "manifest_hash IS NOT NULL AND manifest_hash <> ''", + }, + }, + { + Name: "batchimagejob_output_expires_at", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[28]}, + }, + { + Name: "batchimagejob_downloaded_at", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[31]}, + }, + { + Name: "batchimagejob_user_deleted_at", + Unique: false, + Columns: []*schema.Column{BatchImageJobsColumns[32]}, + }, + }, + } // ChannelMonitorsColumns holds the columns for the "channel_monitors" table. ChannelMonitorsColumns = []*schema.Column{ {Name: "id", Type: field.TypeInt64, Increment: true}, @@ -670,11 +852,19 @@ var ( {Name: "monthly_limit_usd", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "default_validity_days", Type: field.TypeInt, Default: 30}, {Name: "allow_image_generation", Type: field.TypeBool, Default: false}, + {Name: "allow_batch_image_generation", Type: field.TypeBool, Default: false}, {Name: "image_rate_independent", Type: field.TypeBool, Default: false}, {Name: "image_rate_multiplier", Type: field.TypeFloat64, Default: 1, SchemaType: map[string]string{"postgres": "decimal(10,4)"}}, {Name: "image_price_1k", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "image_price_2k", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "image_price_4k", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, + {Name: "batch_image_discount_multiplier", Type: field.TypeFloat64, Default: 0.5, SchemaType: map[string]string{"postgres": "decimal(10,4)"}}, + {Name: "batch_image_hold_multiplier", Type: field.TypeFloat64, Default: 0.6, SchemaType: map[string]string{"postgres": "decimal(10,4)"}}, + {Name: "video_rate_independent", Type: field.TypeBool, Default: false}, + {Name: "video_rate_multiplier", Type: field.TypeFloat64, Default: 1, SchemaType: map[string]string{"postgres": "decimal(10,4)"}}, + {Name: "video_price_480p", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, + {Name: "video_price_720p", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, + {Name: "video_price_1080p", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "claude_code_only", Type: field.TypeBool, Default: false}, {Name: "fallback_group_id", Type: field.TypeInt64, Nullable: true}, {Name: "fallback_group_id_on_invalid_request", Type: field.TypeInt64, Nullable: true}, @@ -725,7 +915,7 @@ var ( { Name: "group_sort_order", Unique: false, - Columns: []*schema.Column{GroupsColumns[32]}, + Columns: []*schema.Column{GroupsColumns[40]}, }, }, } @@ -1382,6 +1572,9 @@ var ( {Name: "image_output_size", Type: field.TypeString, Nullable: true, Size: 32}, {Name: "image_size_source", Type: field.TypeString, Nullable: true, Size: 16}, {Name: "image_size_breakdown", Type: field.TypeJSON, Nullable: true, SchemaType: map[string]string{"postgres": "jsonb"}}, + {Name: "video_count", Type: field.TypeInt, Default: 0}, + {Name: "video_resolution", Type: field.TypeString, Nullable: true, Size: 10}, + {Name: "video_duration_seconds", Type: field.TypeInt, Nullable: true}, {Name: "cache_ttl_overridden", Type: field.TypeBool, Default: false}, {Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}}, {Name: "api_key_id", Type: field.TypeInt64}, @@ -1398,31 +1591,31 @@ var ( ForeignKeys: []*schema.ForeignKey{ { Symbol: "usage_logs_api_keys_usage_logs", - Columns: []*schema.Column{UsageLogsColumns[37]}, + Columns: []*schema.Column{UsageLogsColumns[40]}, RefColumns: []*schema.Column{APIKeysColumns[0]}, OnDelete: schema.NoAction, }, { Symbol: "usage_logs_accounts_usage_logs", - Columns: []*schema.Column{UsageLogsColumns[38]}, + Columns: []*schema.Column{UsageLogsColumns[41]}, RefColumns: []*schema.Column{AccountsColumns[0]}, OnDelete: schema.NoAction, }, { Symbol: "usage_logs_groups_usage_logs", - Columns: []*schema.Column{UsageLogsColumns[39]}, + Columns: []*schema.Column{UsageLogsColumns[42]}, RefColumns: []*schema.Column{GroupsColumns[0]}, OnDelete: schema.SetNull, }, { Symbol: "usage_logs_users_usage_logs", - Columns: []*schema.Column{UsageLogsColumns[40]}, + Columns: []*schema.Column{UsageLogsColumns[43]}, RefColumns: []*schema.Column{UsersColumns[0]}, OnDelete: schema.NoAction, }, { Symbol: "usage_logs_user_subscriptions_usage_logs", - Columns: []*schema.Column{UsageLogsColumns[41]}, + Columns: []*schema.Column{UsageLogsColumns[44]}, RefColumns: []*schema.Column{UserSubscriptionsColumns[0]}, OnDelete: schema.SetNull, }, @@ -1431,32 +1624,32 @@ var ( { Name: "usagelog_user_id", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[40]}, + Columns: []*schema.Column{UsageLogsColumns[43]}, }, { Name: "usagelog_api_key_id", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[37]}, + Columns: []*schema.Column{UsageLogsColumns[40]}, }, { Name: "usagelog_account_id", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[38]}, + Columns: []*schema.Column{UsageLogsColumns[41]}, }, { Name: "usagelog_group_id", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[39]}, + Columns: []*schema.Column{UsageLogsColumns[42]}, }, { Name: "usagelog_subscription_id", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[41]}, + Columns: []*schema.Column{UsageLogsColumns[44]}, }, { Name: "usagelog_created_at", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[36]}, + Columns: []*schema.Column{UsageLogsColumns[39]}, }, { Name: "usagelog_model", @@ -1476,17 +1669,17 @@ var ( { Name: "usagelog_user_id_created_at", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[40], UsageLogsColumns[36]}, + Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[39]}, }, { Name: "usagelog_api_key_id_created_at", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[37], UsageLogsColumns[36]}, + Columns: []*schema.Column{UsageLogsColumns[40], UsageLogsColumns[39]}, }, { Name: "usagelog_group_id_created_at", Unique: false, - Columns: []*schema.Column{UsageLogsColumns[39], UsageLogsColumns[36]}, + Columns: []*schema.Column{UsageLogsColumns[42], UsageLogsColumns[39]}, }, }, } @@ -1500,6 +1693,7 @@ var ( {Name: "password_hash", Type: field.TypeString, Size: 255}, {Name: "role", Type: field.TypeString, Size: 20, Default: "user"}, {Name: "balance", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, + {Name: "frozen_balance", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "concurrency", Type: field.TypeInt, Default: 5}, {Name: "status", Type: field.TypeString, Size: 20, Default: "active"}, {Name: "username", Type: field.TypeString, Size: 100, Default: ""}, @@ -1526,7 +1720,7 @@ var ( { Name: "user_status", Unique: false, - Columns: []*schema.Column{UsersColumns[9]}, + Columns: []*schema.Column{UsersColumns[10]}, }, { Name: "user_deleted_at", @@ -1799,6 +1993,9 @@ var ( AnnouncementReadsTable, AuthIdentitiesTable, AuthIdentityChannelsTable, + BatchImageEventsTable, + BatchImageItemsTable, + BatchImageJobsTable, ChannelMonitorsTable, ChannelMonitorDailyRollupsTable, ChannelMonitorHistoriesTable, @@ -1862,6 +2059,15 @@ func init() { AuthIdentityChannelsTable.Annotation = &entsql.Annotation{ Table: "auth_identity_channels", } + BatchImageEventsTable.Annotation = &entsql.Annotation{ + Table: "batch_image_events", + } + BatchImageItemsTable.Annotation = &entsql.Annotation{ + Table: "batch_image_items", + } + BatchImageJobsTable.Annotation = &entsql.Annotation{ + Table: "batch_image_jobs", + } ChannelMonitorsTable.ForeignKeys[0].RefTable = ChannelMonitorRequestTemplatesTable ChannelMonitorsTable.Annotation = &entsql.Annotation{ Table: "channel_monitors", diff --git a/backend/ent/mutation.go b/backend/ent/mutation.go index c71850e4ce..8d32773050 100644 --- a/backend/ent/mutation.go +++ b/backend/ent/mutation.go @@ -19,6 +19,9 @@ import ( "github.com/Wei-Shaw/sub2api/ent/apikey" "github.com/Wei-Shaw/sub2api/ent/authidentity" "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" "github.com/Wei-Shaw/sub2api/ent/channelmonitor" "github.com/Wei-Shaw/sub2api/ent/channelmonitordailyrollup" "github.com/Wei-Shaw/sub2api/ent/channelmonitorhistory" @@ -67,6 +70,9 @@ const ( TypeAnnouncementRead = "AnnouncementRead" TypeAuthIdentity = "AuthIdentity" TypeAuthIdentityChannel = "AuthIdentityChannel" + TypeBatchImageEvent = "BatchImageEvent" + TypeBatchImageItem = "BatchImageItem" + TypeBatchImageJob = "BatchImageJob" TypeChannelMonitor = "ChannelMonitor" TypeChannelMonitorDailyRollup = "ChannelMonitorDailyRollup" TypeChannelMonitorHistory = "ChannelMonitorHistory" @@ -9120,6 +9126,5476 @@ func (m *AuthIdentityChannelMutation) ResetEdge(name string) error { return fmt.Errorf("unknown AuthIdentityChannel edge %s", name) } +// BatchImageEventMutation represents an operation that mutates the BatchImageEvent nodes in the graph. +type BatchImageEventMutation struct { + config + op Op + typ string + id *int64 + job_id *string + event_type *string + payload *map[string]interface{} + event_hash *string + created_at *time.Time + clearedFields map[string]struct{} + done bool + oldValue func(context.Context) (*BatchImageEvent, error) + predicates []predicate.BatchImageEvent +} + +var _ ent.Mutation = (*BatchImageEventMutation)(nil) + +// batchimageeventOption allows management of the mutation configuration using functional options. +type batchimageeventOption func(*BatchImageEventMutation) + +// newBatchImageEventMutation creates new mutation for the BatchImageEvent entity. +func newBatchImageEventMutation(c config, op Op, opts ...batchimageeventOption) *BatchImageEventMutation { + m := &BatchImageEventMutation{ + config: c, + op: op, + typ: TypeBatchImageEvent, + clearedFields: make(map[string]struct{}), + } + for _, opt := range opts { + opt(m) + } + return m +} + +// withBatchImageEventID sets the ID field of the mutation. +func withBatchImageEventID(id int64) batchimageeventOption { + return func(m *BatchImageEventMutation) { + var ( + err error + once sync.Once + value *BatchImageEvent + ) + m.oldValue = func(ctx context.Context) (*BatchImageEvent, error) { + once.Do(func() { + if m.done { + err = errors.New("querying old values post mutation is not allowed") + } else { + value, err = m.Client().BatchImageEvent.Get(ctx, id) + } + }) + return value, err + } + m.id = &id + } +} + +// withBatchImageEvent sets the old BatchImageEvent of the mutation. +func withBatchImageEvent(node *BatchImageEvent) batchimageeventOption { + return func(m *BatchImageEventMutation) { + m.oldValue = func(context.Context) (*BatchImageEvent, error) { + return node, nil + } + m.id = &node.ID + } +} + +// Client returns a new `ent.Client` from the mutation. If the mutation was +// executed in a transaction (ent.Tx), a transactional client is returned. +func (m BatchImageEventMutation) Client() *Client { + client := &Client{config: m.config} + client.init() + return client +} + +// Tx returns an `ent.Tx` for mutations that were executed in transactions; +// it returns an error otherwise. +func (m BatchImageEventMutation) Tx() (*Tx, error) { + if _, ok := m.driver.(*txDriver); !ok { + return nil, errors.New("ent: mutation is not running in a transaction") + } + tx := &Tx{config: m.config} + tx.init() + return tx, nil +} + +// ID returns the ID value in the mutation. Note that the ID is only available +// if it was provided to the builder or after it was returned from the database. +func (m *BatchImageEventMutation) ID() (id int64, exists bool) { + if m.id == nil { + return + } + return *m.id, true +} + +// IDs queries the database and returns the entity ids that match the mutation's predicate. +// That means, if the mutation is applied within a transaction with an isolation level such +// as sql.LevelSerializable, the returned ids match the ids of the rows that will be updated +// or updated by the mutation. +func (m *BatchImageEventMutation) IDs(ctx context.Context) ([]int64, error) { + switch { + case m.op.Is(OpUpdateOne | OpDeleteOne): + id, exists := m.ID() + if exists { + return []int64{id}, nil + } + fallthrough + case m.op.Is(OpUpdate | OpDelete): + return m.Client().BatchImageEvent.Query().Where(m.predicates...).IDs(ctx) + default: + return nil, fmt.Errorf("IDs is not allowed on %s operations", m.op) + } +} + +// SetJobID sets the "job_id" field. +func (m *BatchImageEventMutation) SetJobID(s string) { + m.job_id = &s +} + +// JobID returns the value of the "job_id" field in the mutation. +func (m *BatchImageEventMutation) JobID() (r string, exists bool) { + v := m.job_id + if v == nil { + return + } + return *v, true +} + +// OldJobID returns the old "job_id" field's value of the BatchImageEvent entity. +// If the BatchImageEvent object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageEventMutation) OldJobID(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldJobID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldJobID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldJobID: %w", err) + } + return oldValue.JobID, nil +} + +// ResetJobID resets all changes to the "job_id" field. +func (m *BatchImageEventMutation) ResetJobID() { + m.job_id = nil +} + +// SetEventType sets the "event_type" field. +func (m *BatchImageEventMutation) SetEventType(s string) { + m.event_type = &s +} + +// EventType returns the value of the "event_type" field in the mutation. +func (m *BatchImageEventMutation) EventType() (r string, exists bool) { + v := m.event_type + if v == nil { + return + } + return *v, true +} + +// OldEventType returns the old "event_type" field's value of the BatchImageEvent entity. +// If the BatchImageEvent object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageEventMutation) OldEventType(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldEventType is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldEventType requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldEventType: %w", err) + } + return oldValue.EventType, nil +} + +// ResetEventType resets all changes to the "event_type" field. +func (m *BatchImageEventMutation) ResetEventType() { + m.event_type = nil +} + +// SetPayload sets the "payload" field. +func (m *BatchImageEventMutation) SetPayload(value map[string]interface{}) { + m.payload = &value +} + +// Payload returns the value of the "payload" field in the mutation. +func (m *BatchImageEventMutation) Payload() (r map[string]interface{}, exists bool) { + v := m.payload + if v == nil { + return + } + return *v, true +} + +// OldPayload returns the old "payload" field's value of the BatchImageEvent entity. +// If the BatchImageEvent object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageEventMutation) OldPayload(ctx context.Context) (v map[string]interface{}, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldPayload is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldPayload requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldPayload: %w", err) + } + return oldValue.Payload, nil +} + +// ClearPayload clears the value of the "payload" field. +func (m *BatchImageEventMutation) ClearPayload() { + m.payload = nil + m.clearedFields[batchimageevent.FieldPayload] = struct{}{} +} + +// PayloadCleared returns if the "payload" field was cleared in this mutation. +func (m *BatchImageEventMutation) PayloadCleared() bool { + _, ok := m.clearedFields[batchimageevent.FieldPayload] + return ok +} + +// ResetPayload resets all changes to the "payload" field. +func (m *BatchImageEventMutation) ResetPayload() { + m.payload = nil + delete(m.clearedFields, batchimageevent.FieldPayload) +} + +// SetEventHash sets the "event_hash" field. +func (m *BatchImageEventMutation) SetEventHash(s string) { + m.event_hash = &s +} + +// EventHash returns the value of the "event_hash" field in the mutation. +func (m *BatchImageEventMutation) EventHash() (r string, exists bool) { + v := m.event_hash + if v == nil { + return + } + return *v, true +} + +// OldEventHash returns the old "event_hash" field's value of the BatchImageEvent entity. +// If the BatchImageEvent object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageEventMutation) OldEventHash(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldEventHash is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldEventHash requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldEventHash: %w", err) + } + return oldValue.EventHash, nil +} + +// ClearEventHash clears the value of the "event_hash" field. +func (m *BatchImageEventMutation) ClearEventHash() { + m.event_hash = nil + m.clearedFields[batchimageevent.FieldEventHash] = struct{}{} +} + +// EventHashCleared returns if the "event_hash" field was cleared in this mutation. +func (m *BatchImageEventMutation) EventHashCleared() bool { + _, ok := m.clearedFields[batchimageevent.FieldEventHash] + return ok +} + +// ResetEventHash resets all changes to the "event_hash" field. +func (m *BatchImageEventMutation) ResetEventHash() { + m.event_hash = nil + delete(m.clearedFields, batchimageevent.FieldEventHash) +} + +// SetCreatedAt sets the "created_at" field. +func (m *BatchImageEventMutation) SetCreatedAt(t time.Time) { + m.created_at = &t +} + +// CreatedAt returns the value of the "created_at" field in the mutation. +func (m *BatchImageEventMutation) CreatedAt() (r time.Time, exists bool) { + v := m.created_at + if v == nil { + return + } + return *v, true +} + +// OldCreatedAt returns the old "created_at" field's value of the BatchImageEvent entity. +// If the BatchImageEvent object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageEventMutation) OldCreatedAt(ctx context.Context) (v time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldCreatedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldCreatedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldCreatedAt: %w", err) + } + return oldValue.CreatedAt, nil +} + +// ResetCreatedAt resets all changes to the "created_at" field. +func (m *BatchImageEventMutation) ResetCreatedAt() { + m.created_at = nil +} + +// Where appends a list predicates to the BatchImageEventMutation builder. +func (m *BatchImageEventMutation) Where(ps ...predicate.BatchImageEvent) { + m.predicates = append(m.predicates, ps...) +} + +// WhereP appends storage-level predicates to the BatchImageEventMutation builder. Using this method, +// users can use type-assertion to append predicates that do not depend on any generated package. +func (m *BatchImageEventMutation) WhereP(ps ...func(*sql.Selector)) { + p := make([]predicate.BatchImageEvent, len(ps)) + for i := range ps { + p[i] = ps[i] + } + m.Where(p...) +} + +// Op returns the operation name. +func (m *BatchImageEventMutation) Op() Op { + return m.op +} + +// SetOp allows setting the mutation operation. +func (m *BatchImageEventMutation) SetOp(op Op) { + m.op = op +} + +// Type returns the node type of this mutation (BatchImageEvent). +func (m *BatchImageEventMutation) Type() string { + return m.typ +} + +// Fields returns all fields that were changed during this mutation. Note that in +// order to get all numeric fields that were incremented/decremented, call +// AddedFields(). +func (m *BatchImageEventMutation) Fields() []string { + fields := make([]string, 0, 5) + if m.job_id != nil { + fields = append(fields, batchimageevent.FieldJobID) + } + if m.event_type != nil { + fields = append(fields, batchimageevent.FieldEventType) + } + if m.payload != nil { + fields = append(fields, batchimageevent.FieldPayload) + } + if m.event_hash != nil { + fields = append(fields, batchimageevent.FieldEventHash) + } + if m.created_at != nil { + fields = append(fields, batchimageevent.FieldCreatedAt) + } + return fields +} + +// Field returns the value of a field with the given name. The second boolean +// return value indicates that this field was not set, or was not defined in the +// schema. +func (m *BatchImageEventMutation) Field(name string) (ent.Value, bool) { + switch name { + case batchimageevent.FieldJobID: + return m.JobID() + case batchimageevent.FieldEventType: + return m.EventType() + case batchimageevent.FieldPayload: + return m.Payload() + case batchimageevent.FieldEventHash: + return m.EventHash() + case batchimageevent.FieldCreatedAt: + return m.CreatedAt() + } + return nil, false +} + +// OldField returns the old value of the field from the database. An error is +// returned if the mutation operation is not UpdateOne, or the query to the +// database failed. +func (m *BatchImageEventMutation) OldField(ctx context.Context, name string) (ent.Value, error) { + switch name { + case batchimageevent.FieldJobID: + return m.OldJobID(ctx) + case batchimageevent.FieldEventType: + return m.OldEventType(ctx) + case batchimageevent.FieldPayload: + return m.OldPayload(ctx) + case batchimageevent.FieldEventHash: + return m.OldEventHash(ctx) + case batchimageevent.FieldCreatedAt: + return m.OldCreatedAt(ctx) + } + return nil, fmt.Errorf("unknown BatchImageEvent field %s", name) +} + +// SetField sets the value of a field with the given name. It returns an error if +// the field is not defined in the schema, or if the type mismatched the field +// type. +func (m *BatchImageEventMutation) SetField(name string, value ent.Value) error { + switch name { + case batchimageevent.FieldJobID: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetJobID(v) + return nil + case batchimageevent.FieldEventType: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetEventType(v) + return nil + case batchimageevent.FieldPayload: + v, ok := value.(map[string]interface{}) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetPayload(v) + return nil + case batchimageevent.FieldEventHash: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetEventHash(v) + return nil + case batchimageevent.FieldCreatedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetCreatedAt(v) + return nil + } + return fmt.Errorf("unknown BatchImageEvent field %s", name) +} + +// AddedFields returns all numeric fields that were incremented/decremented during +// this mutation. +func (m *BatchImageEventMutation) AddedFields() []string { + return nil +} + +// AddedField returns the numeric value that was incremented/decremented on a field +// with the given name. The second boolean return value indicates that this field +// was not set, or was not defined in the schema. +func (m *BatchImageEventMutation) AddedField(name string) (ent.Value, bool) { + return nil, false +} + +// AddField adds the value to the field with the given name. It returns an error if +// the field is not defined in the schema, or if the type mismatched the field +// type. +func (m *BatchImageEventMutation) AddField(name string, value ent.Value) error { + switch name { + } + return fmt.Errorf("unknown BatchImageEvent numeric field %s", name) +} + +// ClearedFields returns all nullable fields that were cleared during this +// mutation. +func (m *BatchImageEventMutation) ClearedFields() []string { + var fields []string + if m.FieldCleared(batchimageevent.FieldPayload) { + fields = append(fields, batchimageevent.FieldPayload) + } + if m.FieldCleared(batchimageevent.FieldEventHash) { + fields = append(fields, batchimageevent.FieldEventHash) + } + return fields +} + +// FieldCleared returns a boolean indicating if a field with the given name was +// cleared in this mutation. +func (m *BatchImageEventMutation) FieldCleared(name string) bool { + _, ok := m.clearedFields[name] + return ok +} + +// ClearField clears the value of the field with the given name. It returns an +// error if the field is not defined in the schema. +func (m *BatchImageEventMutation) ClearField(name string) error { + switch name { + case batchimageevent.FieldPayload: + m.ClearPayload() + return nil + case batchimageevent.FieldEventHash: + m.ClearEventHash() + return nil + } + return fmt.Errorf("unknown BatchImageEvent nullable field %s", name) +} + +// ResetField resets all changes in the mutation for the field with the given name. +// It returns an error if the field is not defined in the schema. +func (m *BatchImageEventMutation) ResetField(name string) error { + switch name { + case batchimageevent.FieldJobID: + m.ResetJobID() + return nil + case batchimageevent.FieldEventType: + m.ResetEventType() + return nil + case batchimageevent.FieldPayload: + m.ResetPayload() + return nil + case batchimageevent.FieldEventHash: + m.ResetEventHash() + return nil + case batchimageevent.FieldCreatedAt: + m.ResetCreatedAt() + return nil + } + return fmt.Errorf("unknown BatchImageEvent field %s", name) +} + +// AddedEdges returns all edge names that were set/added in this mutation. +func (m *BatchImageEventMutation) AddedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// AddedIDs returns all IDs (to other nodes) that were added for the given edge +// name in this mutation. +func (m *BatchImageEventMutation) AddedIDs(name string) []ent.Value { + return nil +} + +// RemovedEdges returns all edge names that were removed in this mutation. +func (m *BatchImageEventMutation) RemovedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// RemovedIDs returns all IDs (to other nodes) that were removed for the edge with +// the given name in this mutation. +func (m *BatchImageEventMutation) RemovedIDs(name string) []ent.Value { + return nil +} + +// ClearedEdges returns all edge names that were cleared in this mutation. +func (m *BatchImageEventMutation) ClearedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// EdgeCleared returns a boolean which indicates if the edge with the given name +// was cleared in this mutation. +func (m *BatchImageEventMutation) EdgeCleared(name string) bool { + return false +} + +// ClearEdge clears the value of the edge with the given name. It returns an error +// if that edge is not defined in the schema. +func (m *BatchImageEventMutation) ClearEdge(name string) error { + return fmt.Errorf("unknown BatchImageEvent unique edge %s", name) +} + +// ResetEdge resets all changes to the edge with the given name in this mutation. +// It returns an error if the edge is not defined in the schema. +func (m *BatchImageEventMutation) ResetEdge(name string) error { + return fmt.Errorf("unknown BatchImageEvent edge %s", name) +} + +// BatchImageItemMutation represents an operation that mutates the BatchImageItem nodes in the graph. +type BatchImageItemMutation struct { + config + op Op + typ string + id *int64 + job_id *string + custom_id *string + status *string + request_hash *string + prompt_preview *string + provider_source_object *string + source_line_number *int + addsource_line_number *int + source_byte_offset *int64 + addsource_byte_offset *int64 + source_byte_length *int64 + addsource_byte_length *int64 + mime_type *string + file_extension *string + image_count *int + addimage_count *int + error_code *string + error_message *string + billed_amount *float64 + addbilled_amount *float64 + created_at *time.Time + indexed_at *time.Time + clearedFields map[string]struct{} + done bool + oldValue func(context.Context) (*BatchImageItem, error) + predicates []predicate.BatchImageItem +} + +var _ ent.Mutation = (*BatchImageItemMutation)(nil) + +// batchimageitemOption allows management of the mutation configuration using functional options. +type batchimageitemOption func(*BatchImageItemMutation) + +// newBatchImageItemMutation creates new mutation for the BatchImageItem entity. +func newBatchImageItemMutation(c config, op Op, opts ...batchimageitemOption) *BatchImageItemMutation { + m := &BatchImageItemMutation{ + config: c, + op: op, + typ: TypeBatchImageItem, + clearedFields: make(map[string]struct{}), + } + for _, opt := range opts { + opt(m) + } + return m +} + +// withBatchImageItemID sets the ID field of the mutation. +func withBatchImageItemID(id int64) batchimageitemOption { + return func(m *BatchImageItemMutation) { + var ( + err error + once sync.Once + value *BatchImageItem + ) + m.oldValue = func(ctx context.Context) (*BatchImageItem, error) { + once.Do(func() { + if m.done { + err = errors.New("querying old values post mutation is not allowed") + } else { + value, err = m.Client().BatchImageItem.Get(ctx, id) + } + }) + return value, err + } + m.id = &id + } +} + +// withBatchImageItem sets the old BatchImageItem of the mutation. +func withBatchImageItem(node *BatchImageItem) batchimageitemOption { + return func(m *BatchImageItemMutation) { + m.oldValue = func(context.Context) (*BatchImageItem, error) { + return node, nil + } + m.id = &node.ID + } +} + +// Client returns a new `ent.Client` from the mutation. If the mutation was +// executed in a transaction (ent.Tx), a transactional client is returned. +func (m BatchImageItemMutation) Client() *Client { + client := &Client{config: m.config} + client.init() + return client +} + +// Tx returns an `ent.Tx` for mutations that were executed in transactions; +// it returns an error otherwise. +func (m BatchImageItemMutation) Tx() (*Tx, error) { + if _, ok := m.driver.(*txDriver); !ok { + return nil, errors.New("ent: mutation is not running in a transaction") + } + tx := &Tx{config: m.config} + tx.init() + return tx, nil +} + +// ID returns the ID value in the mutation. Note that the ID is only available +// if it was provided to the builder or after it was returned from the database. +func (m *BatchImageItemMutation) ID() (id int64, exists bool) { + if m.id == nil { + return + } + return *m.id, true +} + +// IDs queries the database and returns the entity ids that match the mutation's predicate. +// That means, if the mutation is applied within a transaction with an isolation level such +// as sql.LevelSerializable, the returned ids match the ids of the rows that will be updated +// or updated by the mutation. +func (m *BatchImageItemMutation) IDs(ctx context.Context) ([]int64, error) { + switch { + case m.op.Is(OpUpdateOne | OpDeleteOne): + id, exists := m.ID() + if exists { + return []int64{id}, nil + } + fallthrough + case m.op.Is(OpUpdate | OpDelete): + return m.Client().BatchImageItem.Query().Where(m.predicates...).IDs(ctx) + default: + return nil, fmt.Errorf("IDs is not allowed on %s operations", m.op) + } +} + +// SetJobID sets the "job_id" field. +func (m *BatchImageItemMutation) SetJobID(s string) { + m.job_id = &s +} + +// JobID returns the value of the "job_id" field in the mutation. +func (m *BatchImageItemMutation) JobID() (r string, exists bool) { + v := m.job_id + if v == nil { + return + } + return *v, true +} + +// OldJobID returns the old "job_id" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldJobID(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldJobID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldJobID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldJobID: %w", err) + } + return oldValue.JobID, nil +} + +// ResetJobID resets all changes to the "job_id" field. +func (m *BatchImageItemMutation) ResetJobID() { + m.job_id = nil +} + +// SetCustomID sets the "custom_id" field. +func (m *BatchImageItemMutation) SetCustomID(s string) { + m.custom_id = &s +} + +// CustomID returns the value of the "custom_id" field in the mutation. +func (m *BatchImageItemMutation) CustomID() (r string, exists bool) { + v := m.custom_id + if v == nil { + return + } + return *v, true +} + +// OldCustomID returns the old "custom_id" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldCustomID(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldCustomID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldCustomID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldCustomID: %w", err) + } + return oldValue.CustomID, nil +} + +// ResetCustomID resets all changes to the "custom_id" field. +func (m *BatchImageItemMutation) ResetCustomID() { + m.custom_id = nil +} + +// SetStatus sets the "status" field. +func (m *BatchImageItemMutation) SetStatus(s string) { + m.status = &s +} + +// Status returns the value of the "status" field in the mutation. +func (m *BatchImageItemMutation) Status() (r string, exists bool) { + v := m.status + if v == nil { + return + } + return *v, true +} + +// OldStatus returns the old "status" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldStatus(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldStatus is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldStatus requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldStatus: %w", err) + } + return oldValue.Status, nil +} + +// ResetStatus resets all changes to the "status" field. +func (m *BatchImageItemMutation) ResetStatus() { + m.status = nil +} + +// SetRequestHash sets the "request_hash" field. +func (m *BatchImageItemMutation) SetRequestHash(s string) { + m.request_hash = &s +} + +// RequestHash returns the value of the "request_hash" field in the mutation. +func (m *BatchImageItemMutation) RequestHash() (r string, exists bool) { + v := m.request_hash + if v == nil { + return + } + return *v, true +} + +// OldRequestHash returns the old "request_hash" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldRequestHash(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldRequestHash is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldRequestHash requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldRequestHash: %w", err) + } + return oldValue.RequestHash, nil +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (m *BatchImageItemMutation) ClearRequestHash() { + m.request_hash = nil + m.clearedFields[batchimageitem.FieldRequestHash] = struct{}{} +} + +// RequestHashCleared returns if the "request_hash" field was cleared in this mutation. +func (m *BatchImageItemMutation) RequestHashCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldRequestHash] + return ok +} + +// ResetRequestHash resets all changes to the "request_hash" field. +func (m *BatchImageItemMutation) ResetRequestHash() { + m.request_hash = nil + delete(m.clearedFields, batchimageitem.FieldRequestHash) +} + +// SetPromptPreview sets the "prompt_preview" field. +func (m *BatchImageItemMutation) SetPromptPreview(s string) { + m.prompt_preview = &s +} + +// PromptPreview returns the value of the "prompt_preview" field in the mutation. +func (m *BatchImageItemMutation) PromptPreview() (r string, exists bool) { + v := m.prompt_preview + if v == nil { + return + } + return *v, true +} + +// OldPromptPreview returns the old "prompt_preview" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldPromptPreview(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldPromptPreview is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldPromptPreview requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldPromptPreview: %w", err) + } + return oldValue.PromptPreview, nil +} + +// ClearPromptPreview clears the value of the "prompt_preview" field. +func (m *BatchImageItemMutation) ClearPromptPreview() { + m.prompt_preview = nil + m.clearedFields[batchimageitem.FieldPromptPreview] = struct{}{} +} + +// PromptPreviewCleared returns if the "prompt_preview" field was cleared in this mutation. +func (m *BatchImageItemMutation) PromptPreviewCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldPromptPreview] + return ok +} + +// ResetPromptPreview resets all changes to the "prompt_preview" field. +func (m *BatchImageItemMutation) ResetPromptPreview() { + m.prompt_preview = nil + delete(m.clearedFields, batchimageitem.FieldPromptPreview) +} + +// SetProviderSourceObject sets the "provider_source_object" field. +func (m *BatchImageItemMutation) SetProviderSourceObject(s string) { + m.provider_source_object = &s +} + +// ProviderSourceObject returns the value of the "provider_source_object" field in the mutation. +func (m *BatchImageItemMutation) ProviderSourceObject() (r string, exists bool) { + v := m.provider_source_object + if v == nil { + return + } + return *v, true +} + +// OldProviderSourceObject returns the old "provider_source_object" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldProviderSourceObject(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldProviderSourceObject is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldProviderSourceObject requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldProviderSourceObject: %w", err) + } + return oldValue.ProviderSourceObject, nil +} + +// ClearProviderSourceObject clears the value of the "provider_source_object" field. +func (m *BatchImageItemMutation) ClearProviderSourceObject() { + m.provider_source_object = nil + m.clearedFields[batchimageitem.FieldProviderSourceObject] = struct{}{} +} + +// ProviderSourceObjectCleared returns if the "provider_source_object" field was cleared in this mutation. +func (m *BatchImageItemMutation) ProviderSourceObjectCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldProviderSourceObject] + return ok +} + +// ResetProviderSourceObject resets all changes to the "provider_source_object" field. +func (m *BatchImageItemMutation) ResetProviderSourceObject() { + m.provider_source_object = nil + delete(m.clearedFields, batchimageitem.FieldProviderSourceObject) +} + +// SetSourceLineNumber sets the "source_line_number" field. +func (m *BatchImageItemMutation) SetSourceLineNumber(i int) { + m.source_line_number = &i + m.addsource_line_number = nil +} + +// SourceLineNumber returns the value of the "source_line_number" field in the mutation. +func (m *BatchImageItemMutation) SourceLineNumber() (r int, exists bool) { + v := m.source_line_number + if v == nil { + return + } + return *v, true +} + +// OldSourceLineNumber returns the old "source_line_number" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldSourceLineNumber(ctx context.Context) (v *int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldSourceLineNumber is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldSourceLineNumber requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldSourceLineNumber: %w", err) + } + return oldValue.SourceLineNumber, nil +} + +// AddSourceLineNumber adds i to the "source_line_number" field. +func (m *BatchImageItemMutation) AddSourceLineNumber(i int) { + if m.addsource_line_number != nil { + *m.addsource_line_number += i + } else { + m.addsource_line_number = &i + } +} + +// AddedSourceLineNumber returns the value that was added to the "source_line_number" field in this mutation. +func (m *BatchImageItemMutation) AddedSourceLineNumber() (r int, exists bool) { + v := m.addsource_line_number + if v == nil { + return + } + return *v, true +} + +// ClearSourceLineNumber clears the value of the "source_line_number" field. +func (m *BatchImageItemMutation) ClearSourceLineNumber() { + m.source_line_number = nil + m.addsource_line_number = nil + m.clearedFields[batchimageitem.FieldSourceLineNumber] = struct{}{} +} + +// SourceLineNumberCleared returns if the "source_line_number" field was cleared in this mutation. +func (m *BatchImageItemMutation) SourceLineNumberCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldSourceLineNumber] + return ok +} + +// ResetSourceLineNumber resets all changes to the "source_line_number" field. +func (m *BatchImageItemMutation) ResetSourceLineNumber() { + m.source_line_number = nil + m.addsource_line_number = nil + delete(m.clearedFields, batchimageitem.FieldSourceLineNumber) +} + +// SetSourceByteOffset sets the "source_byte_offset" field. +func (m *BatchImageItemMutation) SetSourceByteOffset(i int64) { + m.source_byte_offset = &i + m.addsource_byte_offset = nil +} + +// SourceByteOffset returns the value of the "source_byte_offset" field in the mutation. +func (m *BatchImageItemMutation) SourceByteOffset() (r int64, exists bool) { + v := m.source_byte_offset + if v == nil { + return + } + return *v, true +} + +// OldSourceByteOffset returns the old "source_byte_offset" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldSourceByteOffset(ctx context.Context) (v *int64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldSourceByteOffset is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldSourceByteOffset requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldSourceByteOffset: %w", err) + } + return oldValue.SourceByteOffset, nil +} + +// AddSourceByteOffset adds i to the "source_byte_offset" field. +func (m *BatchImageItemMutation) AddSourceByteOffset(i int64) { + if m.addsource_byte_offset != nil { + *m.addsource_byte_offset += i + } else { + m.addsource_byte_offset = &i + } +} + +// AddedSourceByteOffset returns the value that was added to the "source_byte_offset" field in this mutation. +func (m *BatchImageItemMutation) AddedSourceByteOffset() (r int64, exists bool) { + v := m.addsource_byte_offset + if v == nil { + return + } + return *v, true +} + +// ClearSourceByteOffset clears the value of the "source_byte_offset" field. +func (m *BatchImageItemMutation) ClearSourceByteOffset() { + m.source_byte_offset = nil + m.addsource_byte_offset = nil + m.clearedFields[batchimageitem.FieldSourceByteOffset] = struct{}{} +} + +// SourceByteOffsetCleared returns if the "source_byte_offset" field was cleared in this mutation. +func (m *BatchImageItemMutation) SourceByteOffsetCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldSourceByteOffset] + return ok +} + +// ResetSourceByteOffset resets all changes to the "source_byte_offset" field. +func (m *BatchImageItemMutation) ResetSourceByteOffset() { + m.source_byte_offset = nil + m.addsource_byte_offset = nil + delete(m.clearedFields, batchimageitem.FieldSourceByteOffset) +} + +// SetSourceByteLength sets the "source_byte_length" field. +func (m *BatchImageItemMutation) SetSourceByteLength(i int64) { + m.source_byte_length = &i + m.addsource_byte_length = nil +} + +// SourceByteLength returns the value of the "source_byte_length" field in the mutation. +func (m *BatchImageItemMutation) SourceByteLength() (r int64, exists bool) { + v := m.source_byte_length + if v == nil { + return + } + return *v, true +} + +// OldSourceByteLength returns the old "source_byte_length" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldSourceByteLength(ctx context.Context) (v *int64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldSourceByteLength is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldSourceByteLength requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldSourceByteLength: %w", err) + } + return oldValue.SourceByteLength, nil +} + +// AddSourceByteLength adds i to the "source_byte_length" field. +func (m *BatchImageItemMutation) AddSourceByteLength(i int64) { + if m.addsource_byte_length != nil { + *m.addsource_byte_length += i + } else { + m.addsource_byte_length = &i + } +} + +// AddedSourceByteLength returns the value that was added to the "source_byte_length" field in this mutation. +func (m *BatchImageItemMutation) AddedSourceByteLength() (r int64, exists bool) { + v := m.addsource_byte_length + if v == nil { + return + } + return *v, true +} + +// ClearSourceByteLength clears the value of the "source_byte_length" field. +func (m *BatchImageItemMutation) ClearSourceByteLength() { + m.source_byte_length = nil + m.addsource_byte_length = nil + m.clearedFields[batchimageitem.FieldSourceByteLength] = struct{}{} +} + +// SourceByteLengthCleared returns if the "source_byte_length" field was cleared in this mutation. +func (m *BatchImageItemMutation) SourceByteLengthCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldSourceByteLength] + return ok +} + +// ResetSourceByteLength resets all changes to the "source_byte_length" field. +func (m *BatchImageItemMutation) ResetSourceByteLength() { + m.source_byte_length = nil + m.addsource_byte_length = nil + delete(m.clearedFields, batchimageitem.FieldSourceByteLength) +} + +// SetMimeType sets the "mime_type" field. +func (m *BatchImageItemMutation) SetMimeType(s string) { + m.mime_type = &s +} + +// MimeType returns the value of the "mime_type" field in the mutation. +func (m *BatchImageItemMutation) MimeType() (r string, exists bool) { + v := m.mime_type + if v == nil { + return + } + return *v, true +} + +// OldMimeType returns the old "mime_type" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldMimeType(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldMimeType is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldMimeType requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldMimeType: %w", err) + } + return oldValue.MimeType, nil +} + +// ClearMimeType clears the value of the "mime_type" field. +func (m *BatchImageItemMutation) ClearMimeType() { + m.mime_type = nil + m.clearedFields[batchimageitem.FieldMimeType] = struct{}{} +} + +// MimeTypeCleared returns if the "mime_type" field was cleared in this mutation. +func (m *BatchImageItemMutation) MimeTypeCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldMimeType] + return ok +} + +// ResetMimeType resets all changes to the "mime_type" field. +func (m *BatchImageItemMutation) ResetMimeType() { + m.mime_type = nil + delete(m.clearedFields, batchimageitem.FieldMimeType) +} + +// SetFileExtension sets the "file_extension" field. +func (m *BatchImageItemMutation) SetFileExtension(s string) { + m.file_extension = &s +} + +// FileExtension returns the value of the "file_extension" field in the mutation. +func (m *BatchImageItemMutation) FileExtension() (r string, exists bool) { + v := m.file_extension + if v == nil { + return + } + return *v, true +} + +// OldFileExtension returns the old "file_extension" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldFileExtension(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldFileExtension is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldFileExtension requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldFileExtension: %w", err) + } + return oldValue.FileExtension, nil +} + +// ClearFileExtension clears the value of the "file_extension" field. +func (m *BatchImageItemMutation) ClearFileExtension() { + m.file_extension = nil + m.clearedFields[batchimageitem.FieldFileExtension] = struct{}{} +} + +// FileExtensionCleared returns if the "file_extension" field was cleared in this mutation. +func (m *BatchImageItemMutation) FileExtensionCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldFileExtension] + return ok +} + +// ResetFileExtension resets all changes to the "file_extension" field. +func (m *BatchImageItemMutation) ResetFileExtension() { + m.file_extension = nil + delete(m.clearedFields, batchimageitem.FieldFileExtension) +} + +// SetImageCount sets the "image_count" field. +func (m *BatchImageItemMutation) SetImageCount(i int) { + m.image_count = &i + m.addimage_count = nil +} + +// ImageCount returns the value of the "image_count" field in the mutation. +func (m *BatchImageItemMutation) ImageCount() (r int, exists bool) { + v := m.image_count + if v == nil { + return + } + return *v, true +} + +// OldImageCount returns the old "image_count" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldImageCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldImageCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldImageCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldImageCount: %w", err) + } + return oldValue.ImageCount, nil +} + +// AddImageCount adds i to the "image_count" field. +func (m *BatchImageItemMutation) AddImageCount(i int) { + if m.addimage_count != nil { + *m.addimage_count += i + } else { + m.addimage_count = &i + } +} + +// AddedImageCount returns the value that was added to the "image_count" field in this mutation. +func (m *BatchImageItemMutation) AddedImageCount() (r int, exists bool) { + v := m.addimage_count + if v == nil { + return + } + return *v, true +} + +// ResetImageCount resets all changes to the "image_count" field. +func (m *BatchImageItemMutation) ResetImageCount() { + m.image_count = nil + m.addimage_count = nil +} + +// SetErrorCode sets the "error_code" field. +func (m *BatchImageItemMutation) SetErrorCode(s string) { + m.error_code = &s +} + +// ErrorCode returns the value of the "error_code" field in the mutation. +func (m *BatchImageItemMutation) ErrorCode() (r string, exists bool) { + v := m.error_code + if v == nil { + return + } + return *v, true +} + +// OldErrorCode returns the old "error_code" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldErrorCode(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldErrorCode is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldErrorCode requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldErrorCode: %w", err) + } + return oldValue.ErrorCode, nil +} + +// ClearErrorCode clears the value of the "error_code" field. +func (m *BatchImageItemMutation) ClearErrorCode() { + m.error_code = nil + m.clearedFields[batchimageitem.FieldErrorCode] = struct{}{} +} + +// ErrorCodeCleared returns if the "error_code" field was cleared in this mutation. +func (m *BatchImageItemMutation) ErrorCodeCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldErrorCode] + return ok +} + +// ResetErrorCode resets all changes to the "error_code" field. +func (m *BatchImageItemMutation) ResetErrorCode() { + m.error_code = nil + delete(m.clearedFields, batchimageitem.FieldErrorCode) +} + +// SetErrorMessage sets the "error_message" field. +func (m *BatchImageItemMutation) SetErrorMessage(s string) { + m.error_message = &s +} + +// ErrorMessage returns the value of the "error_message" field in the mutation. +func (m *BatchImageItemMutation) ErrorMessage() (r string, exists bool) { + v := m.error_message + if v == nil { + return + } + return *v, true +} + +// OldErrorMessage returns the old "error_message" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldErrorMessage(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldErrorMessage is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldErrorMessage requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldErrorMessage: %w", err) + } + return oldValue.ErrorMessage, nil +} + +// ClearErrorMessage clears the value of the "error_message" field. +func (m *BatchImageItemMutation) ClearErrorMessage() { + m.error_message = nil + m.clearedFields[batchimageitem.FieldErrorMessage] = struct{}{} +} + +// ErrorMessageCleared returns if the "error_message" field was cleared in this mutation. +func (m *BatchImageItemMutation) ErrorMessageCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldErrorMessage] + return ok +} + +// ResetErrorMessage resets all changes to the "error_message" field. +func (m *BatchImageItemMutation) ResetErrorMessage() { + m.error_message = nil + delete(m.clearedFields, batchimageitem.FieldErrorMessage) +} + +// SetBilledAmount sets the "billed_amount" field. +func (m *BatchImageItemMutation) SetBilledAmount(f float64) { + m.billed_amount = &f + m.addbilled_amount = nil +} + +// BilledAmount returns the value of the "billed_amount" field in the mutation. +func (m *BatchImageItemMutation) BilledAmount() (r float64, exists bool) { + v := m.billed_amount + if v == nil { + return + } + return *v, true +} + +// OldBilledAmount returns the old "billed_amount" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldBilledAmount(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldBilledAmount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldBilledAmount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldBilledAmount: %w", err) + } + return oldValue.BilledAmount, nil +} + +// AddBilledAmount adds f to the "billed_amount" field. +func (m *BatchImageItemMutation) AddBilledAmount(f float64) { + if m.addbilled_amount != nil { + *m.addbilled_amount += f + } else { + m.addbilled_amount = &f + } +} + +// AddedBilledAmount returns the value that was added to the "billed_amount" field in this mutation. +func (m *BatchImageItemMutation) AddedBilledAmount() (r float64, exists bool) { + v := m.addbilled_amount + if v == nil { + return + } + return *v, true +} + +// ClearBilledAmount clears the value of the "billed_amount" field. +func (m *BatchImageItemMutation) ClearBilledAmount() { + m.billed_amount = nil + m.addbilled_amount = nil + m.clearedFields[batchimageitem.FieldBilledAmount] = struct{}{} +} + +// BilledAmountCleared returns if the "billed_amount" field was cleared in this mutation. +func (m *BatchImageItemMutation) BilledAmountCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldBilledAmount] + return ok +} + +// ResetBilledAmount resets all changes to the "billed_amount" field. +func (m *BatchImageItemMutation) ResetBilledAmount() { + m.billed_amount = nil + m.addbilled_amount = nil + delete(m.clearedFields, batchimageitem.FieldBilledAmount) +} + +// SetCreatedAt sets the "created_at" field. +func (m *BatchImageItemMutation) SetCreatedAt(t time.Time) { + m.created_at = &t +} + +// CreatedAt returns the value of the "created_at" field in the mutation. +func (m *BatchImageItemMutation) CreatedAt() (r time.Time, exists bool) { + v := m.created_at + if v == nil { + return + } + return *v, true +} + +// OldCreatedAt returns the old "created_at" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldCreatedAt(ctx context.Context) (v time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldCreatedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldCreatedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldCreatedAt: %w", err) + } + return oldValue.CreatedAt, nil +} + +// ResetCreatedAt resets all changes to the "created_at" field. +func (m *BatchImageItemMutation) ResetCreatedAt() { + m.created_at = nil +} + +// SetIndexedAt sets the "indexed_at" field. +func (m *BatchImageItemMutation) SetIndexedAt(t time.Time) { + m.indexed_at = &t +} + +// IndexedAt returns the value of the "indexed_at" field in the mutation. +func (m *BatchImageItemMutation) IndexedAt() (r time.Time, exists bool) { + v := m.indexed_at + if v == nil { + return + } + return *v, true +} + +// OldIndexedAt returns the old "indexed_at" field's value of the BatchImageItem entity. +// If the BatchImageItem object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageItemMutation) OldIndexedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldIndexedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldIndexedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldIndexedAt: %w", err) + } + return oldValue.IndexedAt, nil +} + +// ClearIndexedAt clears the value of the "indexed_at" field. +func (m *BatchImageItemMutation) ClearIndexedAt() { + m.indexed_at = nil + m.clearedFields[batchimageitem.FieldIndexedAt] = struct{}{} +} + +// IndexedAtCleared returns if the "indexed_at" field was cleared in this mutation. +func (m *BatchImageItemMutation) IndexedAtCleared() bool { + _, ok := m.clearedFields[batchimageitem.FieldIndexedAt] + return ok +} + +// ResetIndexedAt resets all changes to the "indexed_at" field. +func (m *BatchImageItemMutation) ResetIndexedAt() { + m.indexed_at = nil + delete(m.clearedFields, batchimageitem.FieldIndexedAt) +} + +// Where appends a list predicates to the BatchImageItemMutation builder. +func (m *BatchImageItemMutation) Where(ps ...predicate.BatchImageItem) { + m.predicates = append(m.predicates, ps...) +} + +// WhereP appends storage-level predicates to the BatchImageItemMutation builder. Using this method, +// users can use type-assertion to append predicates that do not depend on any generated package. +func (m *BatchImageItemMutation) WhereP(ps ...func(*sql.Selector)) { + p := make([]predicate.BatchImageItem, len(ps)) + for i := range ps { + p[i] = ps[i] + } + m.Where(p...) +} + +// Op returns the operation name. +func (m *BatchImageItemMutation) Op() Op { + return m.op +} + +// SetOp allows setting the mutation operation. +func (m *BatchImageItemMutation) SetOp(op Op) { + m.op = op +} + +// Type returns the node type of this mutation (BatchImageItem). +func (m *BatchImageItemMutation) Type() string { + return m.typ +} + +// Fields returns all fields that were changed during this mutation. Note that in +// order to get all numeric fields that were incremented/decremented, call +// AddedFields(). +func (m *BatchImageItemMutation) Fields() []string { + fields := make([]string, 0, 17) + if m.job_id != nil { + fields = append(fields, batchimageitem.FieldJobID) + } + if m.custom_id != nil { + fields = append(fields, batchimageitem.FieldCustomID) + } + if m.status != nil { + fields = append(fields, batchimageitem.FieldStatus) + } + if m.request_hash != nil { + fields = append(fields, batchimageitem.FieldRequestHash) + } + if m.prompt_preview != nil { + fields = append(fields, batchimageitem.FieldPromptPreview) + } + if m.provider_source_object != nil { + fields = append(fields, batchimageitem.FieldProviderSourceObject) + } + if m.source_line_number != nil { + fields = append(fields, batchimageitem.FieldSourceLineNumber) + } + if m.source_byte_offset != nil { + fields = append(fields, batchimageitem.FieldSourceByteOffset) + } + if m.source_byte_length != nil { + fields = append(fields, batchimageitem.FieldSourceByteLength) + } + if m.mime_type != nil { + fields = append(fields, batchimageitem.FieldMimeType) + } + if m.file_extension != nil { + fields = append(fields, batchimageitem.FieldFileExtension) + } + if m.image_count != nil { + fields = append(fields, batchimageitem.FieldImageCount) + } + if m.error_code != nil { + fields = append(fields, batchimageitem.FieldErrorCode) + } + if m.error_message != nil { + fields = append(fields, batchimageitem.FieldErrorMessage) + } + if m.billed_amount != nil { + fields = append(fields, batchimageitem.FieldBilledAmount) + } + if m.created_at != nil { + fields = append(fields, batchimageitem.FieldCreatedAt) + } + if m.indexed_at != nil { + fields = append(fields, batchimageitem.FieldIndexedAt) + } + return fields +} + +// Field returns the value of a field with the given name. The second boolean +// return value indicates that this field was not set, or was not defined in the +// schema. +func (m *BatchImageItemMutation) Field(name string) (ent.Value, bool) { + switch name { + case batchimageitem.FieldJobID: + return m.JobID() + case batchimageitem.FieldCustomID: + return m.CustomID() + case batchimageitem.FieldStatus: + return m.Status() + case batchimageitem.FieldRequestHash: + return m.RequestHash() + case batchimageitem.FieldPromptPreview: + return m.PromptPreview() + case batchimageitem.FieldProviderSourceObject: + return m.ProviderSourceObject() + case batchimageitem.FieldSourceLineNumber: + return m.SourceLineNumber() + case batchimageitem.FieldSourceByteOffset: + return m.SourceByteOffset() + case batchimageitem.FieldSourceByteLength: + return m.SourceByteLength() + case batchimageitem.FieldMimeType: + return m.MimeType() + case batchimageitem.FieldFileExtension: + return m.FileExtension() + case batchimageitem.FieldImageCount: + return m.ImageCount() + case batchimageitem.FieldErrorCode: + return m.ErrorCode() + case batchimageitem.FieldErrorMessage: + return m.ErrorMessage() + case batchimageitem.FieldBilledAmount: + return m.BilledAmount() + case batchimageitem.FieldCreatedAt: + return m.CreatedAt() + case batchimageitem.FieldIndexedAt: + return m.IndexedAt() + } + return nil, false +} + +// OldField returns the old value of the field from the database. An error is +// returned if the mutation operation is not UpdateOne, or the query to the +// database failed. +func (m *BatchImageItemMutation) OldField(ctx context.Context, name string) (ent.Value, error) { + switch name { + case batchimageitem.FieldJobID: + return m.OldJobID(ctx) + case batchimageitem.FieldCustomID: + return m.OldCustomID(ctx) + case batchimageitem.FieldStatus: + return m.OldStatus(ctx) + case batchimageitem.FieldRequestHash: + return m.OldRequestHash(ctx) + case batchimageitem.FieldPromptPreview: + return m.OldPromptPreview(ctx) + case batchimageitem.FieldProviderSourceObject: + return m.OldProviderSourceObject(ctx) + case batchimageitem.FieldSourceLineNumber: + return m.OldSourceLineNumber(ctx) + case batchimageitem.FieldSourceByteOffset: + return m.OldSourceByteOffset(ctx) + case batchimageitem.FieldSourceByteLength: + return m.OldSourceByteLength(ctx) + case batchimageitem.FieldMimeType: + return m.OldMimeType(ctx) + case batchimageitem.FieldFileExtension: + return m.OldFileExtension(ctx) + case batchimageitem.FieldImageCount: + return m.OldImageCount(ctx) + case batchimageitem.FieldErrorCode: + return m.OldErrorCode(ctx) + case batchimageitem.FieldErrorMessage: + return m.OldErrorMessage(ctx) + case batchimageitem.FieldBilledAmount: + return m.OldBilledAmount(ctx) + case batchimageitem.FieldCreatedAt: + return m.OldCreatedAt(ctx) + case batchimageitem.FieldIndexedAt: + return m.OldIndexedAt(ctx) + } + return nil, fmt.Errorf("unknown BatchImageItem field %s", name) +} + +// SetField sets the value of a field with the given name. It returns an error if +// the field is not defined in the schema, or if the type mismatched the field +// type. +func (m *BatchImageItemMutation) SetField(name string, value ent.Value) error { + switch name { + case batchimageitem.FieldJobID: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetJobID(v) + return nil + case batchimageitem.FieldCustomID: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetCustomID(v) + return nil + case batchimageitem.FieldStatus: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetStatus(v) + return nil + case batchimageitem.FieldRequestHash: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetRequestHash(v) + return nil + case batchimageitem.FieldPromptPreview: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetPromptPreview(v) + return nil + case batchimageitem.FieldProviderSourceObject: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetProviderSourceObject(v) + return nil + case batchimageitem.FieldSourceLineNumber: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetSourceLineNumber(v) + return nil + case batchimageitem.FieldSourceByteOffset: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetSourceByteOffset(v) + return nil + case batchimageitem.FieldSourceByteLength: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetSourceByteLength(v) + return nil + case batchimageitem.FieldMimeType: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetMimeType(v) + return nil + case batchimageitem.FieldFileExtension: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetFileExtension(v) + return nil + case batchimageitem.FieldImageCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetImageCount(v) + return nil + case batchimageitem.FieldErrorCode: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetErrorCode(v) + return nil + case batchimageitem.FieldErrorMessage: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetErrorMessage(v) + return nil + case batchimageitem.FieldBilledAmount: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetBilledAmount(v) + return nil + case batchimageitem.FieldCreatedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetCreatedAt(v) + return nil + case batchimageitem.FieldIndexedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetIndexedAt(v) + return nil + } + return fmt.Errorf("unknown BatchImageItem field %s", name) +} + +// AddedFields returns all numeric fields that were incremented/decremented during +// this mutation. +func (m *BatchImageItemMutation) AddedFields() []string { + var fields []string + if m.addsource_line_number != nil { + fields = append(fields, batchimageitem.FieldSourceLineNumber) + } + if m.addsource_byte_offset != nil { + fields = append(fields, batchimageitem.FieldSourceByteOffset) + } + if m.addsource_byte_length != nil { + fields = append(fields, batchimageitem.FieldSourceByteLength) + } + if m.addimage_count != nil { + fields = append(fields, batchimageitem.FieldImageCount) + } + if m.addbilled_amount != nil { + fields = append(fields, batchimageitem.FieldBilledAmount) + } + return fields +} + +// AddedField returns the numeric value that was incremented/decremented on a field +// with the given name. The second boolean return value indicates that this field +// was not set, or was not defined in the schema. +func (m *BatchImageItemMutation) AddedField(name string) (ent.Value, bool) { + switch name { + case batchimageitem.FieldSourceLineNumber: + return m.AddedSourceLineNumber() + case batchimageitem.FieldSourceByteOffset: + return m.AddedSourceByteOffset() + case batchimageitem.FieldSourceByteLength: + return m.AddedSourceByteLength() + case batchimageitem.FieldImageCount: + return m.AddedImageCount() + case batchimageitem.FieldBilledAmount: + return m.AddedBilledAmount() + } + return nil, false +} + +// AddField adds the value to the field with the given name. It returns an error if +// the field is not defined in the schema, or if the type mismatched the field +// type. +func (m *BatchImageItemMutation) AddField(name string, value ent.Value) error { + switch name { + case batchimageitem.FieldSourceLineNumber: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddSourceLineNumber(v) + return nil + case batchimageitem.FieldSourceByteOffset: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddSourceByteOffset(v) + return nil + case batchimageitem.FieldSourceByteLength: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddSourceByteLength(v) + return nil + case batchimageitem.FieldImageCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddImageCount(v) + return nil + case batchimageitem.FieldBilledAmount: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddBilledAmount(v) + return nil + } + return fmt.Errorf("unknown BatchImageItem numeric field %s", name) +} + +// ClearedFields returns all nullable fields that were cleared during this +// mutation. +func (m *BatchImageItemMutation) ClearedFields() []string { + var fields []string + if m.FieldCleared(batchimageitem.FieldRequestHash) { + fields = append(fields, batchimageitem.FieldRequestHash) + } + if m.FieldCleared(batchimageitem.FieldPromptPreview) { + fields = append(fields, batchimageitem.FieldPromptPreview) + } + if m.FieldCleared(batchimageitem.FieldProviderSourceObject) { + fields = append(fields, batchimageitem.FieldProviderSourceObject) + } + if m.FieldCleared(batchimageitem.FieldSourceLineNumber) { + fields = append(fields, batchimageitem.FieldSourceLineNumber) + } + if m.FieldCleared(batchimageitem.FieldSourceByteOffset) { + fields = append(fields, batchimageitem.FieldSourceByteOffset) + } + if m.FieldCleared(batchimageitem.FieldSourceByteLength) { + fields = append(fields, batchimageitem.FieldSourceByteLength) + } + if m.FieldCleared(batchimageitem.FieldMimeType) { + fields = append(fields, batchimageitem.FieldMimeType) + } + if m.FieldCleared(batchimageitem.FieldFileExtension) { + fields = append(fields, batchimageitem.FieldFileExtension) + } + if m.FieldCleared(batchimageitem.FieldErrorCode) { + fields = append(fields, batchimageitem.FieldErrorCode) + } + if m.FieldCleared(batchimageitem.FieldErrorMessage) { + fields = append(fields, batchimageitem.FieldErrorMessage) + } + if m.FieldCleared(batchimageitem.FieldBilledAmount) { + fields = append(fields, batchimageitem.FieldBilledAmount) + } + if m.FieldCleared(batchimageitem.FieldIndexedAt) { + fields = append(fields, batchimageitem.FieldIndexedAt) + } + return fields +} + +// FieldCleared returns a boolean indicating if a field with the given name was +// cleared in this mutation. +func (m *BatchImageItemMutation) FieldCleared(name string) bool { + _, ok := m.clearedFields[name] + return ok +} + +// ClearField clears the value of the field with the given name. It returns an +// error if the field is not defined in the schema. +func (m *BatchImageItemMutation) ClearField(name string) error { + switch name { + case batchimageitem.FieldRequestHash: + m.ClearRequestHash() + return nil + case batchimageitem.FieldPromptPreview: + m.ClearPromptPreview() + return nil + case batchimageitem.FieldProviderSourceObject: + m.ClearProviderSourceObject() + return nil + case batchimageitem.FieldSourceLineNumber: + m.ClearSourceLineNumber() + return nil + case batchimageitem.FieldSourceByteOffset: + m.ClearSourceByteOffset() + return nil + case batchimageitem.FieldSourceByteLength: + m.ClearSourceByteLength() + return nil + case batchimageitem.FieldMimeType: + m.ClearMimeType() + return nil + case batchimageitem.FieldFileExtension: + m.ClearFileExtension() + return nil + case batchimageitem.FieldErrorCode: + m.ClearErrorCode() + return nil + case batchimageitem.FieldErrorMessage: + m.ClearErrorMessage() + return nil + case batchimageitem.FieldBilledAmount: + m.ClearBilledAmount() + return nil + case batchimageitem.FieldIndexedAt: + m.ClearIndexedAt() + return nil + } + return fmt.Errorf("unknown BatchImageItem nullable field %s", name) +} + +// ResetField resets all changes in the mutation for the field with the given name. +// It returns an error if the field is not defined in the schema. +func (m *BatchImageItemMutation) ResetField(name string) error { + switch name { + case batchimageitem.FieldJobID: + m.ResetJobID() + return nil + case batchimageitem.FieldCustomID: + m.ResetCustomID() + return nil + case batchimageitem.FieldStatus: + m.ResetStatus() + return nil + case batchimageitem.FieldRequestHash: + m.ResetRequestHash() + return nil + case batchimageitem.FieldPromptPreview: + m.ResetPromptPreview() + return nil + case batchimageitem.FieldProviderSourceObject: + m.ResetProviderSourceObject() + return nil + case batchimageitem.FieldSourceLineNumber: + m.ResetSourceLineNumber() + return nil + case batchimageitem.FieldSourceByteOffset: + m.ResetSourceByteOffset() + return nil + case batchimageitem.FieldSourceByteLength: + m.ResetSourceByteLength() + return nil + case batchimageitem.FieldMimeType: + m.ResetMimeType() + return nil + case batchimageitem.FieldFileExtension: + m.ResetFileExtension() + return nil + case batchimageitem.FieldImageCount: + m.ResetImageCount() + return nil + case batchimageitem.FieldErrorCode: + m.ResetErrorCode() + return nil + case batchimageitem.FieldErrorMessage: + m.ResetErrorMessage() + return nil + case batchimageitem.FieldBilledAmount: + m.ResetBilledAmount() + return nil + case batchimageitem.FieldCreatedAt: + m.ResetCreatedAt() + return nil + case batchimageitem.FieldIndexedAt: + m.ResetIndexedAt() + return nil + } + return fmt.Errorf("unknown BatchImageItem field %s", name) +} + +// AddedEdges returns all edge names that were set/added in this mutation. +func (m *BatchImageItemMutation) AddedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// AddedIDs returns all IDs (to other nodes) that were added for the given edge +// name in this mutation. +func (m *BatchImageItemMutation) AddedIDs(name string) []ent.Value { + return nil +} + +// RemovedEdges returns all edge names that were removed in this mutation. +func (m *BatchImageItemMutation) RemovedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// RemovedIDs returns all IDs (to other nodes) that were removed for the edge with +// the given name in this mutation. +func (m *BatchImageItemMutation) RemovedIDs(name string) []ent.Value { + return nil +} + +// ClearedEdges returns all edge names that were cleared in this mutation. +func (m *BatchImageItemMutation) ClearedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// EdgeCleared returns a boolean which indicates if the edge with the given name +// was cleared in this mutation. +func (m *BatchImageItemMutation) EdgeCleared(name string) bool { + return false +} + +// ClearEdge clears the value of the edge with the given name. It returns an error +// if that edge is not defined in the schema. +func (m *BatchImageItemMutation) ClearEdge(name string) error { + return fmt.Errorf("unknown BatchImageItem unique edge %s", name) +} + +// ResetEdge resets all changes to the edge with the given name in this mutation. +// It returns an error if the edge is not defined in the schema. +func (m *BatchImageItemMutation) ResetEdge(name string) error { + return fmt.Errorf("unknown BatchImageItem edge %s", name) +} + +// BatchImageJobMutation represents an operation that mutates the BatchImageJob nodes in the graph. +type BatchImageJobMutation struct { + config + op Op + typ string + id *int64 + batch_id *string + user_id *int64 + adduser_id *int64 + api_key_id *int64 + addapi_key_id *int64 + account_id *int64 + addaccount_id *int64 + provider *string + model *string + task_name *string + status *string + provider_job_name *string + provider_input_ref *string + provider_output_ref *string + gcs_input_uri *string + gcs_output_uri *string + item_count *int + additem_count *int + success_count *int + addsuccess_count *int + fail_count *int + addfail_count *int + cancelled_count *int + addcancelled_count *int + estimated_cost *float64 + addestimated_cost *float64 + hold_amount *float64 + addhold_amount *float64 + actual_cost *float64 + addactual_cost *float64 + currency *string + hold_id *string + idempotency_key *string + request_hash *string + manifest_hash *string + retry_count *int + addretry_count *int + version *int + addversion *int + output_expires_at *time.Time + input_deleted_at *time.Time + output_deleted_at *time.Time + downloaded_at *time.Time + user_deleted_at *time.Time + last_error_code *string + last_error_message *string + created_at *time.Time + updated_at *time.Time + submitted_at *time.Time + started_at *time.Time + finished_at *time.Time + settled_at *time.Time + clearedFields map[string]struct{} + done bool + oldValue func(context.Context) (*BatchImageJob, error) + predicates []predicate.BatchImageJob +} + +var _ ent.Mutation = (*BatchImageJobMutation)(nil) + +// batchimagejobOption allows management of the mutation configuration using functional options. +type batchimagejobOption func(*BatchImageJobMutation) + +// newBatchImageJobMutation creates new mutation for the BatchImageJob entity. +func newBatchImageJobMutation(c config, op Op, opts ...batchimagejobOption) *BatchImageJobMutation { + m := &BatchImageJobMutation{ + config: c, + op: op, + typ: TypeBatchImageJob, + clearedFields: make(map[string]struct{}), + } + for _, opt := range opts { + opt(m) + } + return m +} + +// withBatchImageJobID sets the ID field of the mutation. +func withBatchImageJobID(id int64) batchimagejobOption { + return func(m *BatchImageJobMutation) { + var ( + err error + once sync.Once + value *BatchImageJob + ) + m.oldValue = func(ctx context.Context) (*BatchImageJob, error) { + once.Do(func() { + if m.done { + err = errors.New("querying old values post mutation is not allowed") + } else { + value, err = m.Client().BatchImageJob.Get(ctx, id) + } + }) + return value, err + } + m.id = &id + } +} + +// withBatchImageJob sets the old BatchImageJob of the mutation. +func withBatchImageJob(node *BatchImageJob) batchimagejobOption { + return func(m *BatchImageJobMutation) { + m.oldValue = func(context.Context) (*BatchImageJob, error) { + return node, nil + } + m.id = &node.ID + } +} + +// Client returns a new `ent.Client` from the mutation. If the mutation was +// executed in a transaction (ent.Tx), a transactional client is returned. +func (m BatchImageJobMutation) Client() *Client { + client := &Client{config: m.config} + client.init() + return client +} + +// Tx returns an `ent.Tx` for mutations that were executed in transactions; +// it returns an error otherwise. +func (m BatchImageJobMutation) Tx() (*Tx, error) { + if _, ok := m.driver.(*txDriver); !ok { + return nil, errors.New("ent: mutation is not running in a transaction") + } + tx := &Tx{config: m.config} + tx.init() + return tx, nil +} + +// ID returns the ID value in the mutation. Note that the ID is only available +// if it was provided to the builder or after it was returned from the database. +func (m *BatchImageJobMutation) ID() (id int64, exists bool) { + if m.id == nil { + return + } + return *m.id, true +} + +// IDs queries the database and returns the entity ids that match the mutation's predicate. +// That means, if the mutation is applied within a transaction with an isolation level such +// as sql.LevelSerializable, the returned ids match the ids of the rows that will be updated +// or updated by the mutation. +func (m *BatchImageJobMutation) IDs(ctx context.Context) ([]int64, error) { + switch { + case m.op.Is(OpUpdateOne | OpDeleteOne): + id, exists := m.ID() + if exists { + return []int64{id}, nil + } + fallthrough + case m.op.Is(OpUpdate | OpDelete): + return m.Client().BatchImageJob.Query().Where(m.predicates...).IDs(ctx) + default: + return nil, fmt.Errorf("IDs is not allowed on %s operations", m.op) + } +} + +// SetBatchID sets the "batch_id" field. +func (m *BatchImageJobMutation) SetBatchID(s string) { + m.batch_id = &s +} + +// BatchID returns the value of the "batch_id" field in the mutation. +func (m *BatchImageJobMutation) BatchID() (r string, exists bool) { + v := m.batch_id + if v == nil { + return + } + return *v, true +} + +// OldBatchID returns the old "batch_id" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldBatchID(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldBatchID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldBatchID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldBatchID: %w", err) + } + return oldValue.BatchID, nil +} + +// ResetBatchID resets all changes to the "batch_id" field. +func (m *BatchImageJobMutation) ResetBatchID() { + m.batch_id = nil +} + +// SetUserID sets the "user_id" field. +func (m *BatchImageJobMutation) SetUserID(i int64) { + m.user_id = &i + m.adduser_id = nil +} + +// UserID returns the value of the "user_id" field in the mutation. +func (m *BatchImageJobMutation) UserID() (r int64, exists bool) { + v := m.user_id + if v == nil { + return + } + return *v, true +} + +// OldUserID returns the old "user_id" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldUserID(ctx context.Context) (v int64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldUserID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldUserID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldUserID: %w", err) + } + return oldValue.UserID, nil +} + +// AddUserID adds i to the "user_id" field. +func (m *BatchImageJobMutation) AddUserID(i int64) { + if m.adduser_id != nil { + *m.adduser_id += i + } else { + m.adduser_id = &i + } +} + +// AddedUserID returns the value that was added to the "user_id" field in this mutation. +func (m *BatchImageJobMutation) AddedUserID() (r int64, exists bool) { + v := m.adduser_id + if v == nil { + return + } + return *v, true +} + +// ResetUserID resets all changes to the "user_id" field. +func (m *BatchImageJobMutation) ResetUserID() { + m.user_id = nil + m.adduser_id = nil +} + +// SetAPIKeyID sets the "api_key_id" field. +func (m *BatchImageJobMutation) SetAPIKeyID(i int64) { + m.api_key_id = &i + m.addapi_key_id = nil +} + +// APIKeyID returns the value of the "api_key_id" field in the mutation. +func (m *BatchImageJobMutation) APIKeyID() (r int64, exists bool) { + v := m.api_key_id + if v == nil { + return + } + return *v, true +} + +// OldAPIKeyID returns the old "api_key_id" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldAPIKeyID(ctx context.Context) (v *int64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldAPIKeyID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldAPIKeyID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldAPIKeyID: %w", err) + } + return oldValue.APIKeyID, nil +} + +// AddAPIKeyID adds i to the "api_key_id" field. +func (m *BatchImageJobMutation) AddAPIKeyID(i int64) { + if m.addapi_key_id != nil { + *m.addapi_key_id += i + } else { + m.addapi_key_id = &i + } +} + +// AddedAPIKeyID returns the value that was added to the "api_key_id" field in this mutation. +func (m *BatchImageJobMutation) AddedAPIKeyID() (r int64, exists bool) { + v := m.addapi_key_id + if v == nil { + return + } + return *v, true +} + +// ClearAPIKeyID clears the value of the "api_key_id" field. +func (m *BatchImageJobMutation) ClearAPIKeyID() { + m.api_key_id = nil + m.addapi_key_id = nil + m.clearedFields[batchimagejob.FieldAPIKeyID] = struct{}{} +} + +// APIKeyIDCleared returns if the "api_key_id" field was cleared in this mutation. +func (m *BatchImageJobMutation) APIKeyIDCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldAPIKeyID] + return ok +} + +// ResetAPIKeyID resets all changes to the "api_key_id" field. +func (m *BatchImageJobMutation) ResetAPIKeyID() { + m.api_key_id = nil + m.addapi_key_id = nil + delete(m.clearedFields, batchimagejob.FieldAPIKeyID) +} + +// SetAccountID sets the "account_id" field. +func (m *BatchImageJobMutation) SetAccountID(i int64) { + m.account_id = &i + m.addaccount_id = nil +} + +// AccountID returns the value of the "account_id" field in the mutation. +func (m *BatchImageJobMutation) AccountID() (r int64, exists bool) { + v := m.account_id + if v == nil { + return + } + return *v, true +} + +// OldAccountID returns the old "account_id" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldAccountID(ctx context.Context) (v *int64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldAccountID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldAccountID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldAccountID: %w", err) + } + return oldValue.AccountID, nil +} + +// AddAccountID adds i to the "account_id" field. +func (m *BatchImageJobMutation) AddAccountID(i int64) { + if m.addaccount_id != nil { + *m.addaccount_id += i + } else { + m.addaccount_id = &i + } +} + +// AddedAccountID returns the value that was added to the "account_id" field in this mutation. +func (m *BatchImageJobMutation) AddedAccountID() (r int64, exists bool) { + v := m.addaccount_id + if v == nil { + return + } + return *v, true +} + +// ClearAccountID clears the value of the "account_id" field. +func (m *BatchImageJobMutation) ClearAccountID() { + m.account_id = nil + m.addaccount_id = nil + m.clearedFields[batchimagejob.FieldAccountID] = struct{}{} +} + +// AccountIDCleared returns if the "account_id" field was cleared in this mutation. +func (m *BatchImageJobMutation) AccountIDCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldAccountID] + return ok +} + +// ResetAccountID resets all changes to the "account_id" field. +func (m *BatchImageJobMutation) ResetAccountID() { + m.account_id = nil + m.addaccount_id = nil + delete(m.clearedFields, batchimagejob.FieldAccountID) +} + +// SetProvider sets the "provider" field. +func (m *BatchImageJobMutation) SetProvider(s string) { + m.provider = &s +} + +// Provider returns the value of the "provider" field in the mutation. +func (m *BatchImageJobMutation) Provider() (r string, exists bool) { + v := m.provider + if v == nil { + return + } + return *v, true +} + +// OldProvider returns the old "provider" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldProvider(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldProvider is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldProvider requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldProvider: %w", err) + } + return oldValue.Provider, nil +} + +// ResetProvider resets all changes to the "provider" field. +func (m *BatchImageJobMutation) ResetProvider() { + m.provider = nil +} + +// SetModel sets the "model" field. +func (m *BatchImageJobMutation) SetModel(s string) { + m.model = &s +} + +// Model returns the value of the "model" field in the mutation. +func (m *BatchImageJobMutation) Model() (r string, exists bool) { + v := m.model + if v == nil { + return + } + return *v, true +} + +// OldModel returns the old "model" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldModel(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldModel is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldModel requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldModel: %w", err) + } + return oldValue.Model, nil +} + +// ResetModel resets all changes to the "model" field. +func (m *BatchImageJobMutation) ResetModel() { + m.model = nil +} + +// SetTaskName sets the "task_name" field. +func (m *BatchImageJobMutation) SetTaskName(s string) { + m.task_name = &s +} + +// TaskName returns the value of the "task_name" field in the mutation. +func (m *BatchImageJobMutation) TaskName() (r string, exists bool) { + v := m.task_name + if v == nil { + return + } + return *v, true +} + +// OldTaskName returns the old "task_name" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldTaskName(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldTaskName is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldTaskName requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldTaskName: %w", err) + } + return oldValue.TaskName, nil +} + +// ResetTaskName resets all changes to the "task_name" field. +func (m *BatchImageJobMutation) ResetTaskName() { + m.task_name = nil +} + +// SetStatus sets the "status" field. +func (m *BatchImageJobMutation) SetStatus(s string) { + m.status = &s +} + +// Status returns the value of the "status" field in the mutation. +func (m *BatchImageJobMutation) Status() (r string, exists bool) { + v := m.status + if v == nil { + return + } + return *v, true +} + +// OldStatus returns the old "status" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldStatus(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldStatus is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldStatus requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldStatus: %w", err) + } + return oldValue.Status, nil +} + +// ResetStatus resets all changes to the "status" field. +func (m *BatchImageJobMutation) ResetStatus() { + m.status = nil +} + +// SetProviderJobName sets the "provider_job_name" field. +func (m *BatchImageJobMutation) SetProviderJobName(s string) { + m.provider_job_name = &s +} + +// ProviderJobName returns the value of the "provider_job_name" field in the mutation. +func (m *BatchImageJobMutation) ProviderJobName() (r string, exists bool) { + v := m.provider_job_name + if v == nil { + return + } + return *v, true +} + +// OldProviderJobName returns the old "provider_job_name" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldProviderJobName(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldProviderJobName is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldProviderJobName requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldProviderJobName: %w", err) + } + return oldValue.ProviderJobName, nil +} + +// ClearProviderJobName clears the value of the "provider_job_name" field. +func (m *BatchImageJobMutation) ClearProviderJobName() { + m.provider_job_name = nil + m.clearedFields[batchimagejob.FieldProviderJobName] = struct{}{} +} + +// ProviderJobNameCleared returns if the "provider_job_name" field was cleared in this mutation. +func (m *BatchImageJobMutation) ProviderJobNameCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldProviderJobName] + return ok +} + +// ResetProviderJobName resets all changes to the "provider_job_name" field. +func (m *BatchImageJobMutation) ResetProviderJobName() { + m.provider_job_name = nil + delete(m.clearedFields, batchimagejob.FieldProviderJobName) +} + +// SetProviderInputRef sets the "provider_input_ref" field. +func (m *BatchImageJobMutation) SetProviderInputRef(s string) { + m.provider_input_ref = &s +} + +// ProviderInputRef returns the value of the "provider_input_ref" field in the mutation. +func (m *BatchImageJobMutation) ProviderInputRef() (r string, exists bool) { + v := m.provider_input_ref + if v == nil { + return + } + return *v, true +} + +// OldProviderInputRef returns the old "provider_input_ref" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldProviderInputRef(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldProviderInputRef is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldProviderInputRef requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldProviderInputRef: %w", err) + } + return oldValue.ProviderInputRef, nil +} + +// ClearProviderInputRef clears the value of the "provider_input_ref" field. +func (m *BatchImageJobMutation) ClearProviderInputRef() { + m.provider_input_ref = nil + m.clearedFields[batchimagejob.FieldProviderInputRef] = struct{}{} +} + +// ProviderInputRefCleared returns if the "provider_input_ref" field was cleared in this mutation. +func (m *BatchImageJobMutation) ProviderInputRefCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldProviderInputRef] + return ok +} + +// ResetProviderInputRef resets all changes to the "provider_input_ref" field. +func (m *BatchImageJobMutation) ResetProviderInputRef() { + m.provider_input_ref = nil + delete(m.clearedFields, batchimagejob.FieldProviderInputRef) +} + +// SetProviderOutputRef sets the "provider_output_ref" field. +func (m *BatchImageJobMutation) SetProviderOutputRef(s string) { + m.provider_output_ref = &s +} + +// ProviderOutputRef returns the value of the "provider_output_ref" field in the mutation. +func (m *BatchImageJobMutation) ProviderOutputRef() (r string, exists bool) { + v := m.provider_output_ref + if v == nil { + return + } + return *v, true +} + +// OldProviderOutputRef returns the old "provider_output_ref" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldProviderOutputRef(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldProviderOutputRef is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldProviderOutputRef requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldProviderOutputRef: %w", err) + } + return oldValue.ProviderOutputRef, nil +} + +// ClearProviderOutputRef clears the value of the "provider_output_ref" field. +func (m *BatchImageJobMutation) ClearProviderOutputRef() { + m.provider_output_ref = nil + m.clearedFields[batchimagejob.FieldProviderOutputRef] = struct{}{} +} + +// ProviderOutputRefCleared returns if the "provider_output_ref" field was cleared in this mutation. +func (m *BatchImageJobMutation) ProviderOutputRefCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldProviderOutputRef] + return ok +} + +// ResetProviderOutputRef resets all changes to the "provider_output_ref" field. +func (m *BatchImageJobMutation) ResetProviderOutputRef() { + m.provider_output_ref = nil + delete(m.clearedFields, batchimagejob.FieldProviderOutputRef) +} + +// SetGcsInputURI sets the "gcs_input_uri" field. +func (m *BatchImageJobMutation) SetGcsInputURI(s string) { + m.gcs_input_uri = &s +} + +// GcsInputURI returns the value of the "gcs_input_uri" field in the mutation. +func (m *BatchImageJobMutation) GcsInputURI() (r string, exists bool) { + v := m.gcs_input_uri + if v == nil { + return + } + return *v, true +} + +// OldGcsInputURI returns the old "gcs_input_uri" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldGcsInputURI(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldGcsInputURI is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldGcsInputURI requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldGcsInputURI: %w", err) + } + return oldValue.GcsInputURI, nil +} + +// ClearGcsInputURI clears the value of the "gcs_input_uri" field. +func (m *BatchImageJobMutation) ClearGcsInputURI() { + m.gcs_input_uri = nil + m.clearedFields[batchimagejob.FieldGcsInputURI] = struct{}{} +} + +// GcsInputURICleared returns if the "gcs_input_uri" field was cleared in this mutation. +func (m *BatchImageJobMutation) GcsInputURICleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldGcsInputURI] + return ok +} + +// ResetGcsInputURI resets all changes to the "gcs_input_uri" field. +func (m *BatchImageJobMutation) ResetGcsInputURI() { + m.gcs_input_uri = nil + delete(m.clearedFields, batchimagejob.FieldGcsInputURI) +} + +// SetGcsOutputURI sets the "gcs_output_uri" field. +func (m *BatchImageJobMutation) SetGcsOutputURI(s string) { + m.gcs_output_uri = &s +} + +// GcsOutputURI returns the value of the "gcs_output_uri" field in the mutation. +func (m *BatchImageJobMutation) GcsOutputURI() (r string, exists bool) { + v := m.gcs_output_uri + if v == nil { + return + } + return *v, true +} + +// OldGcsOutputURI returns the old "gcs_output_uri" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldGcsOutputURI(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldGcsOutputURI is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldGcsOutputURI requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldGcsOutputURI: %w", err) + } + return oldValue.GcsOutputURI, nil +} + +// ClearGcsOutputURI clears the value of the "gcs_output_uri" field. +func (m *BatchImageJobMutation) ClearGcsOutputURI() { + m.gcs_output_uri = nil + m.clearedFields[batchimagejob.FieldGcsOutputURI] = struct{}{} +} + +// GcsOutputURICleared returns if the "gcs_output_uri" field was cleared in this mutation. +func (m *BatchImageJobMutation) GcsOutputURICleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldGcsOutputURI] + return ok +} + +// ResetGcsOutputURI resets all changes to the "gcs_output_uri" field. +func (m *BatchImageJobMutation) ResetGcsOutputURI() { + m.gcs_output_uri = nil + delete(m.clearedFields, batchimagejob.FieldGcsOutputURI) +} + +// SetItemCount sets the "item_count" field. +func (m *BatchImageJobMutation) SetItemCount(i int) { + m.item_count = &i + m.additem_count = nil +} + +// ItemCount returns the value of the "item_count" field in the mutation. +func (m *BatchImageJobMutation) ItemCount() (r int, exists bool) { + v := m.item_count + if v == nil { + return + } + return *v, true +} + +// OldItemCount returns the old "item_count" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldItemCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldItemCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldItemCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldItemCount: %w", err) + } + return oldValue.ItemCount, nil +} + +// AddItemCount adds i to the "item_count" field. +func (m *BatchImageJobMutation) AddItemCount(i int) { + if m.additem_count != nil { + *m.additem_count += i + } else { + m.additem_count = &i + } +} + +// AddedItemCount returns the value that was added to the "item_count" field in this mutation. +func (m *BatchImageJobMutation) AddedItemCount() (r int, exists bool) { + v := m.additem_count + if v == nil { + return + } + return *v, true +} + +// ResetItemCount resets all changes to the "item_count" field. +func (m *BatchImageJobMutation) ResetItemCount() { + m.item_count = nil + m.additem_count = nil +} + +// SetSuccessCount sets the "success_count" field. +func (m *BatchImageJobMutation) SetSuccessCount(i int) { + m.success_count = &i + m.addsuccess_count = nil +} + +// SuccessCount returns the value of the "success_count" field in the mutation. +func (m *BatchImageJobMutation) SuccessCount() (r int, exists bool) { + v := m.success_count + if v == nil { + return + } + return *v, true +} + +// OldSuccessCount returns the old "success_count" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldSuccessCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldSuccessCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldSuccessCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldSuccessCount: %w", err) + } + return oldValue.SuccessCount, nil +} + +// AddSuccessCount adds i to the "success_count" field. +func (m *BatchImageJobMutation) AddSuccessCount(i int) { + if m.addsuccess_count != nil { + *m.addsuccess_count += i + } else { + m.addsuccess_count = &i + } +} + +// AddedSuccessCount returns the value that was added to the "success_count" field in this mutation. +func (m *BatchImageJobMutation) AddedSuccessCount() (r int, exists bool) { + v := m.addsuccess_count + if v == nil { + return + } + return *v, true +} + +// ResetSuccessCount resets all changes to the "success_count" field. +func (m *BatchImageJobMutation) ResetSuccessCount() { + m.success_count = nil + m.addsuccess_count = nil +} + +// SetFailCount sets the "fail_count" field. +func (m *BatchImageJobMutation) SetFailCount(i int) { + m.fail_count = &i + m.addfail_count = nil +} + +// FailCount returns the value of the "fail_count" field in the mutation. +func (m *BatchImageJobMutation) FailCount() (r int, exists bool) { + v := m.fail_count + if v == nil { + return + } + return *v, true +} + +// OldFailCount returns the old "fail_count" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldFailCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldFailCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldFailCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldFailCount: %w", err) + } + return oldValue.FailCount, nil +} + +// AddFailCount adds i to the "fail_count" field. +func (m *BatchImageJobMutation) AddFailCount(i int) { + if m.addfail_count != nil { + *m.addfail_count += i + } else { + m.addfail_count = &i + } +} + +// AddedFailCount returns the value that was added to the "fail_count" field in this mutation. +func (m *BatchImageJobMutation) AddedFailCount() (r int, exists bool) { + v := m.addfail_count + if v == nil { + return + } + return *v, true +} + +// ResetFailCount resets all changes to the "fail_count" field. +func (m *BatchImageJobMutation) ResetFailCount() { + m.fail_count = nil + m.addfail_count = nil +} + +// SetCancelledCount sets the "cancelled_count" field. +func (m *BatchImageJobMutation) SetCancelledCount(i int) { + m.cancelled_count = &i + m.addcancelled_count = nil +} + +// CancelledCount returns the value of the "cancelled_count" field in the mutation. +func (m *BatchImageJobMutation) CancelledCount() (r int, exists bool) { + v := m.cancelled_count + if v == nil { + return + } + return *v, true +} + +// OldCancelledCount returns the old "cancelled_count" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldCancelledCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldCancelledCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldCancelledCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldCancelledCount: %w", err) + } + return oldValue.CancelledCount, nil +} + +// AddCancelledCount adds i to the "cancelled_count" field. +func (m *BatchImageJobMutation) AddCancelledCount(i int) { + if m.addcancelled_count != nil { + *m.addcancelled_count += i + } else { + m.addcancelled_count = &i + } +} + +// AddedCancelledCount returns the value that was added to the "cancelled_count" field in this mutation. +func (m *BatchImageJobMutation) AddedCancelledCount() (r int, exists bool) { + v := m.addcancelled_count + if v == nil { + return + } + return *v, true +} + +// ResetCancelledCount resets all changes to the "cancelled_count" field. +func (m *BatchImageJobMutation) ResetCancelledCount() { + m.cancelled_count = nil + m.addcancelled_count = nil +} + +// SetEstimatedCost sets the "estimated_cost" field. +func (m *BatchImageJobMutation) SetEstimatedCost(f float64) { + m.estimated_cost = &f + m.addestimated_cost = nil +} + +// EstimatedCost returns the value of the "estimated_cost" field in the mutation. +func (m *BatchImageJobMutation) EstimatedCost() (r float64, exists bool) { + v := m.estimated_cost + if v == nil { + return + } + return *v, true +} + +// OldEstimatedCost returns the old "estimated_cost" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldEstimatedCost(ctx context.Context) (v float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldEstimatedCost is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldEstimatedCost requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldEstimatedCost: %w", err) + } + return oldValue.EstimatedCost, nil +} + +// AddEstimatedCost adds f to the "estimated_cost" field. +func (m *BatchImageJobMutation) AddEstimatedCost(f float64) { + if m.addestimated_cost != nil { + *m.addestimated_cost += f + } else { + m.addestimated_cost = &f + } +} + +// AddedEstimatedCost returns the value that was added to the "estimated_cost" field in this mutation. +func (m *BatchImageJobMutation) AddedEstimatedCost() (r float64, exists bool) { + v := m.addestimated_cost + if v == nil { + return + } + return *v, true +} + +// ResetEstimatedCost resets all changes to the "estimated_cost" field. +func (m *BatchImageJobMutation) ResetEstimatedCost() { + m.estimated_cost = nil + m.addestimated_cost = nil +} + +// SetHoldAmount sets the "hold_amount" field. +func (m *BatchImageJobMutation) SetHoldAmount(f float64) { + m.hold_amount = &f + m.addhold_amount = nil +} + +// HoldAmount returns the value of the "hold_amount" field in the mutation. +func (m *BatchImageJobMutation) HoldAmount() (r float64, exists bool) { + v := m.hold_amount + if v == nil { + return + } + return *v, true +} + +// OldHoldAmount returns the old "hold_amount" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldHoldAmount(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldHoldAmount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldHoldAmount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldHoldAmount: %w", err) + } + return oldValue.HoldAmount, nil +} + +// AddHoldAmount adds f to the "hold_amount" field. +func (m *BatchImageJobMutation) AddHoldAmount(f float64) { + if m.addhold_amount != nil { + *m.addhold_amount += f + } else { + m.addhold_amount = &f + } +} + +// AddedHoldAmount returns the value that was added to the "hold_amount" field in this mutation. +func (m *BatchImageJobMutation) AddedHoldAmount() (r float64, exists bool) { + v := m.addhold_amount + if v == nil { + return + } + return *v, true +} + +// ClearHoldAmount clears the value of the "hold_amount" field. +func (m *BatchImageJobMutation) ClearHoldAmount() { + m.hold_amount = nil + m.addhold_amount = nil + m.clearedFields[batchimagejob.FieldHoldAmount] = struct{}{} +} + +// HoldAmountCleared returns if the "hold_amount" field was cleared in this mutation. +func (m *BatchImageJobMutation) HoldAmountCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldHoldAmount] + return ok +} + +// ResetHoldAmount resets all changes to the "hold_amount" field. +func (m *BatchImageJobMutation) ResetHoldAmount() { + m.hold_amount = nil + m.addhold_amount = nil + delete(m.clearedFields, batchimagejob.FieldHoldAmount) +} + +// SetActualCost sets the "actual_cost" field. +func (m *BatchImageJobMutation) SetActualCost(f float64) { + m.actual_cost = &f + m.addactual_cost = nil +} + +// ActualCost returns the value of the "actual_cost" field in the mutation. +func (m *BatchImageJobMutation) ActualCost() (r float64, exists bool) { + v := m.actual_cost + if v == nil { + return + } + return *v, true +} + +// OldActualCost returns the old "actual_cost" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldActualCost(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldActualCost is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldActualCost requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldActualCost: %w", err) + } + return oldValue.ActualCost, nil +} + +// AddActualCost adds f to the "actual_cost" field. +func (m *BatchImageJobMutation) AddActualCost(f float64) { + if m.addactual_cost != nil { + *m.addactual_cost += f + } else { + m.addactual_cost = &f + } +} + +// AddedActualCost returns the value that was added to the "actual_cost" field in this mutation. +func (m *BatchImageJobMutation) AddedActualCost() (r float64, exists bool) { + v := m.addactual_cost + if v == nil { + return + } + return *v, true +} + +// ClearActualCost clears the value of the "actual_cost" field. +func (m *BatchImageJobMutation) ClearActualCost() { + m.actual_cost = nil + m.addactual_cost = nil + m.clearedFields[batchimagejob.FieldActualCost] = struct{}{} +} + +// ActualCostCleared returns if the "actual_cost" field was cleared in this mutation. +func (m *BatchImageJobMutation) ActualCostCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldActualCost] + return ok +} + +// ResetActualCost resets all changes to the "actual_cost" field. +func (m *BatchImageJobMutation) ResetActualCost() { + m.actual_cost = nil + m.addactual_cost = nil + delete(m.clearedFields, batchimagejob.FieldActualCost) +} + +// SetCurrency sets the "currency" field. +func (m *BatchImageJobMutation) SetCurrency(s string) { + m.currency = &s +} + +// Currency returns the value of the "currency" field in the mutation. +func (m *BatchImageJobMutation) Currency() (r string, exists bool) { + v := m.currency + if v == nil { + return + } + return *v, true +} + +// OldCurrency returns the old "currency" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldCurrency(ctx context.Context) (v string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldCurrency is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldCurrency requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldCurrency: %w", err) + } + return oldValue.Currency, nil +} + +// ResetCurrency resets all changes to the "currency" field. +func (m *BatchImageJobMutation) ResetCurrency() { + m.currency = nil +} + +// SetHoldID sets the "hold_id" field. +func (m *BatchImageJobMutation) SetHoldID(s string) { + m.hold_id = &s +} + +// HoldID returns the value of the "hold_id" field in the mutation. +func (m *BatchImageJobMutation) HoldID() (r string, exists bool) { + v := m.hold_id + if v == nil { + return + } + return *v, true +} + +// OldHoldID returns the old "hold_id" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldHoldID(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldHoldID is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldHoldID requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldHoldID: %w", err) + } + return oldValue.HoldID, nil +} + +// ClearHoldID clears the value of the "hold_id" field. +func (m *BatchImageJobMutation) ClearHoldID() { + m.hold_id = nil + m.clearedFields[batchimagejob.FieldHoldID] = struct{}{} +} + +// HoldIDCleared returns if the "hold_id" field was cleared in this mutation. +func (m *BatchImageJobMutation) HoldIDCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldHoldID] + return ok +} + +// ResetHoldID resets all changes to the "hold_id" field. +func (m *BatchImageJobMutation) ResetHoldID() { + m.hold_id = nil + delete(m.clearedFields, batchimagejob.FieldHoldID) +} + +// SetIdempotencyKey sets the "idempotency_key" field. +func (m *BatchImageJobMutation) SetIdempotencyKey(s string) { + m.idempotency_key = &s +} + +// IdempotencyKey returns the value of the "idempotency_key" field in the mutation. +func (m *BatchImageJobMutation) IdempotencyKey() (r string, exists bool) { + v := m.idempotency_key + if v == nil { + return + } + return *v, true +} + +// OldIdempotencyKey returns the old "idempotency_key" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldIdempotencyKey(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldIdempotencyKey is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldIdempotencyKey requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldIdempotencyKey: %w", err) + } + return oldValue.IdempotencyKey, nil +} + +// ClearIdempotencyKey clears the value of the "idempotency_key" field. +func (m *BatchImageJobMutation) ClearIdempotencyKey() { + m.idempotency_key = nil + m.clearedFields[batchimagejob.FieldIdempotencyKey] = struct{}{} +} + +// IdempotencyKeyCleared returns if the "idempotency_key" field was cleared in this mutation. +func (m *BatchImageJobMutation) IdempotencyKeyCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldIdempotencyKey] + return ok +} + +// ResetIdempotencyKey resets all changes to the "idempotency_key" field. +func (m *BatchImageJobMutation) ResetIdempotencyKey() { + m.idempotency_key = nil + delete(m.clearedFields, batchimagejob.FieldIdempotencyKey) +} + +// SetRequestHash sets the "request_hash" field. +func (m *BatchImageJobMutation) SetRequestHash(s string) { + m.request_hash = &s +} + +// RequestHash returns the value of the "request_hash" field in the mutation. +func (m *BatchImageJobMutation) RequestHash() (r string, exists bool) { + v := m.request_hash + if v == nil { + return + } + return *v, true +} + +// OldRequestHash returns the old "request_hash" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldRequestHash(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldRequestHash is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldRequestHash requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldRequestHash: %w", err) + } + return oldValue.RequestHash, nil +} + +// ClearRequestHash clears the value of the "request_hash" field. +func (m *BatchImageJobMutation) ClearRequestHash() { + m.request_hash = nil + m.clearedFields[batchimagejob.FieldRequestHash] = struct{}{} +} + +// RequestHashCleared returns if the "request_hash" field was cleared in this mutation. +func (m *BatchImageJobMutation) RequestHashCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldRequestHash] + return ok +} + +// ResetRequestHash resets all changes to the "request_hash" field. +func (m *BatchImageJobMutation) ResetRequestHash() { + m.request_hash = nil + delete(m.clearedFields, batchimagejob.FieldRequestHash) +} + +// SetManifestHash sets the "manifest_hash" field. +func (m *BatchImageJobMutation) SetManifestHash(s string) { + m.manifest_hash = &s +} + +// ManifestHash returns the value of the "manifest_hash" field in the mutation. +func (m *BatchImageJobMutation) ManifestHash() (r string, exists bool) { + v := m.manifest_hash + if v == nil { + return + } + return *v, true +} + +// OldManifestHash returns the old "manifest_hash" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldManifestHash(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldManifestHash is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldManifestHash requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldManifestHash: %w", err) + } + return oldValue.ManifestHash, nil +} + +// ClearManifestHash clears the value of the "manifest_hash" field. +func (m *BatchImageJobMutation) ClearManifestHash() { + m.manifest_hash = nil + m.clearedFields[batchimagejob.FieldManifestHash] = struct{}{} +} + +// ManifestHashCleared returns if the "manifest_hash" field was cleared in this mutation. +func (m *BatchImageJobMutation) ManifestHashCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldManifestHash] + return ok +} + +// ResetManifestHash resets all changes to the "manifest_hash" field. +func (m *BatchImageJobMutation) ResetManifestHash() { + m.manifest_hash = nil + delete(m.clearedFields, batchimagejob.FieldManifestHash) +} + +// SetRetryCount sets the "retry_count" field. +func (m *BatchImageJobMutation) SetRetryCount(i int) { + m.retry_count = &i + m.addretry_count = nil +} + +// RetryCount returns the value of the "retry_count" field in the mutation. +func (m *BatchImageJobMutation) RetryCount() (r int, exists bool) { + v := m.retry_count + if v == nil { + return + } + return *v, true +} + +// OldRetryCount returns the old "retry_count" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldRetryCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldRetryCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldRetryCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldRetryCount: %w", err) + } + return oldValue.RetryCount, nil +} + +// AddRetryCount adds i to the "retry_count" field. +func (m *BatchImageJobMutation) AddRetryCount(i int) { + if m.addretry_count != nil { + *m.addretry_count += i + } else { + m.addretry_count = &i + } +} + +// AddedRetryCount returns the value that was added to the "retry_count" field in this mutation. +func (m *BatchImageJobMutation) AddedRetryCount() (r int, exists bool) { + v := m.addretry_count + if v == nil { + return + } + return *v, true +} + +// ResetRetryCount resets all changes to the "retry_count" field. +func (m *BatchImageJobMutation) ResetRetryCount() { + m.retry_count = nil + m.addretry_count = nil +} + +// SetVersion sets the "version" field. +func (m *BatchImageJobMutation) SetVersion(i int) { + m.version = &i + m.addversion = nil +} + +// Version returns the value of the "version" field in the mutation. +func (m *BatchImageJobMutation) Version() (r int, exists bool) { + v := m.version + if v == nil { + return + } + return *v, true +} + +// OldVersion returns the old "version" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldVersion(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVersion is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVersion requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVersion: %w", err) + } + return oldValue.Version, nil +} + +// AddVersion adds i to the "version" field. +func (m *BatchImageJobMutation) AddVersion(i int) { + if m.addversion != nil { + *m.addversion += i + } else { + m.addversion = &i + } +} + +// AddedVersion returns the value that was added to the "version" field in this mutation. +func (m *BatchImageJobMutation) AddedVersion() (r int, exists bool) { + v := m.addversion + if v == nil { + return + } + return *v, true +} + +// ResetVersion resets all changes to the "version" field. +func (m *BatchImageJobMutation) ResetVersion() { + m.version = nil + m.addversion = nil +} + +// SetOutputExpiresAt sets the "output_expires_at" field. +func (m *BatchImageJobMutation) SetOutputExpiresAt(t time.Time) { + m.output_expires_at = &t +} + +// OutputExpiresAt returns the value of the "output_expires_at" field in the mutation. +func (m *BatchImageJobMutation) OutputExpiresAt() (r time.Time, exists bool) { + v := m.output_expires_at + if v == nil { + return + } + return *v, true +} + +// OldOutputExpiresAt returns the old "output_expires_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldOutputExpiresAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldOutputExpiresAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldOutputExpiresAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldOutputExpiresAt: %w", err) + } + return oldValue.OutputExpiresAt, nil +} + +// ClearOutputExpiresAt clears the value of the "output_expires_at" field. +func (m *BatchImageJobMutation) ClearOutputExpiresAt() { + m.output_expires_at = nil + m.clearedFields[batchimagejob.FieldOutputExpiresAt] = struct{}{} +} + +// OutputExpiresAtCleared returns if the "output_expires_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) OutputExpiresAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldOutputExpiresAt] + return ok +} + +// ResetOutputExpiresAt resets all changes to the "output_expires_at" field. +func (m *BatchImageJobMutation) ResetOutputExpiresAt() { + m.output_expires_at = nil + delete(m.clearedFields, batchimagejob.FieldOutputExpiresAt) +} + +// SetInputDeletedAt sets the "input_deleted_at" field. +func (m *BatchImageJobMutation) SetInputDeletedAt(t time.Time) { + m.input_deleted_at = &t +} + +// InputDeletedAt returns the value of the "input_deleted_at" field in the mutation. +func (m *BatchImageJobMutation) InputDeletedAt() (r time.Time, exists bool) { + v := m.input_deleted_at + if v == nil { + return + } + return *v, true +} + +// OldInputDeletedAt returns the old "input_deleted_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldInputDeletedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldInputDeletedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldInputDeletedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldInputDeletedAt: %w", err) + } + return oldValue.InputDeletedAt, nil +} + +// ClearInputDeletedAt clears the value of the "input_deleted_at" field. +func (m *BatchImageJobMutation) ClearInputDeletedAt() { + m.input_deleted_at = nil + m.clearedFields[batchimagejob.FieldInputDeletedAt] = struct{}{} +} + +// InputDeletedAtCleared returns if the "input_deleted_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) InputDeletedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldInputDeletedAt] + return ok +} + +// ResetInputDeletedAt resets all changes to the "input_deleted_at" field. +func (m *BatchImageJobMutation) ResetInputDeletedAt() { + m.input_deleted_at = nil + delete(m.clearedFields, batchimagejob.FieldInputDeletedAt) +} + +// SetOutputDeletedAt sets the "output_deleted_at" field. +func (m *BatchImageJobMutation) SetOutputDeletedAt(t time.Time) { + m.output_deleted_at = &t +} + +// OutputDeletedAt returns the value of the "output_deleted_at" field in the mutation. +func (m *BatchImageJobMutation) OutputDeletedAt() (r time.Time, exists bool) { + v := m.output_deleted_at + if v == nil { + return + } + return *v, true +} + +// OldOutputDeletedAt returns the old "output_deleted_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldOutputDeletedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldOutputDeletedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldOutputDeletedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldOutputDeletedAt: %w", err) + } + return oldValue.OutputDeletedAt, nil +} + +// ClearOutputDeletedAt clears the value of the "output_deleted_at" field. +func (m *BatchImageJobMutation) ClearOutputDeletedAt() { + m.output_deleted_at = nil + m.clearedFields[batchimagejob.FieldOutputDeletedAt] = struct{}{} +} + +// OutputDeletedAtCleared returns if the "output_deleted_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) OutputDeletedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldOutputDeletedAt] + return ok +} + +// ResetOutputDeletedAt resets all changes to the "output_deleted_at" field. +func (m *BatchImageJobMutation) ResetOutputDeletedAt() { + m.output_deleted_at = nil + delete(m.clearedFields, batchimagejob.FieldOutputDeletedAt) +} + +// SetDownloadedAt sets the "downloaded_at" field. +func (m *BatchImageJobMutation) SetDownloadedAt(t time.Time) { + m.downloaded_at = &t +} + +// DownloadedAt returns the value of the "downloaded_at" field in the mutation. +func (m *BatchImageJobMutation) DownloadedAt() (r time.Time, exists bool) { + v := m.downloaded_at + if v == nil { + return + } + return *v, true +} + +// OldDownloadedAt returns the old "downloaded_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldDownloadedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldDownloadedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldDownloadedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldDownloadedAt: %w", err) + } + return oldValue.DownloadedAt, nil +} + +// ClearDownloadedAt clears the value of the "downloaded_at" field. +func (m *BatchImageJobMutation) ClearDownloadedAt() { + m.downloaded_at = nil + m.clearedFields[batchimagejob.FieldDownloadedAt] = struct{}{} +} + +// DownloadedAtCleared returns if the "downloaded_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) DownloadedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldDownloadedAt] + return ok +} + +// ResetDownloadedAt resets all changes to the "downloaded_at" field. +func (m *BatchImageJobMutation) ResetDownloadedAt() { + m.downloaded_at = nil + delete(m.clearedFields, batchimagejob.FieldDownloadedAt) +} + +// SetUserDeletedAt sets the "user_deleted_at" field. +func (m *BatchImageJobMutation) SetUserDeletedAt(t time.Time) { + m.user_deleted_at = &t +} + +// UserDeletedAt returns the value of the "user_deleted_at" field in the mutation. +func (m *BatchImageJobMutation) UserDeletedAt() (r time.Time, exists bool) { + v := m.user_deleted_at + if v == nil { + return + } + return *v, true +} + +// OldUserDeletedAt returns the old "user_deleted_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldUserDeletedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldUserDeletedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldUserDeletedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldUserDeletedAt: %w", err) + } + return oldValue.UserDeletedAt, nil +} + +// ClearUserDeletedAt clears the value of the "user_deleted_at" field. +func (m *BatchImageJobMutation) ClearUserDeletedAt() { + m.user_deleted_at = nil + m.clearedFields[batchimagejob.FieldUserDeletedAt] = struct{}{} +} + +// UserDeletedAtCleared returns if the "user_deleted_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) UserDeletedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldUserDeletedAt] + return ok +} + +// ResetUserDeletedAt resets all changes to the "user_deleted_at" field. +func (m *BatchImageJobMutation) ResetUserDeletedAt() { + m.user_deleted_at = nil + delete(m.clearedFields, batchimagejob.FieldUserDeletedAt) +} + +// SetLastErrorCode sets the "last_error_code" field. +func (m *BatchImageJobMutation) SetLastErrorCode(s string) { + m.last_error_code = &s +} + +// LastErrorCode returns the value of the "last_error_code" field in the mutation. +func (m *BatchImageJobMutation) LastErrorCode() (r string, exists bool) { + v := m.last_error_code + if v == nil { + return + } + return *v, true +} + +// OldLastErrorCode returns the old "last_error_code" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldLastErrorCode(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldLastErrorCode is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldLastErrorCode requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldLastErrorCode: %w", err) + } + return oldValue.LastErrorCode, nil +} + +// ClearLastErrorCode clears the value of the "last_error_code" field. +func (m *BatchImageJobMutation) ClearLastErrorCode() { + m.last_error_code = nil + m.clearedFields[batchimagejob.FieldLastErrorCode] = struct{}{} +} + +// LastErrorCodeCleared returns if the "last_error_code" field was cleared in this mutation. +func (m *BatchImageJobMutation) LastErrorCodeCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldLastErrorCode] + return ok +} + +// ResetLastErrorCode resets all changes to the "last_error_code" field. +func (m *BatchImageJobMutation) ResetLastErrorCode() { + m.last_error_code = nil + delete(m.clearedFields, batchimagejob.FieldLastErrorCode) +} + +// SetLastErrorMessage sets the "last_error_message" field. +func (m *BatchImageJobMutation) SetLastErrorMessage(s string) { + m.last_error_message = &s +} + +// LastErrorMessage returns the value of the "last_error_message" field in the mutation. +func (m *BatchImageJobMutation) LastErrorMessage() (r string, exists bool) { + v := m.last_error_message + if v == nil { + return + } + return *v, true +} + +// OldLastErrorMessage returns the old "last_error_message" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldLastErrorMessage(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldLastErrorMessage is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldLastErrorMessage requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldLastErrorMessage: %w", err) + } + return oldValue.LastErrorMessage, nil +} + +// ClearLastErrorMessage clears the value of the "last_error_message" field. +func (m *BatchImageJobMutation) ClearLastErrorMessage() { + m.last_error_message = nil + m.clearedFields[batchimagejob.FieldLastErrorMessage] = struct{}{} +} + +// LastErrorMessageCleared returns if the "last_error_message" field was cleared in this mutation. +func (m *BatchImageJobMutation) LastErrorMessageCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldLastErrorMessage] + return ok +} + +// ResetLastErrorMessage resets all changes to the "last_error_message" field. +func (m *BatchImageJobMutation) ResetLastErrorMessage() { + m.last_error_message = nil + delete(m.clearedFields, batchimagejob.FieldLastErrorMessage) +} + +// SetCreatedAt sets the "created_at" field. +func (m *BatchImageJobMutation) SetCreatedAt(t time.Time) { + m.created_at = &t +} + +// CreatedAt returns the value of the "created_at" field in the mutation. +func (m *BatchImageJobMutation) CreatedAt() (r time.Time, exists bool) { + v := m.created_at + if v == nil { + return + } + return *v, true +} + +// OldCreatedAt returns the old "created_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldCreatedAt(ctx context.Context) (v time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldCreatedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldCreatedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldCreatedAt: %w", err) + } + return oldValue.CreatedAt, nil +} + +// ResetCreatedAt resets all changes to the "created_at" field. +func (m *BatchImageJobMutation) ResetCreatedAt() { + m.created_at = nil +} + +// SetUpdatedAt sets the "updated_at" field. +func (m *BatchImageJobMutation) SetUpdatedAt(t time.Time) { + m.updated_at = &t +} + +// UpdatedAt returns the value of the "updated_at" field in the mutation. +func (m *BatchImageJobMutation) UpdatedAt() (r time.Time, exists bool) { + v := m.updated_at + if v == nil { + return + } + return *v, true +} + +// OldUpdatedAt returns the old "updated_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldUpdatedAt(ctx context.Context) (v time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldUpdatedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldUpdatedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldUpdatedAt: %w", err) + } + return oldValue.UpdatedAt, nil +} + +// ResetUpdatedAt resets all changes to the "updated_at" field. +func (m *BatchImageJobMutation) ResetUpdatedAt() { + m.updated_at = nil +} + +// SetSubmittedAt sets the "submitted_at" field. +func (m *BatchImageJobMutation) SetSubmittedAt(t time.Time) { + m.submitted_at = &t +} + +// SubmittedAt returns the value of the "submitted_at" field in the mutation. +func (m *BatchImageJobMutation) SubmittedAt() (r time.Time, exists bool) { + v := m.submitted_at + if v == nil { + return + } + return *v, true +} + +// OldSubmittedAt returns the old "submitted_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldSubmittedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldSubmittedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldSubmittedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldSubmittedAt: %w", err) + } + return oldValue.SubmittedAt, nil +} + +// ClearSubmittedAt clears the value of the "submitted_at" field. +func (m *BatchImageJobMutation) ClearSubmittedAt() { + m.submitted_at = nil + m.clearedFields[batchimagejob.FieldSubmittedAt] = struct{}{} +} + +// SubmittedAtCleared returns if the "submitted_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) SubmittedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldSubmittedAt] + return ok +} + +// ResetSubmittedAt resets all changes to the "submitted_at" field. +func (m *BatchImageJobMutation) ResetSubmittedAt() { + m.submitted_at = nil + delete(m.clearedFields, batchimagejob.FieldSubmittedAt) +} + +// SetStartedAt sets the "started_at" field. +func (m *BatchImageJobMutation) SetStartedAt(t time.Time) { + m.started_at = &t +} + +// StartedAt returns the value of the "started_at" field in the mutation. +func (m *BatchImageJobMutation) StartedAt() (r time.Time, exists bool) { + v := m.started_at + if v == nil { + return + } + return *v, true +} + +// OldStartedAt returns the old "started_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldStartedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldStartedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldStartedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldStartedAt: %w", err) + } + return oldValue.StartedAt, nil +} + +// ClearStartedAt clears the value of the "started_at" field. +func (m *BatchImageJobMutation) ClearStartedAt() { + m.started_at = nil + m.clearedFields[batchimagejob.FieldStartedAt] = struct{}{} +} + +// StartedAtCleared returns if the "started_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) StartedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldStartedAt] + return ok +} + +// ResetStartedAt resets all changes to the "started_at" field. +func (m *BatchImageJobMutation) ResetStartedAt() { + m.started_at = nil + delete(m.clearedFields, batchimagejob.FieldStartedAt) +} + +// SetFinishedAt sets the "finished_at" field. +func (m *BatchImageJobMutation) SetFinishedAt(t time.Time) { + m.finished_at = &t +} + +// FinishedAt returns the value of the "finished_at" field in the mutation. +func (m *BatchImageJobMutation) FinishedAt() (r time.Time, exists bool) { + v := m.finished_at + if v == nil { + return + } + return *v, true +} + +// OldFinishedAt returns the old "finished_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldFinishedAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldFinishedAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldFinishedAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldFinishedAt: %w", err) + } + return oldValue.FinishedAt, nil +} + +// ClearFinishedAt clears the value of the "finished_at" field. +func (m *BatchImageJobMutation) ClearFinishedAt() { + m.finished_at = nil + m.clearedFields[batchimagejob.FieldFinishedAt] = struct{}{} +} + +// FinishedAtCleared returns if the "finished_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) FinishedAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldFinishedAt] + return ok +} + +// ResetFinishedAt resets all changes to the "finished_at" field. +func (m *BatchImageJobMutation) ResetFinishedAt() { + m.finished_at = nil + delete(m.clearedFields, batchimagejob.FieldFinishedAt) +} + +// SetSettledAt sets the "settled_at" field. +func (m *BatchImageJobMutation) SetSettledAt(t time.Time) { + m.settled_at = &t +} + +// SettledAt returns the value of the "settled_at" field in the mutation. +func (m *BatchImageJobMutation) SettledAt() (r time.Time, exists bool) { + v := m.settled_at + if v == nil { + return + } + return *v, true +} + +// OldSettledAt returns the old "settled_at" field's value of the BatchImageJob entity. +// If the BatchImageJob object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *BatchImageJobMutation) OldSettledAt(ctx context.Context) (v *time.Time, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldSettledAt is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldSettledAt requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldSettledAt: %w", err) + } + return oldValue.SettledAt, nil +} + +// ClearSettledAt clears the value of the "settled_at" field. +func (m *BatchImageJobMutation) ClearSettledAt() { + m.settled_at = nil + m.clearedFields[batchimagejob.FieldSettledAt] = struct{}{} +} + +// SettledAtCleared returns if the "settled_at" field was cleared in this mutation. +func (m *BatchImageJobMutation) SettledAtCleared() bool { + _, ok := m.clearedFields[batchimagejob.FieldSettledAt] + return ok +} + +// ResetSettledAt resets all changes to the "settled_at" field. +func (m *BatchImageJobMutation) ResetSettledAt() { + m.settled_at = nil + delete(m.clearedFields, batchimagejob.FieldSettledAt) +} + +// Where appends a list predicates to the BatchImageJobMutation builder. +func (m *BatchImageJobMutation) Where(ps ...predicate.BatchImageJob) { + m.predicates = append(m.predicates, ps...) +} + +// WhereP appends storage-level predicates to the BatchImageJobMutation builder. Using this method, +// users can use type-assertion to append predicates that do not depend on any generated package. +func (m *BatchImageJobMutation) WhereP(ps ...func(*sql.Selector)) { + p := make([]predicate.BatchImageJob, len(ps)) + for i := range ps { + p[i] = ps[i] + } + m.Where(p...) +} + +// Op returns the operation name. +func (m *BatchImageJobMutation) Op() Op { + return m.op +} + +// SetOp allows setting the mutation operation. +func (m *BatchImageJobMutation) SetOp(op Op) { + m.op = op +} + +// Type returns the node type of this mutation (BatchImageJob). +func (m *BatchImageJobMutation) Type() string { + return m.typ +} + +// Fields returns all fields that were changed during this mutation. Note that in +// order to get all numeric fields that were incremented/decremented, call +// AddedFields(). +func (m *BatchImageJobMutation) Fields() []string { + fields := make([]string, 0, 40) + if m.batch_id != nil { + fields = append(fields, batchimagejob.FieldBatchID) + } + if m.user_id != nil { + fields = append(fields, batchimagejob.FieldUserID) + } + if m.api_key_id != nil { + fields = append(fields, batchimagejob.FieldAPIKeyID) + } + if m.account_id != nil { + fields = append(fields, batchimagejob.FieldAccountID) + } + if m.provider != nil { + fields = append(fields, batchimagejob.FieldProvider) + } + if m.model != nil { + fields = append(fields, batchimagejob.FieldModel) + } + if m.task_name != nil { + fields = append(fields, batchimagejob.FieldTaskName) + } + if m.status != nil { + fields = append(fields, batchimagejob.FieldStatus) + } + if m.provider_job_name != nil { + fields = append(fields, batchimagejob.FieldProviderJobName) + } + if m.provider_input_ref != nil { + fields = append(fields, batchimagejob.FieldProviderInputRef) + } + if m.provider_output_ref != nil { + fields = append(fields, batchimagejob.FieldProviderOutputRef) + } + if m.gcs_input_uri != nil { + fields = append(fields, batchimagejob.FieldGcsInputURI) + } + if m.gcs_output_uri != nil { + fields = append(fields, batchimagejob.FieldGcsOutputURI) + } + if m.item_count != nil { + fields = append(fields, batchimagejob.FieldItemCount) + } + if m.success_count != nil { + fields = append(fields, batchimagejob.FieldSuccessCount) + } + if m.fail_count != nil { + fields = append(fields, batchimagejob.FieldFailCount) + } + if m.cancelled_count != nil { + fields = append(fields, batchimagejob.FieldCancelledCount) + } + if m.estimated_cost != nil { + fields = append(fields, batchimagejob.FieldEstimatedCost) + } + if m.hold_amount != nil { + fields = append(fields, batchimagejob.FieldHoldAmount) + } + if m.actual_cost != nil { + fields = append(fields, batchimagejob.FieldActualCost) + } + if m.currency != nil { + fields = append(fields, batchimagejob.FieldCurrency) + } + if m.hold_id != nil { + fields = append(fields, batchimagejob.FieldHoldID) + } + if m.idempotency_key != nil { + fields = append(fields, batchimagejob.FieldIdempotencyKey) + } + if m.request_hash != nil { + fields = append(fields, batchimagejob.FieldRequestHash) + } + if m.manifest_hash != nil { + fields = append(fields, batchimagejob.FieldManifestHash) + } + if m.retry_count != nil { + fields = append(fields, batchimagejob.FieldRetryCount) + } + if m.version != nil { + fields = append(fields, batchimagejob.FieldVersion) + } + if m.output_expires_at != nil { + fields = append(fields, batchimagejob.FieldOutputExpiresAt) + } + if m.input_deleted_at != nil { + fields = append(fields, batchimagejob.FieldInputDeletedAt) + } + if m.output_deleted_at != nil { + fields = append(fields, batchimagejob.FieldOutputDeletedAt) + } + if m.downloaded_at != nil { + fields = append(fields, batchimagejob.FieldDownloadedAt) + } + if m.user_deleted_at != nil { + fields = append(fields, batchimagejob.FieldUserDeletedAt) + } + if m.last_error_code != nil { + fields = append(fields, batchimagejob.FieldLastErrorCode) + } + if m.last_error_message != nil { + fields = append(fields, batchimagejob.FieldLastErrorMessage) + } + if m.created_at != nil { + fields = append(fields, batchimagejob.FieldCreatedAt) + } + if m.updated_at != nil { + fields = append(fields, batchimagejob.FieldUpdatedAt) + } + if m.submitted_at != nil { + fields = append(fields, batchimagejob.FieldSubmittedAt) + } + if m.started_at != nil { + fields = append(fields, batchimagejob.FieldStartedAt) + } + if m.finished_at != nil { + fields = append(fields, batchimagejob.FieldFinishedAt) + } + if m.settled_at != nil { + fields = append(fields, batchimagejob.FieldSettledAt) + } + return fields +} + +// Field returns the value of a field with the given name. The second boolean +// return value indicates that this field was not set, or was not defined in the +// schema. +func (m *BatchImageJobMutation) Field(name string) (ent.Value, bool) { + switch name { + case batchimagejob.FieldBatchID: + return m.BatchID() + case batchimagejob.FieldUserID: + return m.UserID() + case batchimagejob.FieldAPIKeyID: + return m.APIKeyID() + case batchimagejob.FieldAccountID: + return m.AccountID() + case batchimagejob.FieldProvider: + return m.Provider() + case batchimagejob.FieldModel: + return m.Model() + case batchimagejob.FieldTaskName: + return m.TaskName() + case batchimagejob.FieldStatus: + return m.Status() + case batchimagejob.FieldProviderJobName: + return m.ProviderJobName() + case batchimagejob.FieldProviderInputRef: + return m.ProviderInputRef() + case batchimagejob.FieldProviderOutputRef: + return m.ProviderOutputRef() + case batchimagejob.FieldGcsInputURI: + return m.GcsInputURI() + case batchimagejob.FieldGcsOutputURI: + return m.GcsOutputURI() + case batchimagejob.FieldItemCount: + return m.ItemCount() + case batchimagejob.FieldSuccessCount: + return m.SuccessCount() + case batchimagejob.FieldFailCount: + return m.FailCount() + case batchimagejob.FieldCancelledCount: + return m.CancelledCount() + case batchimagejob.FieldEstimatedCost: + return m.EstimatedCost() + case batchimagejob.FieldHoldAmount: + return m.HoldAmount() + case batchimagejob.FieldActualCost: + return m.ActualCost() + case batchimagejob.FieldCurrency: + return m.Currency() + case batchimagejob.FieldHoldID: + return m.HoldID() + case batchimagejob.FieldIdempotencyKey: + return m.IdempotencyKey() + case batchimagejob.FieldRequestHash: + return m.RequestHash() + case batchimagejob.FieldManifestHash: + return m.ManifestHash() + case batchimagejob.FieldRetryCount: + return m.RetryCount() + case batchimagejob.FieldVersion: + return m.Version() + case batchimagejob.FieldOutputExpiresAt: + return m.OutputExpiresAt() + case batchimagejob.FieldInputDeletedAt: + return m.InputDeletedAt() + case batchimagejob.FieldOutputDeletedAt: + return m.OutputDeletedAt() + case batchimagejob.FieldDownloadedAt: + return m.DownloadedAt() + case batchimagejob.FieldUserDeletedAt: + return m.UserDeletedAt() + case batchimagejob.FieldLastErrorCode: + return m.LastErrorCode() + case batchimagejob.FieldLastErrorMessage: + return m.LastErrorMessage() + case batchimagejob.FieldCreatedAt: + return m.CreatedAt() + case batchimagejob.FieldUpdatedAt: + return m.UpdatedAt() + case batchimagejob.FieldSubmittedAt: + return m.SubmittedAt() + case batchimagejob.FieldStartedAt: + return m.StartedAt() + case batchimagejob.FieldFinishedAt: + return m.FinishedAt() + case batchimagejob.FieldSettledAt: + return m.SettledAt() + } + return nil, false +} + +// OldField returns the old value of the field from the database. An error is +// returned if the mutation operation is not UpdateOne, or the query to the +// database failed. +func (m *BatchImageJobMutation) OldField(ctx context.Context, name string) (ent.Value, error) { + switch name { + case batchimagejob.FieldBatchID: + return m.OldBatchID(ctx) + case batchimagejob.FieldUserID: + return m.OldUserID(ctx) + case batchimagejob.FieldAPIKeyID: + return m.OldAPIKeyID(ctx) + case batchimagejob.FieldAccountID: + return m.OldAccountID(ctx) + case batchimagejob.FieldProvider: + return m.OldProvider(ctx) + case batchimagejob.FieldModel: + return m.OldModel(ctx) + case batchimagejob.FieldTaskName: + return m.OldTaskName(ctx) + case batchimagejob.FieldStatus: + return m.OldStatus(ctx) + case batchimagejob.FieldProviderJobName: + return m.OldProviderJobName(ctx) + case batchimagejob.FieldProviderInputRef: + return m.OldProviderInputRef(ctx) + case batchimagejob.FieldProviderOutputRef: + return m.OldProviderOutputRef(ctx) + case batchimagejob.FieldGcsInputURI: + return m.OldGcsInputURI(ctx) + case batchimagejob.FieldGcsOutputURI: + return m.OldGcsOutputURI(ctx) + case batchimagejob.FieldItemCount: + return m.OldItemCount(ctx) + case batchimagejob.FieldSuccessCount: + return m.OldSuccessCount(ctx) + case batchimagejob.FieldFailCount: + return m.OldFailCount(ctx) + case batchimagejob.FieldCancelledCount: + return m.OldCancelledCount(ctx) + case batchimagejob.FieldEstimatedCost: + return m.OldEstimatedCost(ctx) + case batchimagejob.FieldHoldAmount: + return m.OldHoldAmount(ctx) + case batchimagejob.FieldActualCost: + return m.OldActualCost(ctx) + case batchimagejob.FieldCurrency: + return m.OldCurrency(ctx) + case batchimagejob.FieldHoldID: + return m.OldHoldID(ctx) + case batchimagejob.FieldIdempotencyKey: + return m.OldIdempotencyKey(ctx) + case batchimagejob.FieldRequestHash: + return m.OldRequestHash(ctx) + case batchimagejob.FieldManifestHash: + return m.OldManifestHash(ctx) + case batchimagejob.FieldRetryCount: + return m.OldRetryCount(ctx) + case batchimagejob.FieldVersion: + return m.OldVersion(ctx) + case batchimagejob.FieldOutputExpiresAt: + return m.OldOutputExpiresAt(ctx) + case batchimagejob.FieldInputDeletedAt: + return m.OldInputDeletedAt(ctx) + case batchimagejob.FieldOutputDeletedAt: + return m.OldOutputDeletedAt(ctx) + case batchimagejob.FieldDownloadedAt: + return m.OldDownloadedAt(ctx) + case batchimagejob.FieldUserDeletedAt: + return m.OldUserDeletedAt(ctx) + case batchimagejob.FieldLastErrorCode: + return m.OldLastErrorCode(ctx) + case batchimagejob.FieldLastErrorMessage: + return m.OldLastErrorMessage(ctx) + case batchimagejob.FieldCreatedAt: + return m.OldCreatedAt(ctx) + case batchimagejob.FieldUpdatedAt: + return m.OldUpdatedAt(ctx) + case batchimagejob.FieldSubmittedAt: + return m.OldSubmittedAt(ctx) + case batchimagejob.FieldStartedAt: + return m.OldStartedAt(ctx) + case batchimagejob.FieldFinishedAt: + return m.OldFinishedAt(ctx) + case batchimagejob.FieldSettledAt: + return m.OldSettledAt(ctx) + } + return nil, fmt.Errorf("unknown BatchImageJob field %s", name) +} + +// SetField sets the value of a field with the given name. It returns an error if +// the field is not defined in the schema, or if the type mismatched the field +// type. +func (m *BatchImageJobMutation) SetField(name string, value ent.Value) error { + switch name { + case batchimagejob.FieldBatchID: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetBatchID(v) + return nil + case batchimagejob.FieldUserID: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetUserID(v) + return nil + case batchimagejob.FieldAPIKeyID: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetAPIKeyID(v) + return nil + case batchimagejob.FieldAccountID: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetAccountID(v) + return nil + case batchimagejob.FieldProvider: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetProvider(v) + return nil + case batchimagejob.FieldModel: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetModel(v) + return nil + case batchimagejob.FieldTaskName: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetTaskName(v) + return nil + case batchimagejob.FieldStatus: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetStatus(v) + return nil + case batchimagejob.FieldProviderJobName: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetProviderJobName(v) + return nil + case batchimagejob.FieldProviderInputRef: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetProviderInputRef(v) + return nil + case batchimagejob.FieldProviderOutputRef: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetProviderOutputRef(v) + return nil + case batchimagejob.FieldGcsInputURI: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetGcsInputURI(v) + return nil + case batchimagejob.FieldGcsOutputURI: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetGcsOutputURI(v) + return nil + case batchimagejob.FieldItemCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetItemCount(v) + return nil + case batchimagejob.FieldSuccessCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetSuccessCount(v) + return nil + case batchimagejob.FieldFailCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetFailCount(v) + return nil + case batchimagejob.FieldCancelledCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetCancelledCount(v) + return nil + case batchimagejob.FieldEstimatedCost: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetEstimatedCost(v) + return nil + case batchimagejob.FieldHoldAmount: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetHoldAmount(v) + return nil + case batchimagejob.FieldActualCost: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetActualCost(v) + return nil + case batchimagejob.FieldCurrency: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetCurrency(v) + return nil + case batchimagejob.FieldHoldID: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetHoldID(v) + return nil + case batchimagejob.FieldIdempotencyKey: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetIdempotencyKey(v) + return nil + case batchimagejob.FieldRequestHash: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetRequestHash(v) + return nil + case batchimagejob.FieldManifestHash: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetManifestHash(v) + return nil + case batchimagejob.FieldRetryCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetRetryCount(v) + return nil + case batchimagejob.FieldVersion: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVersion(v) + return nil + case batchimagejob.FieldOutputExpiresAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetOutputExpiresAt(v) + return nil + case batchimagejob.FieldInputDeletedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetInputDeletedAt(v) + return nil + case batchimagejob.FieldOutputDeletedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetOutputDeletedAt(v) + return nil + case batchimagejob.FieldDownloadedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetDownloadedAt(v) + return nil + case batchimagejob.FieldUserDeletedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetUserDeletedAt(v) + return nil + case batchimagejob.FieldLastErrorCode: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetLastErrorCode(v) + return nil + case batchimagejob.FieldLastErrorMessage: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetLastErrorMessage(v) + return nil + case batchimagejob.FieldCreatedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetCreatedAt(v) + return nil + case batchimagejob.FieldUpdatedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetUpdatedAt(v) + return nil + case batchimagejob.FieldSubmittedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetSubmittedAt(v) + return nil + case batchimagejob.FieldStartedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetStartedAt(v) + return nil + case batchimagejob.FieldFinishedAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetFinishedAt(v) + return nil + case batchimagejob.FieldSettledAt: + v, ok := value.(time.Time) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetSettledAt(v) + return nil + } + return fmt.Errorf("unknown BatchImageJob field %s", name) +} + +// AddedFields returns all numeric fields that were incremented/decremented during +// this mutation. +func (m *BatchImageJobMutation) AddedFields() []string { + var fields []string + if m.adduser_id != nil { + fields = append(fields, batchimagejob.FieldUserID) + } + if m.addapi_key_id != nil { + fields = append(fields, batchimagejob.FieldAPIKeyID) + } + if m.addaccount_id != nil { + fields = append(fields, batchimagejob.FieldAccountID) + } + if m.additem_count != nil { + fields = append(fields, batchimagejob.FieldItemCount) + } + if m.addsuccess_count != nil { + fields = append(fields, batchimagejob.FieldSuccessCount) + } + if m.addfail_count != nil { + fields = append(fields, batchimagejob.FieldFailCount) + } + if m.addcancelled_count != nil { + fields = append(fields, batchimagejob.FieldCancelledCount) + } + if m.addestimated_cost != nil { + fields = append(fields, batchimagejob.FieldEstimatedCost) + } + if m.addhold_amount != nil { + fields = append(fields, batchimagejob.FieldHoldAmount) + } + if m.addactual_cost != nil { + fields = append(fields, batchimagejob.FieldActualCost) + } + if m.addretry_count != nil { + fields = append(fields, batchimagejob.FieldRetryCount) + } + if m.addversion != nil { + fields = append(fields, batchimagejob.FieldVersion) + } + return fields +} + +// AddedField returns the numeric value that was incremented/decremented on a field +// with the given name. The second boolean return value indicates that this field +// was not set, or was not defined in the schema. +func (m *BatchImageJobMutation) AddedField(name string) (ent.Value, bool) { + switch name { + case batchimagejob.FieldUserID: + return m.AddedUserID() + case batchimagejob.FieldAPIKeyID: + return m.AddedAPIKeyID() + case batchimagejob.FieldAccountID: + return m.AddedAccountID() + case batchimagejob.FieldItemCount: + return m.AddedItemCount() + case batchimagejob.FieldSuccessCount: + return m.AddedSuccessCount() + case batchimagejob.FieldFailCount: + return m.AddedFailCount() + case batchimagejob.FieldCancelledCount: + return m.AddedCancelledCount() + case batchimagejob.FieldEstimatedCost: + return m.AddedEstimatedCost() + case batchimagejob.FieldHoldAmount: + return m.AddedHoldAmount() + case batchimagejob.FieldActualCost: + return m.AddedActualCost() + case batchimagejob.FieldRetryCount: + return m.AddedRetryCount() + case batchimagejob.FieldVersion: + return m.AddedVersion() + } + return nil, false +} + +// AddField adds the value to the field with the given name. It returns an error if +// the field is not defined in the schema, or if the type mismatched the field +// type. +func (m *BatchImageJobMutation) AddField(name string, value ent.Value) error { + switch name { + case batchimagejob.FieldUserID: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddUserID(v) + return nil + case batchimagejob.FieldAPIKeyID: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddAPIKeyID(v) + return nil + case batchimagejob.FieldAccountID: + v, ok := value.(int64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddAccountID(v) + return nil + case batchimagejob.FieldItemCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddItemCount(v) + return nil + case batchimagejob.FieldSuccessCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddSuccessCount(v) + return nil + case batchimagejob.FieldFailCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddFailCount(v) + return nil + case batchimagejob.FieldCancelledCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddCancelledCount(v) + return nil + case batchimagejob.FieldEstimatedCost: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddEstimatedCost(v) + return nil + case batchimagejob.FieldHoldAmount: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddHoldAmount(v) + return nil + case batchimagejob.FieldActualCost: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddActualCost(v) + return nil + case batchimagejob.FieldRetryCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddRetryCount(v) + return nil + case batchimagejob.FieldVersion: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVersion(v) + return nil + } + return fmt.Errorf("unknown BatchImageJob numeric field %s", name) +} + +// ClearedFields returns all nullable fields that were cleared during this +// mutation. +func (m *BatchImageJobMutation) ClearedFields() []string { + var fields []string + if m.FieldCleared(batchimagejob.FieldAPIKeyID) { + fields = append(fields, batchimagejob.FieldAPIKeyID) + } + if m.FieldCleared(batchimagejob.FieldAccountID) { + fields = append(fields, batchimagejob.FieldAccountID) + } + if m.FieldCleared(batchimagejob.FieldProviderJobName) { + fields = append(fields, batchimagejob.FieldProviderJobName) + } + if m.FieldCleared(batchimagejob.FieldProviderInputRef) { + fields = append(fields, batchimagejob.FieldProviderInputRef) + } + if m.FieldCleared(batchimagejob.FieldProviderOutputRef) { + fields = append(fields, batchimagejob.FieldProviderOutputRef) + } + if m.FieldCleared(batchimagejob.FieldGcsInputURI) { + fields = append(fields, batchimagejob.FieldGcsInputURI) + } + if m.FieldCleared(batchimagejob.FieldGcsOutputURI) { + fields = append(fields, batchimagejob.FieldGcsOutputURI) + } + if m.FieldCleared(batchimagejob.FieldHoldAmount) { + fields = append(fields, batchimagejob.FieldHoldAmount) + } + if m.FieldCleared(batchimagejob.FieldActualCost) { + fields = append(fields, batchimagejob.FieldActualCost) + } + if m.FieldCleared(batchimagejob.FieldHoldID) { + fields = append(fields, batchimagejob.FieldHoldID) + } + if m.FieldCleared(batchimagejob.FieldIdempotencyKey) { + fields = append(fields, batchimagejob.FieldIdempotencyKey) + } + if m.FieldCleared(batchimagejob.FieldRequestHash) { + fields = append(fields, batchimagejob.FieldRequestHash) + } + if m.FieldCleared(batchimagejob.FieldManifestHash) { + fields = append(fields, batchimagejob.FieldManifestHash) + } + if m.FieldCleared(batchimagejob.FieldOutputExpiresAt) { + fields = append(fields, batchimagejob.FieldOutputExpiresAt) + } + if m.FieldCleared(batchimagejob.FieldInputDeletedAt) { + fields = append(fields, batchimagejob.FieldInputDeletedAt) + } + if m.FieldCleared(batchimagejob.FieldOutputDeletedAt) { + fields = append(fields, batchimagejob.FieldOutputDeletedAt) + } + if m.FieldCleared(batchimagejob.FieldDownloadedAt) { + fields = append(fields, batchimagejob.FieldDownloadedAt) + } + if m.FieldCleared(batchimagejob.FieldUserDeletedAt) { + fields = append(fields, batchimagejob.FieldUserDeletedAt) + } + if m.FieldCleared(batchimagejob.FieldLastErrorCode) { + fields = append(fields, batchimagejob.FieldLastErrorCode) + } + if m.FieldCleared(batchimagejob.FieldLastErrorMessage) { + fields = append(fields, batchimagejob.FieldLastErrorMessage) + } + if m.FieldCleared(batchimagejob.FieldSubmittedAt) { + fields = append(fields, batchimagejob.FieldSubmittedAt) + } + if m.FieldCleared(batchimagejob.FieldStartedAt) { + fields = append(fields, batchimagejob.FieldStartedAt) + } + if m.FieldCleared(batchimagejob.FieldFinishedAt) { + fields = append(fields, batchimagejob.FieldFinishedAt) + } + if m.FieldCleared(batchimagejob.FieldSettledAt) { + fields = append(fields, batchimagejob.FieldSettledAt) + } + return fields +} + +// FieldCleared returns a boolean indicating if a field with the given name was +// cleared in this mutation. +func (m *BatchImageJobMutation) FieldCleared(name string) bool { + _, ok := m.clearedFields[name] + return ok +} + +// ClearField clears the value of the field with the given name. It returns an +// error if the field is not defined in the schema. +func (m *BatchImageJobMutation) ClearField(name string) error { + switch name { + case batchimagejob.FieldAPIKeyID: + m.ClearAPIKeyID() + return nil + case batchimagejob.FieldAccountID: + m.ClearAccountID() + return nil + case batchimagejob.FieldProviderJobName: + m.ClearProviderJobName() + return nil + case batchimagejob.FieldProviderInputRef: + m.ClearProviderInputRef() + return nil + case batchimagejob.FieldProviderOutputRef: + m.ClearProviderOutputRef() + return nil + case batchimagejob.FieldGcsInputURI: + m.ClearGcsInputURI() + return nil + case batchimagejob.FieldGcsOutputURI: + m.ClearGcsOutputURI() + return nil + case batchimagejob.FieldHoldAmount: + m.ClearHoldAmount() + return nil + case batchimagejob.FieldActualCost: + m.ClearActualCost() + return nil + case batchimagejob.FieldHoldID: + m.ClearHoldID() + return nil + case batchimagejob.FieldIdempotencyKey: + m.ClearIdempotencyKey() + return nil + case batchimagejob.FieldRequestHash: + m.ClearRequestHash() + return nil + case batchimagejob.FieldManifestHash: + m.ClearManifestHash() + return nil + case batchimagejob.FieldOutputExpiresAt: + m.ClearOutputExpiresAt() + return nil + case batchimagejob.FieldInputDeletedAt: + m.ClearInputDeletedAt() + return nil + case batchimagejob.FieldOutputDeletedAt: + m.ClearOutputDeletedAt() + return nil + case batchimagejob.FieldDownloadedAt: + m.ClearDownloadedAt() + return nil + case batchimagejob.FieldUserDeletedAt: + m.ClearUserDeletedAt() + return nil + case batchimagejob.FieldLastErrorCode: + m.ClearLastErrorCode() + return nil + case batchimagejob.FieldLastErrorMessage: + m.ClearLastErrorMessage() + return nil + case batchimagejob.FieldSubmittedAt: + m.ClearSubmittedAt() + return nil + case batchimagejob.FieldStartedAt: + m.ClearStartedAt() + return nil + case batchimagejob.FieldFinishedAt: + m.ClearFinishedAt() + return nil + case batchimagejob.FieldSettledAt: + m.ClearSettledAt() + return nil + } + return fmt.Errorf("unknown BatchImageJob nullable field %s", name) +} + +// ResetField resets all changes in the mutation for the field with the given name. +// It returns an error if the field is not defined in the schema. +func (m *BatchImageJobMutation) ResetField(name string) error { + switch name { + case batchimagejob.FieldBatchID: + m.ResetBatchID() + return nil + case batchimagejob.FieldUserID: + m.ResetUserID() + return nil + case batchimagejob.FieldAPIKeyID: + m.ResetAPIKeyID() + return nil + case batchimagejob.FieldAccountID: + m.ResetAccountID() + return nil + case batchimagejob.FieldProvider: + m.ResetProvider() + return nil + case batchimagejob.FieldModel: + m.ResetModel() + return nil + case batchimagejob.FieldTaskName: + m.ResetTaskName() + return nil + case batchimagejob.FieldStatus: + m.ResetStatus() + return nil + case batchimagejob.FieldProviderJobName: + m.ResetProviderJobName() + return nil + case batchimagejob.FieldProviderInputRef: + m.ResetProviderInputRef() + return nil + case batchimagejob.FieldProviderOutputRef: + m.ResetProviderOutputRef() + return nil + case batchimagejob.FieldGcsInputURI: + m.ResetGcsInputURI() + return nil + case batchimagejob.FieldGcsOutputURI: + m.ResetGcsOutputURI() + return nil + case batchimagejob.FieldItemCount: + m.ResetItemCount() + return nil + case batchimagejob.FieldSuccessCount: + m.ResetSuccessCount() + return nil + case batchimagejob.FieldFailCount: + m.ResetFailCount() + return nil + case batchimagejob.FieldCancelledCount: + m.ResetCancelledCount() + return nil + case batchimagejob.FieldEstimatedCost: + m.ResetEstimatedCost() + return nil + case batchimagejob.FieldHoldAmount: + m.ResetHoldAmount() + return nil + case batchimagejob.FieldActualCost: + m.ResetActualCost() + return nil + case batchimagejob.FieldCurrency: + m.ResetCurrency() + return nil + case batchimagejob.FieldHoldID: + m.ResetHoldID() + return nil + case batchimagejob.FieldIdempotencyKey: + m.ResetIdempotencyKey() + return nil + case batchimagejob.FieldRequestHash: + m.ResetRequestHash() + return nil + case batchimagejob.FieldManifestHash: + m.ResetManifestHash() + return nil + case batchimagejob.FieldRetryCount: + m.ResetRetryCount() + return nil + case batchimagejob.FieldVersion: + m.ResetVersion() + return nil + case batchimagejob.FieldOutputExpiresAt: + m.ResetOutputExpiresAt() + return nil + case batchimagejob.FieldInputDeletedAt: + m.ResetInputDeletedAt() + return nil + case batchimagejob.FieldOutputDeletedAt: + m.ResetOutputDeletedAt() + return nil + case batchimagejob.FieldDownloadedAt: + m.ResetDownloadedAt() + return nil + case batchimagejob.FieldUserDeletedAt: + m.ResetUserDeletedAt() + return nil + case batchimagejob.FieldLastErrorCode: + m.ResetLastErrorCode() + return nil + case batchimagejob.FieldLastErrorMessage: + m.ResetLastErrorMessage() + return nil + case batchimagejob.FieldCreatedAt: + m.ResetCreatedAt() + return nil + case batchimagejob.FieldUpdatedAt: + m.ResetUpdatedAt() + return nil + case batchimagejob.FieldSubmittedAt: + m.ResetSubmittedAt() + return nil + case batchimagejob.FieldStartedAt: + m.ResetStartedAt() + return nil + case batchimagejob.FieldFinishedAt: + m.ResetFinishedAt() + return nil + case batchimagejob.FieldSettledAt: + m.ResetSettledAt() + return nil + } + return fmt.Errorf("unknown BatchImageJob field %s", name) +} + +// AddedEdges returns all edge names that were set/added in this mutation. +func (m *BatchImageJobMutation) AddedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// AddedIDs returns all IDs (to other nodes) that were added for the given edge +// name in this mutation. +func (m *BatchImageJobMutation) AddedIDs(name string) []ent.Value { + return nil +} + +// RemovedEdges returns all edge names that were removed in this mutation. +func (m *BatchImageJobMutation) RemovedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// RemovedIDs returns all IDs (to other nodes) that were removed for the edge with +// the given name in this mutation. +func (m *BatchImageJobMutation) RemovedIDs(name string) []ent.Value { + return nil +} + +// ClearedEdges returns all edge names that were cleared in this mutation. +func (m *BatchImageJobMutation) ClearedEdges() []string { + edges := make([]string, 0, 0) + return edges +} + +// EdgeCleared returns a boolean which indicates if the edge with the given name +// was cleared in this mutation. +func (m *BatchImageJobMutation) EdgeCleared(name string) bool { + return false +} + +// ClearEdge clears the value of the edge with the given name. It returns an error +// if that edge is not defined in the schema. +func (m *BatchImageJobMutation) ClearEdge(name string) error { + return fmt.Errorf("unknown BatchImageJob unique edge %s", name) +} + +// ResetEdge resets all changes to the edge with the given name in this mutation. +// It returns an error if the edge is not defined in the schema. +func (m *BatchImageJobMutation) ResetEdge(name string) error { + return fmt.Errorf("unknown BatchImageJob edge %s", name) +} + // ChannelMonitorMutation represents an operation that mutates the ChannelMonitor nodes in the graph. type ChannelMonitorMutation struct { config @@ -15343,6 +20819,7 @@ type GroupMutation struct { default_validity_days *int adddefault_validity_days *int allow_image_generation *bool + allow_batch_image_generation *bool image_rate_independent *bool image_rate_multiplier *float64 addimage_rate_multiplier *float64 @@ -15352,6 +20829,19 @@ type GroupMutation struct { addimage_price_2k *float64 image_price_4k *float64 addimage_price_4k *float64 + batch_image_discount_multiplier *float64 + addbatch_image_discount_multiplier *float64 + batch_image_hold_multiplier *float64 + addbatch_image_hold_multiplier *float64 + video_rate_independent *bool + video_rate_multiplier *float64 + addvideo_rate_multiplier *float64 + video_price_480p *float64 + addvideo_price_480p *float64 + video_price_720p *float64 + addvideo_price_720p *float64 + video_price_1080p *float64 + addvideo_price_1080p *float64 claude_code_only *bool fallback_group_id *int64 addfallback_group_id *int64 @@ -16366,6 +21856,42 @@ func (m *GroupMutation) ResetAllowImageGeneration() { m.allow_image_generation = nil } +// SetAllowBatchImageGeneration sets the "allow_batch_image_generation" field. +func (m *GroupMutation) SetAllowBatchImageGeneration(b bool) { + m.allow_batch_image_generation = &b +} + +// AllowBatchImageGeneration returns the value of the "allow_batch_image_generation" field in the mutation. +func (m *GroupMutation) AllowBatchImageGeneration() (r bool, exists bool) { + v := m.allow_batch_image_generation + if v == nil { + return + } + return *v, true +} + +// OldAllowBatchImageGeneration returns the old "allow_batch_image_generation" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldAllowBatchImageGeneration(ctx context.Context) (v bool, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldAllowBatchImageGeneration is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldAllowBatchImageGeneration requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldAllowBatchImageGeneration: %w", err) + } + return oldValue.AllowBatchImageGeneration, nil +} + +// ResetAllowBatchImageGeneration resets all changes to the "allow_batch_image_generation" field. +func (m *GroupMutation) ResetAllowBatchImageGeneration() { + m.allow_batch_image_generation = nil +} + // SetImageRateIndependent sets the "image_rate_independent" field. func (m *GroupMutation) SetImageRateIndependent(b bool) { m.image_rate_independent = &b @@ -16668,6 +22194,420 @@ func (m *GroupMutation) ResetImagePrice4k() { delete(m.clearedFields, group.FieldImagePrice4k) } +// SetBatchImageDiscountMultiplier sets the "batch_image_discount_multiplier" field. +func (m *GroupMutation) SetBatchImageDiscountMultiplier(f float64) { + m.batch_image_discount_multiplier = &f + m.addbatch_image_discount_multiplier = nil +} + +// BatchImageDiscountMultiplier returns the value of the "batch_image_discount_multiplier" field in the mutation. +func (m *GroupMutation) BatchImageDiscountMultiplier() (r float64, exists bool) { + v := m.batch_image_discount_multiplier + if v == nil { + return + } + return *v, true +} + +// OldBatchImageDiscountMultiplier returns the old "batch_image_discount_multiplier" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldBatchImageDiscountMultiplier(ctx context.Context) (v float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldBatchImageDiscountMultiplier is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldBatchImageDiscountMultiplier requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldBatchImageDiscountMultiplier: %w", err) + } + return oldValue.BatchImageDiscountMultiplier, nil +} + +// AddBatchImageDiscountMultiplier adds f to the "batch_image_discount_multiplier" field. +func (m *GroupMutation) AddBatchImageDiscountMultiplier(f float64) { + if m.addbatch_image_discount_multiplier != nil { + *m.addbatch_image_discount_multiplier += f + } else { + m.addbatch_image_discount_multiplier = &f + } +} + +// AddedBatchImageDiscountMultiplier returns the value that was added to the "batch_image_discount_multiplier" field in this mutation. +func (m *GroupMutation) AddedBatchImageDiscountMultiplier() (r float64, exists bool) { + v := m.addbatch_image_discount_multiplier + if v == nil { + return + } + return *v, true +} + +// ResetBatchImageDiscountMultiplier resets all changes to the "batch_image_discount_multiplier" field. +func (m *GroupMutation) ResetBatchImageDiscountMultiplier() { + m.batch_image_discount_multiplier = nil + m.addbatch_image_discount_multiplier = nil +} + +// SetBatchImageHoldMultiplier sets the "batch_image_hold_multiplier" field. +func (m *GroupMutation) SetBatchImageHoldMultiplier(f float64) { + m.batch_image_hold_multiplier = &f + m.addbatch_image_hold_multiplier = nil +} + +// BatchImageHoldMultiplier returns the value of the "batch_image_hold_multiplier" field in the mutation. +func (m *GroupMutation) BatchImageHoldMultiplier() (r float64, exists bool) { + v := m.batch_image_hold_multiplier + if v == nil { + return + } + return *v, true +} + +// OldBatchImageHoldMultiplier returns the old "batch_image_hold_multiplier" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldBatchImageHoldMultiplier(ctx context.Context) (v float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldBatchImageHoldMultiplier is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldBatchImageHoldMultiplier requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldBatchImageHoldMultiplier: %w", err) + } + return oldValue.BatchImageHoldMultiplier, nil +} + +// AddBatchImageHoldMultiplier adds f to the "batch_image_hold_multiplier" field. +func (m *GroupMutation) AddBatchImageHoldMultiplier(f float64) { + if m.addbatch_image_hold_multiplier != nil { + *m.addbatch_image_hold_multiplier += f + } else { + m.addbatch_image_hold_multiplier = &f + } +} + +// AddedBatchImageHoldMultiplier returns the value that was added to the "batch_image_hold_multiplier" field in this mutation. +func (m *GroupMutation) AddedBatchImageHoldMultiplier() (r float64, exists bool) { + v := m.addbatch_image_hold_multiplier + if v == nil { + return + } + return *v, true +} + +// ResetBatchImageHoldMultiplier resets all changes to the "batch_image_hold_multiplier" field. +func (m *GroupMutation) ResetBatchImageHoldMultiplier() { + m.batch_image_hold_multiplier = nil + m.addbatch_image_hold_multiplier = nil +} + +// SetVideoRateIndependent sets the "video_rate_independent" field. +func (m *GroupMutation) SetVideoRateIndependent(b bool) { + m.video_rate_independent = &b +} + +// VideoRateIndependent returns the value of the "video_rate_independent" field in the mutation. +func (m *GroupMutation) VideoRateIndependent() (r bool, exists bool) { + v := m.video_rate_independent + if v == nil { + return + } + return *v, true +} + +// OldVideoRateIndependent returns the old "video_rate_independent" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldVideoRateIndependent(ctx context.Context) (v bool, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoRateIndependent is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoRateIndependent requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoRateIndependent: %w", err) + } + return oldValue.VideoRateIndependent, nil +} + +// ResetVideoRateIndependent resets all changes to the "video_rate_independent" field. +func (m *GroupMutation) ResetVideoRateIndependent() { + m.video_rate_independent = nil +} + +// SetVideoRateMultiplier sets the "video_rate_multiplier" field. +func (m *GroupMutation) SetVideoRateMultiplier(f float64) { + m.video_rate_multiplier = &f + m.addvideo_rate_multiplier = nil +} + +// VideoRateMultiplier returns the value of the "video_rate_multiplier" field in the mutation. +func (m *GroupMutation) VideoRateMultiplier() (r float64, exists bool) { + v := m.video_rate_multiplier + if v == nil { + return + } + return *v, true +} + +// OldVideoRateMultiplier returns the old "video_rate_multiplier" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldVideoRateMultiplier(ctx context.Context) (v float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoRateMultiplier is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoRateMultiplier requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoRateMultiplier: %w", err) + } + return oldValue.VideoRateMultiplier, nil +} + +// AddVideoRateMultiplier adds f to the "video_rate_multiplier" field. +func (m *GroupMutation) AddVideoRateMultiplier(f float64) { + if m.addvideo_rate_multiplier != nil { + *m.addvideo_rate_multiplier += f + } else { + m.addvideo_rate_multiplier = &f + } +} + +// AddedVideoRateMultiplier returns the value that was added to the "video_rate_multiplier" field in this mutation. +func (m *GroupMutation) AddedVideoRateMultiplier() (r float64, exists bool) { + v := m.addvideo_rate_multiplier + if v == nil { + return + } + return *v, true +} + +// ResetVideoRateMultiplier resets all changes to the "video_rate_multiplier" field. +func (m *GroupMutation) ResetVideoRateMultiplier() { + m.video_rate_multiplier = nil + m.addvideo_rate_multiplier = nil +} + +// SetVideoPrice480p sets the "video_price_480p" field. +func (m *GroupMutation) SetVideoPrice480p(f float64) { + m.video_price_480p = &f + m.addvideo_price_480p = nil +} + +// VideoPrice480p returns the value of the "video_price_480p" field in the mutation. +func (m *GroupMutation) VideoPrice480p() (r float64, exists bool) { + v := m.video_price_480p + if v == nil { + return + } + return *v, true +} + +// OldVideoPrice480p returns the old "video_price_480p" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldVideoPrice480p(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoPrice480p is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoPrice480p requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoPrice480p: %w", err) + } + return oldValue.VideoPrice480p, nil +} + +// AddVideoPrice480p adds f to the "video_price_480p" field. +func (m *GroupMutation) AddVideoPrice480p(f float64) { + if m.addvideo_price_480p != nil { + *m.addvideo_price_480p += f + } else { + m.addvideo_price_480p = &f + } +} + +// AddedVideoPrice480p returns the value that was added to the "video_price_480p" field in this mutation. +func (m *GroupMutation) AddedVideoPrice480p() (r float64, exists bool) { + v := m.addvideo_price_480p + if v == nil { + return + } + return *v, true +} + +// ClearVideoPrice480p clears the value of the "video_price_480p" field. +func (m *GroupMutation) ClearVideoPrice480p() { + m.video_price_480p = nil + m.addvideo_price_480p = nil + m.clearedFields[group.FieldVideoPrice480p] = struct{}{} +} + +// VideoPrice480pCleared returns if the "video_price_480p" field was cleared in this mutation. +func (m *GroupMutation) VideoPrice480pCleared() bool { + _, ok := m.clearedFields[group.FieldVideoPrice480p] + return ok +} + +// ResetVideoPrice480p resets all changes to the "video_price_480p" field. +func (m *GroupMutation) ResetVideoPrice480p() { + m.video_price_480p = nil + m.addvideo_price_480p = nil + delete(m.clearedFields, group.FieldVideoPrice480p) +} + +// SetVideoPrice720p sets the "video_price_720p" field. +func (m *GroupMutation) SetVideoPrice720p(f float64) { + m.video_price_720p = &f + m.addvideo_price_720p = nil +} + +// VideoPrice720p returns the value of the "video_price_720p" field in the mutation. +func (m *GroupMutation) VideoPrice720p() (r float64, exists bool) { + v := m.video_price_720p + if v == nil { + return + } + return *v, true +} + +// OldVideoPrice720p returns the old "video_price_720p" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldVideoPrice720p(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoPrice720p is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoPrice720p requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoPrice720p: %w", err) + } + return oldValue.VideoPrice720p, nil +} + +// AddVideoPrice720p adds f to the "video_price_720p" field. +func (m *GroupMutation) AddVideoPrice720p(f float64) { + if m.addvideo_price_720p != nil { + *m.addvideo_price_720p += f + } else { + m.addvideo_price_720p = &f + } +} + +// AddedVideoPrice720p returns the value that was added to the "video_price_720p" field in this mutation. +func (m *GroupMutation) AddedVideoPrice720p() (r float64, exists bool) { + v := m.addvideo_price_720p + if v == nil { + return + } + return *v, true +} + +// ClearVideoPrice720p clears the value of the "video_price_720p" field. +func (m *GroupMutation) ClearVideoPrice720p() { + m.video_price_720p = nil + m.addvideo_price_720p = nil + m.clearedFields[group.FieldVideoPrice720p] = struct{}{} +} + +// VideoPrice720pCleared returns if the "video_price_720p" field was cleared in this mutation. +func (m *GroupMutation) VideoPrice720pCleared() bool { + _, ok := m.clearedFields[group.FieldVideoPrice720p] + return ok +} + +// ResetVideoPrice720p resets all changes to the "video_price_720p" field. +func (m *GroupMutation) ResetVideoPrice720p() { + m.video_price_720p = nil + m.addvideo_price_720p = nil + delete(m.clearedFields, group.FieldVideoPrice720p) +} + +// SetVideoPrice1080p sets the "video_price_1080p" field. +func (m *GroupMutation) SetVideoPrice1080p(f float64) { + m.video_price_1080p = &f + m.addvideo_price_1080p = nil +} + +// VideoPrice1080p returns the value of the "video_price_1080p" field in the mutation. +func (m *GroupMutation) VideoPrice1080p() (r float64, exists bool) { + v := m.video_price_1080p + if v == nil { + return + } + return *v, true +} + +// OldVideoPrice1080p returns the old "video_price_1080p" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldVideoPrice1080p(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoPrice1080p is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoPrice1080p requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoPrice1080p: %w", err) + } + return oldValue.VideoPrice1080p, nil +} + +// AddVideoPrice1080p adds f to the "video_price_1080p" field. +func (m *GroupMutation) AddVideoPrice1080p(f float64) { + if m.addvideo_price_1080p != nil { + *m.addvideo_price_1080p += f + } else { + m.addvideo_price_1080p = &f + } +} + +// AddedVideoPrice1080p returns the value that was added to the "video_price_1080p" field in this mutation. +func (m *GroupMutation) AddedVideoPrice1080p() (r float64, exists bool) { + v := m.addvideo_price_1080p + if v == nil { + return + } + return *v, true +} + +// ClearVideoPrice1080p clears the value of the "video_price_1080p" field. +func (m *GroupMutation) ClearVideoPrice1080p() { + m.video_price_1080p = nil + m.addvideo_price_1080p = nil + m.clearedFields[group.FieldVideoPrice1080p] = struct{}{} +} + +// VideoPrice1080pCleared returns if the "video_price_1080p" field was cleared in this mutation. +func (m *GroupMutation) VideoPrice1080pCleared() bool { + _, ok := m.clearedFields[group.FieldVideoPrice1080p] + return ok +} + +// ResetVideoPrice1080p resets all changes to the "video_price_1080p" field. +func (m *GroupMutation) ResetVideoPrice1080p() { + m.video_price_1080p = nil + m.addvideo_price_1080p = nil + delete(m.clearedFields, group.FieldVideoPrice1080p) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (m *GroupMutation) SetClaudeCodeOnly(b bool) { m.claude_code_only = &b @@ -17702,7 +23642,7 @@ func (m *GroupMutation) Type() string { // order to get all numeric fields that were incremented/decremented, call // AddedFields(). func (m *GroupMutation) Fields() []string { - fields := make([]string, 0, 39) + fields := make([]string, 0, 47) if m.created_at != nil { fields = append(fields, group.FieldCreatedAt) } @@ -17760,6 +23700,9 @@ func (m *GroupMutation) Fields() []string { if m.allow_image_generation != nil { fields = append(fields, group.FieldAllowImageGeneration) } + if m.allow_batch_image_generation != nil { + fields = append(fields, group.FieldAllowBatchImageGeneration) + } if m.image_rate_independent != nil { fields = append(fields, group.FieldImageRateIndependent) } @@ -17775,6 +23718,27 @@ func (m *GroupMutation) Fields() []string { if m.image_price_4k != nil { fields = append(fields, group.FieldImagePrice4k) } + if m.batch_image_discount_multiplier != nil { + fields = append(fields, group.FieldBatchImageDiscountMultiplier) + } + if m.batch_image_hold_multiplier != nil { + fields = append(fields, group.FieldBatchImageHoldMultiplier) + } + if m.video_rate_independent != nil { + fields = append(fields, group.FieldVideoRateIndependent) + } + if m.video_rate_multiplier != nil { + fields = append(fields, group.FieldVideoRateMultiplier) + } + if m.video_price_480p != nil { + fields = append(fields, group.FieldVideoPrice480p) + } + if m.video_price_720p != nil { + fields = append(fields, group.FieldVideoPrice720p) + } + if m.video_price_1080p != nil { + fields = append(fields, group.FieldVideoPrice1080p) + } if m.claude_code_only != nil { fields = append(fields, group.FieldClaudeCodeOnly) } @@ -17866,6 +23830,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) { return m.DefaultValidityDays() case group.FieldAllowImageGeneration: return m.AllowImageGeneration() + case group.FieldAllowBatchImageGeneration: + return m.AllowBatchImageGeneration() case group.FieldImageRateIndependent: return m.ImageRateIndependent() case group.FieldImageRateMultiplier: @@ -17876,6 +23842,20 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) { return m.ImagePrice2k() case group.FieldImagePrice4k: return m.ImagePrice4k() + case group.FieldBatchImageDiscountMultiplier: + return m.BatchImageDiscountMultiplier() + case group.FieldBatchImageHoldMultiplier: + return m.BatchImageHoldMultiplier() + case group.FieldVideoRateIndependent: + return m.VideoRateIndependent() + case group.FieldVideoRateMultiplier: + return m.VideoRateMultiplier() + case group.FieldVideoPrice480p: + return m.VideoPrice480p() + case group.FieldVideoPrice720p: + return m.VideoPrice720p() + case group.FieldVideoPrice1080p: + return m.VideoPrice1080p() case group.FieldClaudeCodeOnly: return m.ClaudeCodeOnly() case group.FieldFallbackGroupID: @@ -17953,6 +23933,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e return m.OldDefaultValidityDays(ctx) case group.FieldAllowImageGeneration: return m.OldAllowImageGeneration(ctx) + case group.FieldAllowBatchImageGeneration: + return m.OldAllowBatchImageGeneration(ctx) case group.FieldImageRateIndependent: return m.OldImageRateIndependent(ctx) case group.FieldImageRateMultiplier: @@ -17963,6 +23945,20 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e return m.OldImagePrice2k(ctx) case group.FieldImagePrice4k: return m.OldImagePrice4k(ctx) + case group.FieldBatchImageDiscountMultiplier: + return m.OldBatchImageDiscountMultiplier(ctx) + case group.FieldBatchImageHoldMultiplier: + return m.OldBatchImageHoldMultiplier(ctx) + case group.FieldVideoRateIndependent: + return m.OldVideoRateIndependent(ctx) + case group.FieldVideoRateMultiplier: + return m.OldVideoRateMultiplier(ctx) + case group.FieldVideoPrice480p: + return m.OldVideoPrice480p(ctx) + case group.FieldVideoPrice720p: + return m.OldVideoPrice720p(ctx) + case group.FieldVideoPrice1080p: + return m.OldVideoPrice1080p(ctx) case group.FieldClaudeCodeOnly: return m.OldClaudeCodeOnly(ctx) case group.FieldFallbackGroupID: @@ -18135,6 +24131,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error { } m.SetAllowImageGeneration(v) return nil + case group.FieldAllowBatchImageGeneration: + v, ok := value.(bool) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetAllowBatchImageGeneration(v) + return nil case group.FieldImageRateIndependent: v, ok := value.(bool) if !ok { @@ -18170,6 +24173,55 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error { } m.SetImagePrice4k(v) return nil + case group.FieldBatchImageDiscountMultiplier: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetBatchImageDiscountMultiplier(v) + return nil + case group.FieldBatchImageHoldMultiplier: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetBatchImageHoldMultiplier(v) + return nil + case group.FieldVideoRateIndependent: + v, ok := value.(bool) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoRateIndependent(v) + return nil + case group.FieldVideoRateMultiplier: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoRateMultiplier(v) + return nil + case group.FieldVideoPrice480p: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoPrice480p(v) + return nil + case group.FieldVideoPrice720p: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoPrice720p(v) + return nil + case group.FieldVideoPrice1080p: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoPrice1080p(v) + return nil case group.FieldClaudeCodeOnly: v, ok := value.(bool) if !ok { @@ -18313,6 +24365,24 @@ func (m *GroupMutation) AddedFields() []string { if m.addimage_price_4k != nil { fields = append(fields, group.FieldImagePrice4k) } + if m.addbatch_image_discount_multiplier != nil { + fields = append(fields, group.FieldBatchImageDiscountMultiplier) + } + if m.addbatch_image_hold_multiplier != nil { + fields = append(fields, group.FieldBatchImageHoldMultiplier) + } + if m.addvideo_rate_multiplier != nil { + fields = append(fields, group.FieldVideoRateMultiplier) + } + if m.addvideo_price_480p != nil { + fields = append(fields, group.FieldVideoPrice480p) + } + if m.addvideo_price_720p != nil { + fields = append(fields, group.FieldVideoPrice720p) + } + if m.addvideo_price_1080p != nil { + fields = append(fields, group.FieldVideoPrice1080p) + } if m.addfallback_group_id != nil { fields = append(fields, group.FieldFallbackGroupID) } @@ -18353,6 +24423,18 @@ func (m *GroupMutation) AddedField(name string) (ent.Value, bool) { return m.AddedImagePrice2k() case group.FieldImagePrice4k: return m.AddedImagePrice4k() + case group.FieldBatchImageDiscountMultiplier: + return m.AddedBatchImageDiscountMultiplier() + case group.FieldBatchImageHoldMultiplier: + return m.AddedBatchImageHoldMultiplier() + case group.FieldVideoRateMultiplier: + return m.AddedVideoRateMultiplier() + case group.FieldVideoPrice480p: + return m.AddedVideoPrice480p() + case group.FieldVideoPrice720p: + return m.AddedVideoPrice720p() + case group.FieldVideoPrice1080p: + return m.AddedVideoPrice1080p() case group.FieldFallbackGroupID: return m.AddedFallbackGroupID() case group.FieldFallbackGroupIDOnInvalidRequest: @@ -18440,6 +24522,48 @@ func (m *GroupMutation) AddField(name string, value ent.Value) error { } m.AddImagePrice4k(v) return nil + case group.FieldBatchImageDiscountMultiplier: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddBatchImageDiscountMultiplier(v) + return nil + case group.FieldBatchImageHoldMultiplier: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddBatchImageHoldMultiplier(v) + return nil + case group.FieldVideoRateMultiplier: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVideoRateMultiplier(v) + return nil + case group.FieldVideoPrice480p: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVideoPrice480p(v) + return nil + case group.FieldVideoPrice720p: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVideoPrice720p(v) + return nil + case group.FieldVideoPrice1080p: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVideoPrice1080p(v) + return nil case group.FieldFallbackGroupID: v, ok := value.(int64) if !ok { @@ -18500,6 +24624,15 @@ func (m *GroupMutation) ClearedFields() []string { if m.FieldCleared(group.FieldImagePrice4k) { fields = append(fields, group.FieldImagePrice4k) } + if m.FieldCleared(group.FieldVideoPrice480p) { + fields = append(fields, group.FieldVideoPrice480p) + } + if m.FieldCleared(group.FieldVideoPrice720p) { + fields = append(fields, group.FieldVideoPrice720p) + } + if m.FieldCleared(group.FieldVideoPrice1080p) { + fields = append(fields, group.FieldVideoPrice1080p) + } if m.FieldCleared(group.FieldFallbackGroupID) { fields = append(fields, group.FieldFallbackGroupID) } @@ -18547,6 +24680,15 @@ func (m *GroupMutation) ClearField(name string) error { case group.FieldImagePrice4k: m.ClearImagePrice4k() return nil + case group.FieldVideoPrice480p: + m.ClearVideoPrice480p() + return nil + case group.FieldVideoPrice720p: + m.ClearVideoPrice720p() + return nil + case group.FieldVideoPrice1080p: + m.ClearVideoPrice1080p() + return nil case group.FieldFallbackGroupID: m.ClearFallbackGroupID() return nil @@ -18621,6 +24763,9 @@ func (m *GroupMutation) ResetField(name string) error { case group.FieldAllowImageGeneration: m.ResetAllowImageGeneration() return nil + case group.FieldAllowBatchImageGeneration: + m.ResetAllowBatchImageGeneration() + return nil case group.FieldImageRateIndependent: m.ResetImageRateIndependent() return nil @@ -18636,6 +24781,27 @@ func (m *GroupMutation) ResetField(name string) error { case group.FieldImagePrice4k: m.ResetImagePrice4k() return nil + case group.FieldBatchImageDiscountMultiplier: + m.ResetBatchImageDiscountMultiplier() + return nil + case group.FieldBatchImageHoldMultiplier: + m.ResetBatchImageHoldMultiplier() + return nil + case group.FieldVideoRateIndependent: + m.ResetVideoRateIndependent() + return nil + case group.FieldVideoRateMultiplier: + m.ResetVideoRateMultiplier() + return nil + case group.FieldVideoPrice480p: + m.ResetVideoPrice480p() + return nil + case group.FieldVideoPrice720p: + m.ResetVideoPrice720p() + return nil + case group.FieldVideoPrice1080p: + m.ResetVideoPrice1080p() + return nil case group.FieldClaudeCodeOnly: m.ResetClaudeCodeOnly() return nil @@ -35546,6 +41712,11 @@ type UsageLogMutation struct { image_output_size *string image_size_source *string image_size_breakdown *map[string]int + video_count *int + addvideo_count *int + video_resolution *string + video_duration_seconds *int + addvideo_duration_seconds *int cache_ttl_overridden *bool created_at *time.Time clearedFields map[string]struct{} @@ -37684,6 +43855,181 @@ func (m *UsageLogMutation) ResetImageSizeBreakdown() { delete(m.clearedFields, usagelog.FieldImageSizeBreakdown) } +// SetVideoCount sets the "video_count" field. +func (m *UsageLogMutation) SetVideoCount(i int) { + m.video_count = &i + m.addvideo_count = nil +} + +// VideoCount returns the value of the "video_count" field in the mutation. +func (m *UsageLogMutation) VideoCount() (r int, exists bool) { + v := m.video_count + if v == nil { + return + } + return *v, true +} + +// OldVideoCount returns the old "video_count" field's value of the UsageLog entity. +// If the UsageLog object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *UsageLogMutation) OldVideoCount(ctx context.Context) (v int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoCount is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoCount requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoCount: %w", err) + } + return oldValue.VideoCount, nil +} + +// AddVideoCount adds i to the "video_count" field. +func (m *UsageLogMutation) AddVideoCount(i int) { + if m.addvideo_count != nil { + *m.addvideo_count += i + } else { + m.addvideo_count = &i + } +} + +// AddedVideoCount returns the value that was added to the "video_count" field in this mutation. +func (m *UsageLogMutation) AddedVideoCount() (r int, exists bool) { + v := m.addvideo_count + if v == nil { + return + } + return *v, true +} + +// ResetVideoCount resets all changes to the "video_count" field. +func (m *UsageLogMutation) ResetVideoCount() { + m.video_count = nil + m.addvideo_count = nil +} + +// SetVideoResolution sets the "video_resolution" field. +func (m *UsageLogMutation) SetVideoResolution(s string) { + m.video_resolution = &s +} + +// VideoResolution returns the value of the "video_resolution" field in the mutation. +func (m *UsageLogMutation) VideoResolution() (r string, exists bool) { + v := m.video_resolution + if v == nil { + return + } + return *v, true +} + +// OldVideoResolution returns the old "video_resolution" field's value of the UsageLog entity. +// If the UsageLog object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *UsageLogMutation) OldVideoResolution(ctx context.Context) (v *string, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoResolution is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoResolution requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoResolution: %w", err) + } + return oldValue.VideoResolution, nil +} + +// ClearVideoResolution clears the value of the "video_resolution" field. +func (m *UsageLogMutation) ClearVideoResolution() { + m.video_resolution = nil + m.clearedFields[usagelog.FieldVideoResolution] = struct{}{} +} + +// VideoResolutionCleared returns if the "video_resolution" field was cleared in this mutation. +func (m *UsageLogMutation) VideoResolutionCleared() bool { + _, ok := m.clearedFields[usagelog.FieldVideoResolution] + return ok +} + +// ResetVideoResolution resets all changes to the "video_resolution" field. +func (m *UsageLogMutation) ResetVideoResolution() { + m.video_resolution = nil + delete(m.clearedFields, usagelog.FieldVideoResolution) +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (m *UsageLogMutation) SetVideoDurationSeconds(i int) { + m.video_duration_seconds = &i + m.addvideo_duration_seconds = nil +} + +// VideoDurationSeconds returns the value of the "video_duration_seconds" field in the mutation. +func (m *UsageLogMutation) VideoDurationSeconds() (r int, exists bool) { + v := m.video_duration_seconds + if v == nil { + return + } + return *v, true +} + +// OldVideoDurationSeconds returns the old "video_duration_seconds" field's value of the UsageLog entity. +// If the UsageLog object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *UsageLogMutation) OldVideoDurationSeconds(ctx context.Context) (v *int, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldVideoDurationSeconds is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldVideoDurationSeconds requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldVideoDurationSeconds: %w", err) + } + return oldValue.VideoDurationSeconds, nil +} + +// AddVideoDurationSeconds adds i to the "video_duration_seconds" field. +func (m *UsageLogMutation) AddVideoDurationSeconds(i int) { + if m.addvideo_duration_seconds != nil { + *m.addvideo_duration_seconds += i + } else { + m.addvideo_duration_seconds = &i + } +} + +// AddedVideoDurationSeconds returns the value that was added to the "video_duration_seconds" field in this mutation. +func (m *UsageLogMutation) AddedVideoDurationSeconds() (r int, exists bool) { + v := m.addvideo_duration_seconds + if v == nil { + return + } + return *v, true +} + +// ClearVideoDurationSeconds clears the value of the "video_duration_seconds" field. +func (m *UsageLogMutation) ClearVideoDurationSeconds() { + m.video_duration_seconds = nil + m.addvideo_duration_seconds = nil + m.clearedFields[usagelog.FieldVideoDurationSeconds] = struct{}{} +} + +// VideoDurationSecondsCleared returns if the "video_duration_seconds" field was cleared in this mutation. +func (m *UsageLogMutation) VideoDurationSecondsCleared() bool { + _, ok := m.clearedFields[usagelog.FieldVideoDurationSeconds] + return ok +} + +// ResetVideoDurationSeconds resets all changes to the "video_duration_seconds" field. +func (m *UsageLogMutation) ResetVideoDurationSeconds() { + m.video_duration_seconds = nil + m.addvideo_duration_seconds = nil + delete(m.clearedFields, usagelog.FieldVideoDurationSeconds) +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (m *UsageLogMutation) SetCacheTTLOverridden(b bool) { m.cache_ttl_overridden = &b @@ -37925,7 +44271,7 @@ func (m *UsageLogMutation) Type() string { // order to get all numeric fields that were incremented/decremented, call // AddedFields(). func (m *UsageLogMutation) Fields() []string { - fields := make([]string, 0, 41) + fields := make([]string, 0, 44) if m.user != nil { fields = append(fields, usagelog.FieldUserID) } @@ -38043,6 +44389,15 @@ func (m *UsageLogMutation) Fields() []string { if m.image_size_breakdown != nil { fields = append(fields, usagelog.FieldImageSizeBreakdown) } + if m.video_count != nil { + fields = append(fields, usagelog.FieldVideoCount) + } + if m.video_resolution != nil { + fields = append(fields, usagelog.FieldVideoResolution) + } + if m.video_duration_seconds != nil { + fields = append(fields, usagelog.FieldVideoDurationSeconds) + } if m.cache_ttl_overridden != nil { fields = append(fields, usagelog.FieldCacheTTLOverridden) } @@ -38135,6 +44490,12 @@ func (m *UsageLogMutation) Field(name string) (ent.Value, bool) { return m.ImageSizeSource() case usagelog.FieldImageSizeBreakdown: return m.ImageSizeBreakdown() + case usagelog.FieldVideoCount: + return m.VideoCount() + case usagelog.FieldVideoResolution: + return m.VideoResolution() + case usagelog.FieldVideoDurationSeconds: + return m.VideoDurationSeconds() case usagelog.FieldCacheTTLOverridden: return m.CacheTTLOverridden() case usagelog.FieldCreatedAt: @@ -38226,6 +44587,12 @@ func (m *UsageLogMutation) OldField(ctx context.Context, name string) (ent.Value return m.OldImageSizeSource(ctx) case usagelog.FieldImageSizeBreakdown: return m.OldImageSizeBreakdown(ctx) + case usagelog.FieldVideoCount: + return m.OldVideoCount(ctx) + case usagelog.FieldVideoResolution: + return m.OldVideoResolution(ctx) + case usagelog.FieldVideoDurationSeconds: + return m.OldVideoDurationSeconds(ctx) case usagelog.FieldCacheTTLOverridden: return m.OldCacheTTLOverridden(ctx) case usagelog.FieldCreatedAt: @@ -38512,6 +44879,27 @@ func (m *UsageLogMutation) SetField(name string, value ent.Value) error { } m.SetImageSizeBreakdown(v) return nil + case usagelog.FieldVideoCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoCount(v) + return nil + case usagelog.FieldVideoResolution: + v, ok := value.(string) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoResolution(v) + return nil + case usagelog.FieldVideoDurationSeconds: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetVideoDurationSeconds(v) + return nil case usagelog.FieldCacheTTLOverridden: v, ok := value.(bool) if !ok { @@ -38591,6 +44979,12 @@ func (m *UsageLogMutation) AddedFields() []string { if m.addimage_count != nil { fields = append(fields, usagelog.FieldImageCount) } + if m.addvideo_count != nil { + fields = append(fields, usagelog.FieldVideoCount) + } + if m.addvideo_duration_seconds != nil { + fields = append(fields, usagelog.FieldVideoDurationSeconds) + } return fields } @@ -38637,6 +45031,10 @@ func (m *UsageLogMutation) AddedField(name string) (ent.Value, bool) { return m.AddedFirstTokenMs() case usagelog.FieldImageCount: return m.AddedImageCount() + case usagelog.FieldVideoCount: + return m.AddedVideoCount() + case usagelog.FieldVideoDurationSeconds: + return m.AddedVideoDurationSeconds() } return nil, false } @@ -38779,6 +45177,20 @@ func (m *UsageLogMutation) AddField(name string, value ent.Value) error { } m.AddImageCount(v) return nil + case usagelog.FieldVideoCount: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVideoCount(v) + return nil + case usagelog.FieldVideoDurationSeconds: + v, ok := value.(int) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddVideoDurationSeconds(v) + return nil } return fmt.Errorf("unknown UsageLog numeric field %s", name) } @@ -38841,6 +45253,12 @@ func (m *UsageLogMutation) ClearedFields() []string { if m.FieldCleared(usagelog.FieldImageSizeBreakdown) { fields = append(fields, usagelog.FieldImageSizeBreakdown) } + if m.FieldCleared(usagelog.FieldVideoResolution) { + fields = append(fields, usagelog.FieldVideoResolution) + } + if m.FieldCleared(usagelog.FieldVideoDurationSeconds) { + fields = append(fields, usagelog.FieldVideoDurationSeconds) + } return fields } @@ -38909,6 +45327,12 @@ func (m *UsageLogMutation) ClearField(name string) error { case usagelog.FieldImageSizeBreakdown: m.ClearImageSizeBreakdown() return nil + case usagelog.FieldVideoResolution: + m.ClearVideoResolution() + return nil + case usagelog.FieldVideoDurationSeconds: + m.ClearVideoDurationSeconds() + return nil } return fmt.Errorf("unknown UsageLog nullable field %s", name) } @@ -39034,6 +45458,15 @@ func (m *UsageLogMutation) ResetField(name string) error { case usagelog.FieldImageSizeBreakdown: m.ResetImageSizeBreakdown() return nil + case usagelog.FieldVideoCount: + m.ResetVideoCount() + return nil + case usagelog.FieldVideoResolution: + m.ResetVideoResolution() + return nil + case usagelog.FieldVideoDurationSeconds: + m.ResetVideoDurationSeconds() + return nil case usagelog.FieldCacheTTLOverridden: m.ResetCacheTTLOverridden() return nil @@ -39204,6 +45637,8 @@ type UserMutation struct { role *string balance *float64 addbalance *float64 + frozen_balance *float64 + addfrozen_balance *float64 concurrency *int addconcurrency *int status *string @@ -39652,6 +46087,62 @@ func (m *UserMutation) ResetBalance() { m.addbalance = nil } +// SetFrozenBalance sets the "frozen_balance" field. +func (m *UserMutation) SetFrozenBalance(f float64) { + m.frozen_balance = &f + m.addfrozen_balance = nil +} + +// FrozenBalance returns the value of the "frozen_balance" field in the mutation. +func (m *UserMutation) FrozenBalance() (r float64, exists bool) { + v := m.frozen_balance + if v == nil { + return + } + return *v, true +} + +// OldFrozenBalance returns the old "frozen_balance" field's value of the User entity. +// If the User object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *UserMutation) OldFrozenBalance(ctx context.Context) (v float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldFrozenBalance is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldFrozenBalance requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldFrozenBalance: %w", err) + } + return oldValue.FrozenBalance, nil +} + +// AddFrozenBalance adds f to the "frozen_balance" field. +func (m *UserMutation) AddFrozenBalance(f float64) { + if m.addfrozen_balance != nil { + *m.addfrozen_balance += f + } else { + m.addfrozen_balance = &f + } +} + +// AddedFrozenBalance returns the value that was added to the "frozen_balance" field in this mutation. +func (m *UserMutation) AddedFrozenBalance() (r float64, exists bool) { + v := m.addfrozen_balance + if v == nil { + return + } + return *v, true +} + +// ResetFrozenBalance resets all changes to the "frozen_balance" field. +func (m *UserMutation) ResetFrozenBalance() { + m.frozen_balance = nil + m.addfrozen_balance = nil +} + // SetConcurrency sets the "concurrency" field. func (m *UserMutation) SetConcurrency(i int) { m.concurrency = &i @@ -41110,7 +47601,7 @@ func (m *UserMutation) Type() string { // order to get all numeric fields that were incremented/decremented, call // AddedFields(). func (m *UserMutation) Fields() []string { - fields := make([]string, 0, 23) + fields := make([]string, 0, 24) if m.created_at != nil { fields = append(fields, user.FieldCreatedAt) } @@ -41132,6 +47623,9 @@ func (m *UserMutation) Fields() []string { if m.balance != nil { fields = append(fields, user.FieldBalance) } + if m.frozen_balance != nil { + fields = append(fields, user.FieldFrozenBalance) + } if m.concurrency != nil { fields = append(fields, user.FieldConcurrency) } @@ -41202,6 +47696,8 @@ func (m *UserMutation) Field(name string) (ent.Value, bool) { return m.Role() case user.FieldBalance: return m.Balance() + case user.FieldFrozenBalance: + return m.FrozenBalance() case user.FieldConcurrency: return m.Concurrency() case user.FieldStatus: @@ -41257,6 +47753,8 @@ func (m *UserMutation) OldField(ctx context.Context, name string) (ent.Value, er return m.OldRole(ctx) case user.FieldBalance: return m.OldBalance(ctx) + case user.FieldFrozenBalance: + return m.OldFrozenBalance(ctx) case user.FieldConcurrency: return m.OldConcurrency(ctx) case user.FieldStatus: @@ -41347,6 +47845,13 @@ func (m *UserMutation) SetField(name string, value ent.Value) error { } m.SetBalance(v) return nil + case user.FieldFrozenBalance: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetFrozenBalance(v) + return nil case user.FieldConcurrency: v, ok := value.(int) if !ok { @@ -41470,6 +47975,9 @@ func (m *UserMutation) AddedFields() []string { if m.addbalance != nil { fields = append(fields, user.FieldBalance) } + if m.addfrozen_balance != nil { + fields = append(fields, user.FieldFrozenBalance) + } if m.addconcurrency != nil { fields = append(fields, user.FieldConcurrency) } @@ -41492,6 +48000,8 @@ func (m *UserMutation) AddedField(name string) (ent.Value, bool) { switch name { case user.FieldBalance: return m.AddedBalance() + case user.FieldFrozenBalance: + return m.AddedFrozenBalance() case user.FieldConcurrency: return m.AddedConcurrency() case user.FieldBalanceNotifyThreshold: @@ -41516,6 +48026,13 @@ func (m *UserMutation) AddField(name string, value ent.Value) error { } m.AddBalance(v) return nil + case user.FieldFrozenBalance: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddFrozenBalance(v) + return nil case user.FieldConcurrency: v, ok := value.(int) if !ok { @@ -41631,6 +48148,9 @@ func (m *UserMutation) ResetField(name string) error { case user.FieldBalance: m.ResetBalance() return nil + case user.FieldFrozenBalance: + m.ResetFrozenBalance() + return nil case user.FieldConcurrency: m.ResetConcurrency() return nil diff --git a/backend/ent/predicate/predicate.go b/backend/ent/predicate/predicate.go index ab4d7d1827..8d18d38151 100644 --- a/backend/ent/predicate/predicate.go +++ b/backend/ent/predicate/predicate.go @@ -27,6 +27,15 @@ type AuthIdentity func(*sql.Selector) // AuthIdentityChannel is the predicate function for authidentitychannel builders. type AuthIdentityChannel func(*sql.Selector) +// BatchImageEvent is the predicate function for batchimageevent builders. +type BatchImageEvent func(*sql.Selector) + +// BatchImageItem is the predicate function for batchimageitem builders. +type BatchImageItem func(*sql.Selector) + +// BatchImageJob is the predicate function for batchimagejob builders. +type BatchImageJob func(*sql.Selector) + // ChannelMonitor is the predicate function for channelmonitor builders. type ChannelMonitor func(*sql.Selector) diff --git a/backend/ent/runtime/runtime.go b/backend/ent/runtime/runtime.go index b86f6a0560..d47e7d143b 100644 --- a/backend/ent/runtime/runtime.go +++ b/backend/ent/runtime/runtime.go @@ -12,6 +12,9 @@ import ( "github.com/Wei-Shaw/sub2api/ent/apikey" "github.com/Wei-Shaw/sub2api/ent/authidentity" "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" + "github.com/Wei-Shaw/sub2api/ent/batchimageevent" + "github.com/Wei-Shaw/sub2api/ent/batchimageitem" + "github.com/Wei-Shaw/sub2api/ent/batchimagejob" "github.com/Wei-Shaw/sub2api/ent/channelmonitor" "github.com/Wei-Shaw/sub2api/ent/channelmonitordailyrollup" "github.com/Wei-Shaw/sub2api/ent/channelmonitorhistory" @@ -432,6 +435,172 @@ func init() { authidentitychannelDescMetadata := authidentitychannelFields[6].Descriptor() // authidentitychannel.DefaultMetadata holds the default value on creation for the metadata field. authidentitychannel.DefaultMetadata = authidentitychannelDescMetadata.Default.(func() map[string]interface{}) + batchimageeventFields := schema.BatchImageEvent{}.Fields() + _ = batchimageeventFields + // batchimageeventDescJobID is the schema descriptor for job_id field. + batchimageeventDescJobID := batchimageeventFields[0].Descriptor() + // batchimageevent.JobIDValidator is a validator for the "job_id" field. It is called by the builders before save. + batchimageevent.JobIDValidator = batchimageeventDescJobID.Validators[0].(func(string) error) + // batchimageeventDescEventType is the schema descriptor for event_type field. + batchimageeventDescEventType := batchimageeventFields[1].Descriptor() + // batchimageevent.EventTypeValidator is a validator for the "event_type" field. It is called by the builders before save. + batchimageevent.EventTypeValidator = batchimageeventDescEventType.Validators[0].(func(string) error) + // batchimageeventDescEventHash is the schema descriptor for event_hash field. + batchimageeventDescEventHash := batchimageeventFields[3].Descriptor() + // batchimageevent.EventHashValidator is a validator for the "event_hash" field. It is called by the builders before save. + batchimageevent.EventHashValidator = batchimageeventDescEventHash.Validators[0].(func(string) error) + // batchimageeventDescCreatedAt is the schema descriptor for created_at field. + batchimageeventDescCreatedAt := batchimageeventFields[4].Descriptor() + // batchimageevent.DefaultCreatedAt holds the default value on creation for the created_at field. + batchimageevent.DefaultCreatedAt = batchimageeventDescCreatedAt.Default.(func() time.Time) + batchimageitemFields := schema.BatchImageItem{}.Fields() + _ = batchimageitemFields + // batchimageitemDescJobID is the schema descriptor for job_id field. + batchimageitemDescJobID := batchimageitemFields[0].Descriptor() + // batchimageitem.JobIDValidator is a validator for the "job_id" field. It is called by the builders before save. + batchimageitem.JobIDValidator = batchimageitemDescJobID.Validators[0].(func(string) error) + // batchimageitemDescCustomID is the schema descriptor for custom_id field. + batchimageitemDescCustomID := batchimageitemFields[1].Descriptor() + // batchimageitem.CustomIDValidator is a validator for the "custom_id" field. It is called by the builders before save. + batchimageitem.CustomIDValidator = batchimageitemDescCustomID.Validators[0].(func(string) error) + // batchimageitemDescStatus is the schema descriptor for status field. + batchimageitemDescStatus := batchimageitemFields[2].Descriptor() + // batchimageitem.StatusValidator is a validator for the "status" field. It is called by the builders before save. + batchimageitem.StatusValidator = batchimageitemDescStatus.Validators[0].(func(string) error) + // batchimageitemDescRequestHash is the schema descriptor for request_hash field. + batchimageitemDescRequestHash := batchimageitemFields[3].Descriptor() + // batchimageitem.RequestHashValidator is a validator for the "request_hash" field. It is called by the builders before save. + batchimageitem.RequestHashValidator = batchimageitemDescRequestHash.Validators[0].(func(string) error) + // batchimageitemDescProviderSourceObject is the schema descriptor for provider_source_object field. + batchimageitemDescProviderSourceObject := batchimageitemFields[5].Descriptor() + // batchimageitem.ProviderSourceObjectValidator is a validator for the "provider_source_object" field. It is called by the builders before save. + batchimageitem.ProviderSourceObjectValidator = batchimageitemDescProviderSourceObject.Validators[0].(func(string) error) + // batchimageitemDescMimeType is the schema descriptor for mime_type field. + batchimageitemDescMimeType := batchimageitemFields[9].Descriptor() + // batchimageitem.MimeTypeValidator is a validator for the "mime_type" field. It is called by the builders before save. + batchimageitem.MimeTypeValidator = batchimageitemDescMimeType.Validators[0].(func(string) error) + // batchimageitemDescFileExtension is the schema descriptor for file_extension field. + batchimageitemDescFileExtension := batchimageitemFields[10].Descriptor() + // batchimageitem.FileExtensionValidator is a validator for the "file_extension" field. It is called by the builders before save. + batchimageitem.FileExtensionValidator = batchimageitemDescFileExtension.Validators[0].(func(string) error) + // batchimageitemDescImageCount is the schema descriptor for image_count field. + batchimageitemDescImageCount := batchimageitemFields[11].Descriptor() + // batchimageitem.DefaultImageCount holds the default value on creation for the image_count field. + batchimageitem.DefaultImageCount = batchimageitemDescImageCount.Default.(int) + // batchimageitemDescErrorCode is the schema descriptor for error_code field. + batchimageitemDescErrorCode := batchimageitemFields[12].Descriptor() + // batchimageitem.ErrorCodeValidator is a validator for the "error_code" field. It is called by the builders before save. + batchimageitem.ErrorCodeValidator = batchimageitemDescErrorCode.Validators[0].(func(string) error) + // batchimageitemDescCreatedAt is the schema descriptor for created_at field. + batchimageitemDescCreatedAt := batchimageitemFields[15].Descriptor() + // batchimageitem.DefaultCreatedAt holds the default value on creation for the created_at field. + batchimageitem.DefaultCreatedAt = batchimageitemDescCreatedAt.Default.(func() time.Time) + batchimagejobFields := schema.BatchImageJob{}.Fields() + _ = batchimagejobFields + // batchimagejobDescBatchID is the schema descriptor for batch_id field. + batchimagejobDescBatchID := batchimagejobFields[0].Descriptor() + // batchimagejob.BatchIDValidator is a validator for the "batch_id" field. It is called by the builders before save. + batchimagejob.BatchIDValidator = batchimagejobDescBatchID.Validators[0].(func(string) error) + // batchimagejobDescProvider is the schema descriptor for provider field. + batchimagejobDescProvider := batchimagejobFields[4].Descriptor() + // batchimagejob.ProviderValidator is a validator for the "provider" field. It is called by the builders before save. + batchimagejob.ProviderValidator = batchimagejobDescProvider.Validators[0].(func(string) error) + // batchimagejobDescModel is the schema descriptor for model field. + batchimagejobDescModel := batchimagejobFields[5].Descriptor() + // batchimagejob.ModelValidator is a validator for the "model" field. It is called by the builders before save. + batchimagejob.ModelValidator = batchimagejobDescModel.Validators[0].(func(string) error) + // batchimagejobDescTaskName is the schema descriptor for task_name field. + batchimagejobDescTaskName := batchimagejobFields[6].Descriptor() + // batchimagejob.DefaultTaskName holds the default value on creation for the task_name field. + batchimagejob.DefaultTaskName = batchimagejobDescTaskName.Default.(string) + // batchimagejob.TaskNameValidator is a validator for the "task_name" field. It is called by the builders before save. + batchimagejob.TaskNameValidator = batchimagejobDescTaskName.Validators[0].(func(string) error) + // batchimagejobDescStatus is the schema descriptor for status field. + batchimagejobDescStatus := batchimagejobFields[7].Descriptor() + // batchimagejob.DefaultStatus holds the default value on creation for the status field. + batchimagejob.DefaultStatus = batchimagejobDescStatus.Default.(string) + // batchimagejob.StatusValidator is a validator for the "status" field. It is called by the builders before save. + batchimagejob.StatusValidator = batchimagejobDescStatus.Validators[0].(func(string) error) + // batchimagejobDescProviderJobName is the schema descriptor for provider_job_name field. + batchimagejobDescProviderJobName := batchimagejobFields[8].Descriptor() + // batchimagejob.ProviderJobNameValidator is a validator for the "provider_job_name" field. It is called by the builders before save. + batchimagejob.ProviderJobNameValidator = batchimagejobDescProviderJobName.Validators[0].(func(string) error) + // batchimagejobDescProviderInputRef is the schema descriptor for provider_input_ref field. + batchimagejobDescProviderInputRef := batchimagejobFields[9].Descriptor() + // batchimagejob.ProviderInputRefValidator is a validator for the "provider_input_ref" field. It is called by the builders before save. + batchimagejob.ProviderInputRefValidator = batchimagejobDescProviderInputRef.Validators[0].(func(string) error) + // batchimagejobDescProviderOutputRef is the schema descriptor for provider_output_ref field. + batchimagejobDescProviderOutputRef := batchimagejobFields[10].Descriptor() + // batchimagejob.ProviderOutputRefValidator is a validator for the "provider_output_ref" field. It is called by the builders before save. + batchimagejob.ProviderOutputRefValidator = batchimagejobDescProviderOutputRef.Validators[0].(func(string) error) + // batchimagejobDescGcsInputURI is the schema descriptor for gcs_input_uri field. + batchimagejobDescGcsInputURI := batchimagejobFields[11].Descriptor() + // batchimagejob.GcsInputURIValidator is a validator for the "gcs_input_uri" field. It is called by the builders before save. + batchimagejob.GcsInputURIValidator = batchimagejobDescGcsInputURI.Validators[0].(func(string) error) + // batchimagejobDescGcsOutputURI is the schema descriptor for gcs_output_uri field. + batchimagejobDescGcsOutputURI := batchimagejobFields[12].Descriptor() + // batchimagejob.GcsOutputURIValidator is a validator for the "gcs_output_uri" field. It is called by the builders before save. + batchimagejob.GcsOutputURIValidator = batchimagejobDescGcsOutputURI.Validators[0].(func(string) error) + // batchimagejobDescSuccessCount is the schema descriptor for success_count field. + batchimagejobDescSuccessCount := batchimagejobFields[14].Descriptor() + // batchimagejob.DefaultSuccessCount holds the default value on creation for the success_count field. + batchimagejob.DefaultSuccessCount = batchimagejobDescSuccessCount.Default.(int) + // batchimagejobDescFailCount is the schema descriptor for fail_count field. + batchimagejobDescFailCount := batchimagejobFields[15].Descriptor() + // batchimagejob.DefaultFailCount holds the default value on creation for the fail_count field. + batchimagejob.DefaultFailCount = batchimagejobDescFailCount.Default.(int) + // batchimagejobDescCancelledCount is the schema descriptor for cancelled_count field. + batchimagejobDescCancelledCount := batchimagejobFields[16].Descriptor() + // batchimagejob.DefaultCancelledCount holds the default value on creation for the cancelled_count field. + batchimagejob.DefaultCancelledCount = batchimagejobDescCancelledCount.Default.(int) + // batchimagejobDescEstimatedCost is the schema descriptor for estimated_cost field. + batchimagejobDescEstimatedCost := batchimagejobFields[17].Descriptor() + // batchimagejob.DefaultEstimatedCost holds the default value on creation for the estimated_cost field. + batchimagejob.DefaultEstimatedCost = batchimagejobDescEstimatedCost.Default.(float64) + // batchimagejobDescCurrency is the schema descriptor for currency field. + batchimagejobDescCurrency := batchimagejobFields[20].Descriptor() + // batchimagejob.DefaultCurrency holds the default value on creation for the currency field. + batchimagejob.DefaultCurrency = batchimagejobDescCurrency.Default.(string) + // batchimagejob.CurrencyValidator is a validator for the "currency" field. It is called by the builders before save. + batchimagejob.CurrencyValidator = batchimagejobDescCurrency.Validators[0].(func(string) error) + // batchimagejobDescHoldID is the schema descriptor for hold_id field. + batchimagejobDescHoldID := batchimagejobFields[21].Descriptor() + // batchimagejob.HoldIDValidator is a validator for the "hold_id" field. It is called by the builders before save. + batchimagejob.HoldIDValidator = batchimagejobDescHoldID.Validators[0].(func(string) error) + // batchimagejobDescIdempotencyKey is the schema descriptor for idempotency_key field. + batchimagejobDescIdempotencyKey := batchimagejobFields[22].Descriptor() + // batchimagejob.IdempotencyKeyValidator is a validator for the "idempotency_key" field. It is called by the builders before save. + batchimagejob.IdempotencyKeyValidator = batchimagejobDescIdempotencyKey.Validators[0].(func(string) error) + // batchimagejobDescRequestHash is the schema descriptor for request_hash field. + batchimagejobDescRequestHash := batchimagejobFields[23].Descriptor() + // batchimagejob.RequestHashValidator is a validator for the "request_hash" field. It is called by the builders before save. + batchimagejob.RequestHashValidator = batchimagejobDescRequestHash.Validators[0].(func(string) error) + // batchimagejobDescManifestHash is the schema descriptor for manifest_hash field. + batchimagejobDescManifestHash := batchimagejobFields[24].Descriptor() + // batchimagejob.ManifestHashValidator is a validator for the "manifest_hash" field. It is called by the builders before save. + batchimagejob.ManifestHashValidator = batchimagejobDescManifestHash.Validators[0].(func(string) error) + // batchimagejobDescRetryCount is the schema descriptor for retry_count field. + batchimagejobDescRetryCount := batchimagejobFields[25].Descriptor() + // batchimagejob.DefaultRetryCount holds the default value on creation for the retry_count field. + batchimagejob.DefaultRetryCount = batchimagejobDescRetryCount.Default.(int) + // batchimagejobDescVersion is the schema descriptor for version field. + batchimagejobDescVersion := batchimagejobFields[26].Descriptor() + // batchimagejob.DefaultVersion holds the default value on creation for the version field. + batchimagejob.DefaultVersion = batchimagejobDescVersion.Default.(int) + // batchimagejobDescLastErrorCode is the schema descriptor for last_error_code field. + batchimagejobDescLastErrorCode := batchimagejobFields[32].Descriptor() + // batchimagejob.LastErrorCodeValidator is a validator for the "last_error_code" field. It is called by the builders before save. + batchimagejob.LastErrorCodeValidator = batchimagejobDescLastErrorCode.Validators[0].(func(string) error) + // batchimagejobDescCreatedAt is the schema descriptor for created_at field. + batchimagejobDescCreatedAt := batchimagejobFields[34].Descriptor() + // batchimagejob.DefaultCreatedAt holds the default value on creation for the created_at field. + batchimagejob.DefaultCreatedAt = batchimagejobDescCreatedAt.Default.(func() time.Time) + // batchimagejobDescUpdatedAt is the schema descriptor for updated_at field. + batchimagejobDescUpdatedAt := batchimagejobFields[35].Descriptor() + // batchimagejob.DefaultUpdatedAt holds the default value on creation for the updated_at field. + batchimagejob.DefaultUpdatedAt = batchimagejobDescUpdatedAt.Default.(func() time.Time) + // batchimagejob.UpdateDefaultUpdatedAt holds the default value on update for the updated_at field. + batchimagejob.UpdateDefaultUpdatedAt = batchimagejobDescUpdatedAt.UpdateDefault.(func() time.Time) channelmonitorMixin := schema.ChannelMonitor{}.Mixin() channelmonitorMixinFields0 := channelmonitorMixin[0].Fields() _ = channelmonitorMixinFields0 @@ -846,62 +1015,82 @@ func init() { groupDescAllowImageGeneration := groupFields[15].Descriptor() // group.DefaultAllowImageGeneration holds the default value on creation for the allow_image_generation field. group.DefaultAllowImageGeneration = groupDescAllowImageGeneration.Default.(bool) + // groupDescAllowBatchImageGeneration is the schema descriptor for allow_batch_image_generation field. + groupDescAllowBatchImageGeneration := groupFields[16].Descriptor() + // group.DefaultAllowBatchImageGeneration holds the default value on creation for the allow_batch_image_generation field. + group.DefaultAllowBatchImageGeneration = groupDescAllowBatchImageGeneration.Default.(bool) // groupDescImageRateIndependent is the schema descriptor for image_rate_independent field. - groupDescImageRateIndependent := groupFields[16].Descriptor() + groupDescImageRateIndependent := groupFields[17].Descriptor() // group.DefaultImageRateIndependent holds the default value on creation for the image_rate_independent field. group.DefaultImageRateIndependent = groupDescImageRateIndependent.Default.(bool) // groupDescImageRateMultiplier is the schema descriptor for image_rate_multiplier field. - groupDescImageRateMultiplier := groupFields[17].Descriptor() + groupDescImageRateMultiplier := groupFields[18].Descriptor() // group.DefaultImageRateMultiplier holds the default value on creation for the image_rate_multiplier field. group.DefaultImageRateMultiplier = groupDescImageRateMultiplier.Default.(float64) + // groupDescBatchImageDiscountMultiplier is the schema descriptor for batch_image_discount_multiplier field. + groupDescBatchImageDiscountMultiplier := groupFields[22].Descriptor() + // group.DefaultBatchImageDiscountMultiplier holds the default value on creation for the batch_image_discount_multiplier field. + group.DefaultBatchImageDiscountMultiplier = groupDescBatchImageDiscountMultiplier.Default.(float64) + // groupDescBatchImageHoldMultiplier is the schema descriptor for batch_image_hold_multiplier field. + groupDescBatchImageHoldMultiplier := groupFields[23].Descriptor() + // group.DefaultBatchImageHoldMultiplier holds the default value on creation for the batch_image_hold_multiplier field. + group.DefaultBatchImageHoldMultiplier = groupDescBatchImageHoldMultiplier.Default.(float64) + // groupDescVideoRateIndependent is the schema descriptor for video_rate_independent field. + groupDescVideoRateIndependent := groupFields[24].Descriptor() + // group.DefaultVideoRateIndependent holds the default value on creation for the video_rate_independent field. + group.DefaultVideoRateIndependent = groupDescVideoRateIndependent.Default.(bool) + // groupDescVideoRateMultiplier is the schema descriptor for video_rate_multiplier field. + groupDescVideoRateMultiplier := groupFields[25].Descriptor() + // group.DefaultVideoRateMultiplier holds the default value on creation for the video_rate_multiplier field. + group.DefaultVideoRateMultiplier = groupDescVideoRateMultiplier.Default.(float64) // groupDescClaudeCodeOnly is the schema descriptor for claude_code_only field. - groupDescClaudeCodeOnly := groupFields[21].Descriptor() + groupDescClaudeCodeOnly := groupFields[29].Descriptor() // group.DefaultClaudeCodeOnly holds the default value on creation for the claude_code_only field. group.DefaultClaudeCodeOnly = groupDescClaudeCodeOnly.Default.(bool) // groupDescModelRoutingEnabled is the schema descriptor for model_routing_enabled field. - groupDescModelRoutingEnabled := groupFields[25].Descriptor() + groupDescModelRoutingEnabled := groupFields[33].Descriptor() // group.DefaultModelRoutingEnabled holds the default value on creation for the model_routing_enabled field. group.DefaultModelRoutingEnabled = groupDescModelRoutingEnabled.Default.(bool) // groupDescMcpXMLInject is the schema descriptor for mcp_xml_inject field. - groupDescMcpXMLInject := groupFields[26].Descriptor() + groupDescMcpXMLInject := groupFields[34].Descriptor() // group.DefaultMcpXMLInject holds the default value on creation for the mcp_xml_inject field. group.DefaultMcpXMLInject = groupDescMcpXMLInject.Default.(bool) // groupDescSupportedModelScopes is the schema descriptor for supported_model_scopes field. - groupDescSupportedModelScopes := groupFields[27].Descriptor() + groupDescSupportedModelScopes := groupFields[35].Descriptor() // group.DefaultSupportedModelScopes holds the default value on creation for the supported_model_scopes field. group.DefaultSupportedModelScopes = groupDescSupportedModelScopes.Default.([]string) // groupDescSortOrder is the schema descriptor for sort_order field. - groupDescSortOrder := groupFields[28].Descriptor() + groupDescSortOrder := groupFields[36].Descriptor() // group.DefaultSortOrder holds the default value on creation for the sort_order field. group.DefaultSortOrder = groupDescSortOrder.Default.(int) // groupDescAllowMessagesDispatch is the schema descriptor for allow_messages_dispatch field. - groupDescAllowMessagesDispatch := groupFields[29].Descriptor() + groupDescAllowMessagesDispatch := groupFields[37].Descriptor() // group.DefaultAllowMessagesDispatch holds the default value on creation for the allow_messages_dispatch field. group.DefaultAllowMessagesDispatch = groupDescAllowMessagesDispatch.Default.(bool) // groupDescRequireOauthOnly is the schema descriptor for require_oauth_only field. - groupDescRequireOauthOnly := groupFields[30].Descriptor() + groupDescRequireOauthOnly := groupFields[38].Descriptor() // group.DefaultRequireOauthOnly holds the default value on creation for the require_oauth_only field. group.DefaultRequireOauthOnly = groupDescRequireOauthOnly.Default.(bool) // groupDescRequirePrivacySet is the schema descriptor for require_privacy_set field. - groupDescRequirePrivacySet := groupFields[31].Descriptor() + groupDescRequirePrivacySet := groupFields[39].Descriptor() // group.DefaultRequirePrivacySet holds the default value on creation for the require_privacy_set field. group.DefaultRequirePrivacySet = groupDescRequirePrivacySet.Default.(bool) // groupDescDefaultMappedModel is the schema descriptor for default_mapped_model field. - groupDescDefaultMappedModel := groupFields[32].Descriptor() + groupDescDefaultMappedModel := groupFields[40].Descriptor() // group.DefaultDefaultMappedModel holds the default value on creation for the default_mapped_model field. group.DefaultDefaultMappedModel = groupDescDefaultMappedModel.Default.(string) // group.DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save. group.DefaultMappedModelValidator = groupDescDefaultMappedModel.Validators[0].(func(string) error) // groupDescMessagesDispatchModelConfig is the schema descriptor for messages_dispatch_model_config field. - groupDescMessagesDispatchModelConfig := groupFields[33].Descriptor() + groupDescMessagesDispatchModelConfig := groupFields[41].Descriptor() // group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field. group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig) // groupDescModelsListConfig is the schema descriptor for models_list_config field. - groupDescModelsListConfig := groupFields[34].Descriptor() + groupDescModelsListConfig := groupFields[42].Descriptor() // group.DefaultModelsListConfig holds the default value on creation for the models_list_config field. group.DefaultModelsListConfig = groupDescModelsListConfig.Default.(domain.GroupModelsListConfig) // groupDescRpmLimit is the schema descriptor for rpm_limit field. - groupDescRpmLimit := groupFields[35].Descriptor() + groupDescRpmLimit := groupFields[43].Descriptor() // group.DefaultRpmLimit holds the default value on creation for the rpm_limit field. group.DefaultRpmLimit = groupDescRpmLimit.Default.(int) idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin() @@ -1787,12 +1976,20 @@ func init() { usagelogDescImageSizeSource := usagelogFields[37].Descriptor() // usagelog.ImageSizeSourceValidator is a validator for the "image_size_source" field. It is called by the builders before save. usagelog.ImageSizeSourceValidator = usagelogDescImageSizeSource.Validators[0].(func(string) error) + // usagelogDescVideoCount is the schema descriptor for video_count field. + usagelogDescVideoCount := usagelogFields[39].Descriptor() + // usagelog.DefaultVideoCount holds the default value on creation for the video_count field. + usagelog.DefaultVideoCount = usagelogDescVideoCount.Default.(int) + // usagelogDescVideoResolution is the schema descriptor for video_resolution field. + usagelogDescVideoResolution := usagelogFields[40].Descriptor() + // usagelog.VideoResolutionValidator is a validator for the "video_resolution" field. It is called by the builders before save. + usagelog.VideoResolutionValidator = usagelogDescVideoResolution.Validators[0].(func(string) error) // usagelogDescCacheTTLOverridden is the schema descriptor for cache_ttl_overridden field. - usagelogDescCacheTTLOverridden := usagelogFields[39].Descriptor() + usagelogDescCacheTTLOverridden := usagelogFields[42].Descriptor() // usagelog.DefaultCacheTTLOverridden holds the default value on creation for the cache_ttl_overridden field. usagelog.DefaultCacheTTLOverridden = usagelogDescCacheTTLOverridden.Default.(bool) // usagelogDescCreatedAt is the schema descriptor for created_at field. - usagelogDescCreatedAt := usagelogFields[40].Descriptor() + usagelogDescCreatedAt := usagelogFields[43].Descriptor() // usagelog.DefaultCreatedAt holds the default value on creation for the created_at field. usagelog.DefaultCreatedAt = usagelogDescCreatedAt.Default.(func() time.Time) userMixin := schema.User{}.Mixin() @@ -1860,54 +2057,58 @@ func init() { userDescBalance := userFields[3].Descriptor() // user.DefaultBalance holds the default value on creation for the balance field. user.DefaultBalance = userDescBalance.Default.(float64) + // userDescFrozenBalance is the schema descriptor for frozen_balance field. + userDescFrozenBalance := userFields[4].Descriptor() + // user.DefaultFrozenBalance holds the default value on creation for the frozen_balance field. + user.DefaultFrozenBalance = userDescFrozenBalance.Default.(float64) // userDescConcurrency is the schema descriptor for concurrency field. - userDescConcurrency := userFields[4].Descriptor() + userDescConcurrency := userFields[5].Descriptor() // user.DefaultConcurrency holds the default value on creation for the concurrency field. user.DefaultConcurrency = userDescConcurrency.Default.(int) // userDescStatus is the schema descriptor for status field. - userDescStatus := userFields[5].Descriptor() + userDescStatus := userFields[6].Descriptor() // user.DefaultStatus holds the default value on creation for the status field. user.DefaultStatus = userDescStatus.Default.(string) // user.StatusValidator is a validator for the "status" field. It is called by the builders before save. user.StatusValidator = userDescStatus.Validators[0].(func(string) error) // userDescUsername is the schema descriptor for username field. - userDescUsername := userFields[6].Descriptor() + userDescUsername := userFields[7].Descriptor() // user.DefaultUsername holds the default value on creation for the username field. user.DefaultUsername = userDescUsername.Default.(string) // user.UsernameValidator is a validator for the "username" field. It is called by the builders before save. user.UsernameValidator = userDescUsername.Validators[0].(func(string) error) // userDescNotes is the schema descriptor for notes field. - userDescNotes := userFields[7].Descriptor() + userDescNotes := userFields[8].Descriptor() // user.DefaultNotes holds the default value on creation for the notes field. user.DefaultNotes = userDescNotes.Default.(string) // userDescTotpEnabled is the schema descriptor for totp_enabled field. - userDescTotpEnabled := userFields[9].Descriptor() + userDescTotpEnabled := userFields[10].Descriptor() // user.DefaultTotpEnabled holds the default value on creation for the totp_enabled field. user.DefaultTotpEnabled = userDescTotpEnabled.Default.(bool) // userDescSignupSource is the schema descriptor for signup_source field. - userDescSignupSource := userFields[11].Descriptor() + userDescSignupSource := userFields[12].Descriptor() // user.DefaultSignupSource holds the default value on creation for the signup_source field. user.DefaultSignupSource = userDescSignupSource.Default.(string) // user.SignupSourceValidator is a validator for the "signup_source" field. It is called by the builders before save. user.SignupSourceValidator = userDescSignupSource.Validators[0].(func(string) error) // userDescBalanceNotifyEnabled is the schema descriptor for balance_notify_enabled field. - userDescBalanceNotifyEnabled := userFields[14].Descriptor() + userDescBalanceNotifyEnabled := userFields[15].Descriptor() // user.DefaultBalanceNotifyEnabled holds the default value on creation for the balance_notify_enabled field. user.DefaultBalanceNotifyEnabled = userDescBalanceNotifyEnabled.Default.(bool) // userDescBalanceNotifyThresholdType is the schema descriptor for balance_notify_threshold_type field. - userDescBalanceNotifyThresholdType := userFields[15].Descriptor() + userDescBalanceNotifyThresholdType := userFields[16].Descriptor() // user.DefaultBalanceNotifyThresholdType holds the default value on creation for the balance_notify_threshold_type field. user.DefaultBalanceNotifyThresholdType = userDescBalanceNotifyThresholdType.Default.(string) // userDescBalanceNotifyExtraEmails is the schema descriptor for balance_notify_extra_emails field. - userDescBalanceNotifyExtraEmails := userFields[17].Descriptor() + userDescBalanceNotifyExtraEmails := userFields[18].Descriptor() // user.DefaultBalanceNotifyExtraEmails holds the default value on creation for the balance_notify_extra_emails field. user.DefaultBalanceNotifyExtraEmails = userDescBalanceNotifyExtraEmails.Default.(string) // userDescTotalRecharged is the schema descriptor for total_recharged field. - userDescTotalRecharged := userFields[18].Descriptor() + userDescTotalRecharged := userFields[19].Descriptor() // user.DefaultTotalRecharged holds the default value on creation for the total_recharged field. user.DefaultTotalRecharged = userDescTotalRecharged.Default.(float64) // userDescRpmLimit is the schema descriptor for rpm_limit field. - userDescRpmLimit := userFields[19].Descriptor() + userDescRpmLimit := userFields[20].Descriptor() // user.DefaultRpmLimit holds the default value on creation for the rpm_limit field. user.DefaultRpmLimit = userDescRpmLimit.Default.(int) userallowedgroupFields := schema.UserAllowedGroup{}.Fields() diff --git a/backend/ent/schema/batch_image_event.go b/backend/ent/schema/batch_image_event.go new file mode 100644 index 0000000000..44af3d1620 --- /dev/null +++ b/backend/ent/schema/batch_image_event.go @@ -0,0 +1,43 @@ +package schema + +import ( + "time" + + "entgo.io/ent" + "entgo.io/ent/dialect" + "entgo.io/ent/dialect/entsql" + "entgo.io/ent/schema" + "entgo.io/ent/schema/field" + "entgo.io/ent/schema/index" +) + +// BatchImageEvent records append-only operational events for batch image jobs. +type BatchImageEvent struct { + ent.Schema +} + +func (BatchImageEvent) Annotations() []schema.Annotation { + return []schema.Annotation{ + entsql.Annotation{Table: "batch_image_events"}, + } +} + +func (BatchImageEvent) Fields() []ent.Field { + return []ent.Field{ + field.String("job_id").MaxLen(64), + field.String("event_type").MaxLen(64), + field.JSON("payload", map[string]any{}). + Optional(). + SchemaType(map[string]string{dialect.Postgres: "jsonb"}), + field.String("event_hash").Optional().Nillable().MaxLen(128), + field.Time("created_at").Immutable().Default(time.Now).SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + } +} + +func (BatchImageEvent) Indexes() []ent.Index { + return []ent.Index{ + index.Fields("job_id", "created_at"), + index.Fields("event_type"), + index.Fields("job_id", "event_hash").Unique().Annotations(entsql.IndexWhere("event_hash IS NOT NULL AND event_hash <> ''")), + } +} diff --git a/backend/ent/schema/batch_image_item.go b/backend/ent/schema/batch_image_item.go new file mode 100644 index 0000000000..6a7a097c6a --- /dev/null +++ b/backend/ent/schema/batch_image_item.go @@ -0,0 +1,53 @@ +package schema + +import ( + "time" + + "entgo.io/ent" + "entgo.io/ent/dialect" + "entgo.io/ent/dialect/entsql" + "entgo.io/ent/schema" + "entgo.io/ent/schema/field" + "entgo.io/ent/schema/index" +) + +// BatchImageItem holds indexed output rows for a batch image job. +type BatchImageItem struct { + ent.Schema +} + +func (BatchImageItem) Annotations() []schema.Annotation { + return []schema.Annotation{ + entsql.Annotation{Table: "batch_image_items"}, + } +} + +func (BatchImageItem) Fields() []ent.Field { + return []ent.Field{ + field.String("job_id").MaxLen(64), + field.String("custom_id").MaxLen(255), + field.String("status").MaxLen(32), + field.String("request_hash").Optional().Nillable().MaxLen(128), + field.String("prompt_preview").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "text"}), + field.String("provider_source_object").Optional().Nillable().MaxLen(1024), + field.Int("source_line_number").Optional().Nillable(), + field.Int64("source_byte_offset").Optional().Nillable(), + field.Int64("source_byte_length").Optional().Nillable(), + field.String("mime_type").Optional().Nillable().MaxLen(128), + field.String("file_extension").Optional().Nillable().MaxLen(32), + field.Int("image_count").Default(0), + field.String("error_code").Optional().Nillable().MaxLen(128), + field.String("error_message").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "text"}), + field.Float("billed_amount").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "decimal(20,10)"}), + field.Time("created_at").Immutable().Default(time.Now).SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("indexed_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + } +} + +func (BatchImageItem) Indexes() []ent.Index { + return []ent.Index{ + index.Fields("job_id", "custom_id").Unique(), + index.Fields("job_id", "status"), + index.Fields("provider_source_object"), + } +} diff --git a/backend/ent/schema/batch_image_job.go b/backend/ent/schema/batch_image_job.go new file mode 100644 index 0000000000..a65156eaea --- /dev/null +++ b/backend/ent/schema/batch_image_job.go @@ -0,0 +1,86 @@ +package schema + +import ( + "time" + + "entgo.io/ent" + "entgo.io/ent/dialect" + "entgo.io/ent/dialect/entsql" + "entgo.io/ent/schema" + "entgo.io/ent/schema/field" + "entgo.io/ent/schema/index" +) + +// BatchImageJob holds the schema definition for asynchronous image batch jobs. +// +// 删除策略:账务源保留 +// 这张表是批量生图任务的账务和状态源;用户侧删除仅通过 user_deleted_at +// 从列表隐藏,输出清理通过 output_deleted 状态和删除时间字段表达。 +type BatchImageJob struct { + ent.Schema +} + +func (BatchImageJob) Annotations() []schema.Annotation { + return []schema.Annotation{ + entsql.Annotation{Table: "batch_image_jobs"}, + } +} + +func (BatchImageJob) Fields() []ent.Field { + return []ent.Field{ + field.String("batch_id").MaxLen(64).Immutable(), + field.Int64("user_id"), + field.Int64("api_key_id").Optional().Nillable(), + field.Int64("account_id").Optional().Nillable(), + field.String("provider").MaxLen(32), + field.String("model").MaxLen(128), + field.String("task_name").MaxLen(255).Default(""), + field.String("status").MaxLen(32).Default("created"), + field.String("provider_job_name").Optional().Nillable().MaxLen(512), + field.String("provider_input_ref").Optional().Nillable().MaxLen(1024), + field.String("provider_output_ref").Optional().Nillable().MaxLen(1024), + field.String("gcs_input_uri").Optional().Nillable().MaxLen(1024), + field.String("gcs_output_uri").Optional().Nillable().MaxLen(1024), + field.Int("item_count"), + field.Int("success_count").Default(0), + field.Int("fail_count").Default(0), + field.Int("cancelled_count").Default(0), + field.Float("estimated_cost").SchemaType(map[string]string{dialect.Postgres: "decimal(20,10)"}).Default(0), + field.Float("hold_amount").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "decimal(20,10)"}), + field.Float("actual_cost").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "decimal(20,10)"}), + field.String("currency").MaxLen(16).Default("USD"), + field.String("hold_id").Optional().Nillable().MaxLen(128), + field.String("idempotency_key").Optional().Nillable().MaxLen(255), + field.String("request_hash").Optional().Nillable().MaxLen(128), + field.String("manifest_hash").Optional().Nillable().MaxLen(128), + field.Int("retry_count").Default(0), + field.Int("version").Default(0), + field.Time("output_expires_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("input_deleted_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("output_deleted_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("downloaded_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("user_deleted_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.String("last_error_code").Optional().Nillable().MaxLen(128), + field.String("last_error_message").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "text"}), + field.Time("created_at").Immutable().Default(time.Now).SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("updated_at").Default(time.Now).UpdateDefault(time.Now).SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("submitted_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("started_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("finished_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + field.Time("settled_at").Optional().Nillable().SchemaType(map[string]string{dialect.Postgres: "timestamptz"}), + } +} + +func (BatchImageJob) Indexes() []ent.Index { + return []ent.Index{ + index.Fields("batch_id").Unique(), + index.Fields("user_id", "created_at"), + index.Fields("status"), + index.Fields("provider", "status"), + index.Fields("idempotency_key").Annotations(entsql.IndexWhere("idempotency_key IS NOT NULL AND idempotency_key <> ''")), + index.Fields("manifest_hash").Unique().Annotations(entsql.IndexWhere("manifest_hash IS NOT NULL AND manifest_hash <> ''")), + index.Fields("output_expires_at"), + index.Fields("downloaded_at"), + index.Fields("user_deleted_at"), + } +} diff --git a/backend/ent/schema/group.go b/backend/ent/schema/group.go index 2b8420db6d..b104609a1b 100644 --- a/backend/ent/schema/group.go +++ b/backend/ent/schema/group.go @@ -93,6 +93,9 @@ func (Group) Fields() []ent.Field { field.Bool("allow_image_generation"). Default(false). Comment("是否允许该分组使用图片生成能力"), + field.Bool("allow_batch_image_generation"). + Default(false). + Comment("是否允许该分组使用批量图片生成能力"), field.Bool("image_rate_independent"). Default(false). Comment("图片生成是否使用独立倍率;false 表示共享分组有效倍率"), @@ -112,6 +115,33 @@ func (Group) Fields() []ent.Field { Optional(). Nillable(). SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}), + field.Float("batch_image_discount_multiplier"). + SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}). + Default(0.5). + Comment("批量图片生成折扣倍率,最终单价会乘以该值;0 表示免费"), + field.Float("batch_image_hold_multiplier"). + SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}). + Default(0.6). + Comment("批量图片生成冻结价格比例,按普通生图原价乘以该比例冻结,结算后释放差额"), + field.Bool("video_rate_independent"). + Default(false). + Comment("视频生成是否使用独立倍率;false 表示共享分组有效倍率"), + field.Float("video_rate_multiplier"). + SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}). + Default(1.0). + Comment("视频生成独立倍率,仅 video_rate_independent=true 时生效"), + field.Float("video_price_480p"). + Optional(). + Nillable(). + SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}), + field.Float("video_price_720p"). + Optional(). + Nillable(). + SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}), + field.Float("video_price_1080p"). + Optional(). + Nillable(). + SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}), // Claude Code 客户端限制 (added by migration 029) field.Bool("claude_code_only"). diff --git a/backend/ent/schema/usage_log.go b/backend/ent/schema/usage_log.go index db9e517892..e84cc1c140 100644 --- a/backend/ent/schema/usage_log.go +++ b/backend/ent/schema/usage_log.go @@ -149,6 +149,20 @@ func (UsageLog) Fields() []ent.Field { field.JSON("image_size_breakdown", map[string]int{}). Optional(). SchemaType(map[string]string{dialect.Postgres: "jsonb"}), + + // 视频生成字段(Grok 视频按秒计费;billing_mode 走 token/其他模式时这些列仍标记视频用量) + field.Int("video_count"). + Default(0). + Comment("视频生成数量;>0 表示本行是视频生成用量"), + field.String("video_resolution"). + MaxLen(10). + Optional(). + Nillable(). + Comment("计费用视频分辨率 480p/720p/1080p"), + field.Int("video_duration_seconds"). + Optional(). + Nillable(). + Comment("提交时请求的视频时长(秒),按秒计费的乘数"), // Cache TTL Override 标记(管理员强制替换了缓存 TTL 计费) field.Bool("cache_ttl_overridden"). Default(false), diff --git a/backend/ent/schema/user.go b/backend/ent/schema/user.go index 127b5af9a7..baa7efbbd9 100644 --- a/backend/ent/schema/user.go +++ b/backend/ent/schema/user.go @@ -49,6 +49,9 @@ func (User) Fields() []ent.Field { field.Float("balance"). SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}). Default(0), + field.Float("frozen_balance"). + SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}). + Default(0), field.Int("concurrency"). Default(5), field.String("status"). diff --git a/backend/ent/tx.go b/backend/ent/tx.go index 846cfcd4da..6de2c2b63b 100644 --- a/backend/ent/tx.go +++ b/backend/ent/tx.go @@ -28,6 +28,12 @@ type Tx struct { AuthIdentity *AuthIdentityClient // AuthIdentityChannel is the client for interacting with the AuthIdentityChannel builders. AuthIdentityChannel *AuthIdentityChannelClient + // BatchImageEvent is the client for interacting with the BatchImageEvent builders. + BatchImageEvent *BatchImageEventClient + // BatchImageItem is the client for interacting with the BatchImageItem builders. + BatchImageItem *BatchImageItemClient + // BatchImageJob is the client for interacting with the BatchImageJob builders. + BatchImageJob *BatchImageJobClient // ChannelMonitor is the client for interacting with the ChannelMonitor builders. ChannelMonitor *ChannelMonitorClient // ChannelMonitorDailyRollup is the client for interacting with the ChannelMonitorDailyRollup builders. @@ -222,6 +228,9 @@ func (tx *Tx) init() { tx.AnnouncementRead = NewAnnouncementReadClient(tx.config) tx.AuthIdentity = NewAuthIdentityClient(tx.config) tx.AuthIdentityChannel = NewAuthIdentityChannelClient(tx.config) + tx.BatchImageEvent = NewBatchImageEventClient(tx.config) + tx.BatchImageItem = NewBatchImageItemClient(tx.config) + tx.BatchImageJob = NewBatchImageJobClient(tx.config) tx.ChannelMonitor = NewChannelMonitorClient(tx.config) tx.ChannelMonitorDailyRollup = NewChannelMonitorDailyRollupClient(tx.config) tx.ChannelMonitorHistory = NewChannelMonitorHistoryClient(tx.config) diff --git a/backend/ent/usagelog.go b/backend/ent/usagelog.go index 283fe828a9..4d374a8495 100644 --- a/backend/ent/usagelog.go +++ b/backend/ent/usagelog.go @@ -101,6 +101,12 @@ type UsageLog struct { ImageSizeSource *string `json:"image_size_source,omitempty"` // ImageSizeBreakdown holds the value of the "image_size_breakdown" field. ImageSizeBreakdown map[string]int `json:"image_size_breakdown,omitempty"` + // 视频生成数量;>0 表示本行是视频生成用量 + VideoCount int `json:"video_count,omitempty"` + // 计费用视频分辨率 480p/720p/1080p + VideoResolution *string `json:"video_resolution,omitempty"` + // 提交时请求的视频时长(秒),按秒计费的乘数 + VideoDurationSeconds *int `json:"video_duration_seconds,omitempty"` // CacheTTLOverridden holds the value of the "cache_ttl_overridden" field. CacheTTLOverridden bool `json:"cache_ttl_overridden,omitempty"` // CreatedAt holds the value of the "created_at" field. @@ -194,9 +200,9 @@ func (*UsageLog) scanValues(columns []string) ([]any, error) { values[i] = new(sql.NullBool) case usagelog.FieldInputCost, usagelog.FieldOutputCost, usagelog.FieldCacheCreationCost, usagelog.FieldCacheReadCost, usagelog.FieldTotalCost, usagelog.FieldActualCost, usagelog.FieldRateMultiplier, usagelog.FieldAccountRateMultiplier: values[i] = new(sql.NullFloat64) - case usagelog.FieldID, usagelog.FieldUserID, usagelog.FieldAPIKeyID, usagelog.FieldAccountID, usagelog.FieldChannelID, usagelog.FieldGroupID, usagelog.FieldSubscriptionID, usagelog.FieldInputTokens, usagelog.FieldOutputTokens, usagelog.FieldCacheCreationTokens, usagelog.FieldCacheReadTokens, usagelog.FieldCacheCreation5mTokens, usagelog.FieldCacheCreation1hTokens, usagelog.FieldBillingType, usagelog.FieldDurationMs, usagelog.FieldFirstTokenMs, usagelog.FieldImageCount: + case usagelog.FieldID, usagelog.FieldUserID, usagelog.FieldAPIKeyID, usagelog.FieldAccountID, usagelog.FieldChannelID, usagelog.FieldGroupID, usagelog.FieldSubscriptionID, usagelog.FieldInputTokens, usagelog.FieldOutputTokens, usagelog.FieldCacheCreationTokens, usagelog.FieldCacheReadTokens, usagelog.FieldCacheCreation5mTokens, usagelog.FieldCacheCreation1hTokens, usagelog.FieldBillingType, usagelog.FieldDurationMs, usagelog.FieldFirstTokenMs, usagelog.FieldImageCount, usagelog.FieldVideoCount, usagelog.FieldVideoDurationSeconds: values[i] = new(sql.NullInt64) - case usagelog.FieldRequestID, usagelog.FieldModel, usagelog.FieldRequestedModel, usagelog.FieldUpstreamModel, usagelog.FieldModelMappingChain, usagelog.FieldBillingTier, usagelog.FieldBillingMode, usagelog.FieldUserAgent, usagelog.FieldIPAddress, usagelog.FieldImageSize, usagelog.FieldImageInputSize, usagelog.FieldImageOutputSize, usagelog.FieldImageSizeSource: + case usagelog.FieldRequestID, usagelog.FieldModel, usagelog.FieldRequestedModel, usagelog.FieldUpstreamModel, usagelog.FieldModelMappingChain, usagelog.FieldBillingTier, usagelog.FieldBillingMode, usagelog.FieldUserAgent, usagelog.FieldIPAddress, usagelog.FieldImageSize, usagelog.FieldImageInputSize, usagelog.FieldImageOutputSize, usagelog.FieldImageSizeSource, usagelog.FieldVideoResolution: values[i] = new(sql.NullString) case usagelog.FieldCreatedAt: values[i] = new(sql.NullTime) @@ -474,6 +480,26 @@ func (_m *UsageLog) assignValues(columns []string, values []any) error { return fmt.Errorf("unmarshal field image_size_breakdown: %w", err) } } + case usagelog.FieldVideoCount: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field video_count", values[i]) + } else if value.Valid { + _m.VideoCount = int(value.Int64) + } + case usagelog.FieldVideoResolution: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field video_resolution", values[i]) + } else if value.Valid { + _m.VideoResolution = new(string) + *_m.VideoResolution = value.String + } + case usagelog.FieldVideoDurationSeconds: + if value, ok := values[i].(*sql.NullInt64); !ok { + return fmt.Errorf("unexpected type %T for field video_duration_seconds", values[i]) + } else if value.Valid { + _m.VideoDurationSeconds = new(int) + *_m.VideoDurationSeconds = int(value.Int64) + } case usagelog.FieldCacheTTLOverridden: if value, ok := values[i].(*sql.NullBool); !ok { return fmt.Errorf("unexpected type %T for field cache_ttl_overridden", values[i]) @@ -698,6 +724,19 @@ func (_m *UsageLog) String() string { builder.WriteString("image_size_breakdown=") builder.WriteString(fmt.Sprintf("%v", _m.ImageSizeBreakdown)) builder.WriteString(", ") + builder.WriteString("video_count=") + builder.WriteString(fmt.Sprintf("%v", _m.VideoCount)) + builder.WriteString(", ") + if v := _m.VideoResolution; v != nil { + builder.WriteString("video_resolution=") + builder.WriteString(*v) + } + builder.WriteString(", ") + if v := _m.VideoDurationSeconds; v != nil { + builder.WriteString("video_duration_seconds=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") builder.WriteString("cache_ttl_overridden=") builder.WriteString(fmt.Sprintf("%v", _m.CacheTTLOverridden)) builder.WriteString(", ") diff --git a/backend/ent/usagelog/usagelog.go b/backend/ent/usagelog/usagelog.go index 297e0b41ad..a74a92c40f 100644 --- a/backend/ent/usagelog/usagelog.go +++ b/backend/ent/usagelog/usagelog.go @@ -92,6 +92,12 @@ const ( FieldImageSizeSource = "image_size_source" // FieldImageSizeBreakdown holds the string denoting the image_size_breakdown field in the database. FieldImageSizeBreakdown = "image_size_breakdown" + // FieldVideoCount holds the string denoting the video_count field in the database. + FieldVideoCount = "video_count" + // FieldVideoResolution holds the string denoting the video_resolution field in the database. + FieldVideoResolution = "video_resolution" + // FieldVideoDurationSeconds holds the string denoting the video_duration_seconds field in the database. + FieldVideoDurationSeconds = "video_duration_seconds" // FieldCacheTTLOverridden holds the string denoting the cache_ttl_overridden field in the database. FieldCacheTTLOverridden = "cache_ttl_overridden" // FieldCreatedAt holds the string denoting the created_at field in the database. @@ -187,6 +193,9 @@ var Columns = []string{ FieldImageOutputSize, FieldImageSizeSource, FieldImageSizeBreakdown, + FieldVideoCount, + FieldVideoResolution, + FieldVideoDurationSeconds, FieldCacheTTLOverridden, FieldCreatedAt, } @@ -260,6 +269,10 @@ var ( ImageOutputSizeValidator func(string) error // ImageSizeSourceValidator is a validator for the "image_size_source" field. It is called by the builders before save. ImageSizeSourceValidator func(string) error + // DefaultVideoCount holds the default value on creation for the "video_count" field. + DefaultVideoCount int + // VideoResolutionValidator is a validator for the "video_resolution" field. It is called by the builders before save. + VideoResolutionValidator func(string) error // DefaultCacheTTLOverridden holds the default value on creation for the "cache_ttl_overridden" field. DefaultCacheTTLOverridden bool // DefaultCreatedAt holds the default value on creation for the "created_at" field. @@ -464,6 +477,21 @@ func ByImageSizeSource(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldImageSizeSource, opts...).ToFunc() } +// ByVideoCount orders the results by the video_count field. +func ByVideoCount(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoCount, opts...).ToFunc() +} + +// ByVideoResolution orders the results by the video_resolution field. +func ByVideoResolution(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoResolution, opts...).ToFunc() +} + +// ByVideoDurationSeconds orders the results by the video_duration_seconds field. +func ByVideoDurationSeconds(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldVideoDurationSeconds, opts...).ToFunc() +} + // ByCacheTTLOverridden orders the results by the cache_ttl_overridden field. func ByCacheTTLOverridden(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldCacheTTLOverridden, opts...).ToFunc() diff --git a/backend/ent/usagelog/where.go b/backend/ent/usagelog/where.go index 2987f17930..4b08cc3425 100644 --- a/backend/ent/usagelog/where.go +++ b/backend/ent/usagelog/where.go @@ -245,6 +245,21 @@ func ImageSizeSource(v string) predicate.UsageLog { return predicate.UsageLog(sql.FieldEQ(FieldImageSizeSource, v)) } +// VideoCount applies equality check predicate on the "video_count" field. It's identical to VideoCountEQ. +func VideoCount(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEQ(FieldVideoCount, v)) +} + +// VideoResolution applies equality check predicate on the "video_resolution" field. It's identical to VideoResolutionEQ. +func VideoResolution(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEQ(FieldVideoResolution, v)) +} + +// VideoDurationSeconds applies equality check predicate on the "video_duration_seconds" field. It's identical to VideoDurationSecondsEQ. +func VideoDurationSeconds(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEQ(FieldVideoDurationSeconds, v)) +} + // CacheTTLOverridden applies equality check predicate on the "cache_ttl_overridden" field. It's identical to CacheTTLOverriddenEQ. func CacheTTLOverridden(v bool) predicate.UsageLog { return predicate.UsageLog(sql.FieldEQ(FieldCacheTTLOverridden, v)) @@ -2150,6 +2165,171 @@ func ImageSizeBreakdownNotNil() predicate.UsageLog { return predicate.UsageLog(sql.FieldNotNull(FieldImageSizeBreakdown)) } +// VideoCountEQ applies the EQ predicate on the "video_count" field. +func VideoCountEQ(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEQ(FieldVideoCount, v)) +} + +// VideoCountNEQ applies the NEQ predicate on the "video_count" field. +func VideoCountNEQ(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldNEQ(FieldVideoCount, v)) +} + +// VideoCountIn applies the In predicate on the "video_count" field. +func VideoCountIn(vs ...int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldIn(FieldVideoCount, vs...)) +} + +// VideoCountNotIn applies the NotIn predicate on the "video_count" field. +func VideoCountNotIn(vs ...int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldNotIn(FieldVideoCount, vs...)) +} + +// VideoCountGT applies the GT predicate on the "video_count" field. +func VideoCountGT(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldGT(FieldVideoCount, v)) +} + +// VideoCountGTE applies the GTE predicate on the "video_count" field. +func VideoCountGTE(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldGTE(FieldVideoCount, v)) +} + +// VideoCountLT applies the LT predicate on the "video_count" field. +func VideoCountLT(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldLT(FieldVideoCount, v)) +} + +// VideoCountLTE applies the LTE predicate on the "video_count" field. +func VideoCountLTE(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldLTE(FieldVideoCount, v)) +} + +// VideoResolutionEQ applies the EQ predicate on the "video_resolution" field. +func VideoResolutionEQ(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEQ(FieldVideoResolution, v)) +} + +// VideoResolutionNEQ applies the NEQ predicate on the "video_resolution" field. +func VideoResolutionNEQ(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldNEQ(FieldVideoResolution, v)) +} + +// VideoResolutionIn applies the In predicate on the "video_resolution" field. +func VideoResolutionIn(vs ...string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldIn(FieldVideoResolution, vs...)) +} + +// VideoResolutionNotIn applies the NotIn predicate on the "video_resolution" field. +func VideoResolutionNotIn(vs ...string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldNotIn(FieldVideoResolution, vs...)) +} + +// VideoResolutionGT applies the GT predicate on the "video_resolution" field. +func VideoResolutionGT(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldGT(FieldVideoResolution, v)) +} + +// VideoResolutionGTE applies the GTE predicate on the "video_resolution" field. +func VideoResolutionGTE(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldGTE(FieldVideoResolution, v)) +} + +// VideoResolutionLT applies the LT predicate on the "video_resolution" field. +func VideoResolutionLT(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldLT(FieldVideoResolution, v)) +} + +// VideoResolutionLTE applies the LTE predicate on the "video_resolution" field. +func VideoResolutionLTE(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldLTE(FieldVideoResolution, v)) +} + +// VideoResolutionContains applies the Contains predicate on the "video_resolution" field. +func VideoResolutionContains(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldContains(FieldVideoResolution, v)) +} + +// VideoResolutionHasPrefix applies the HasPrefix predicate on the "video_resolution" field. +func VideoResolutionHasPrefix(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldHasPrefix(FieldVideoResolution, v)) +} + +// VideoResolutionHasSuffix applies the HasSuffix predicate on the "video_resolution" field. +func VideoResolutionHasSuffix(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldHasSuffix(FieldVideoResolution, v)) +} + +// VideoResolutionIsNil applies the IsNil predicate on the "video_resolution" field. +func VideoResolutionIsNil() predicate.UsageLog { + return predicate.UsageLog(sql.FieldIsNull(FieldVideoResolution)) +} + +// VideoResolutionNotNil applies the NotNil predicate on the "video_resolution" field. +func VideoResolutionNotNil() predicate.UsageLog { + return predicate.UsageLog(sql.FieldNotNull(FieldVideoResolution)) +} + +// VideoResolutionEqualFold applies the EqualFold predicate on the "video_resolution" field. +func VideoResolutionEqualFold(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEqualFold(FieldVideoResolution, v)) +} + +// VideoResolutionContainsFold applies the ContainsFold predicate on the "video_resolution" field. +func VideoResolutionContainsFold(v string) predicate.UsageLog { + return predicate.UsageLog(sql.FieldContainsFold(FieldVideoResolution, v)) +} + +// VideoDurationSecondsEQ applies the EQ predicate on the "video_duration_seconds" field. +func VideoDurationSecondsEQ(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldEQ(FieldVideoDurationSeconds, v)) +} + +// VideoDurationSecondsNEQ applies the NEQ predicate on the "video_duration_seconds" field. +func VideoDurationSecondsNEQ(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldNEQ(FieldVideoDurationSeconds, v)) +} + +// VideoDurationSecondsIn applies the In predicate on the "video_duration_seconds" field. +func VideoDurationSecondsIn(vs ...int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldIn(FieldVideoDurationSeconds, vs...)) +} + +// VideoDurationSecondsNotIn applies the NotIn predicate on the "video_duration_seconds" field. +func VideoDurationSecondsNotIn(vs ...int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldNotIn(FieldVideoDurationSeconds, vs...)) +} + +// VideoDurationSecondsGT applies the GT predicate on the "video_duration_seconds" field. +func VideoDurationSecondsGT(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldGT(FieldVideoDurationSeconds, v)) +} + +// VideoDurationSecondsGTE applies the GTE predicate on the "video_duration_seconds" field. +func VideoDurationSecondsGTE(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldGTE(FieldVideoDurationSeconds, v)) +} + +// VideoDurationSecondsLT applies the LT predicate on the "video_duration_seconds" field. +func VideoDurationSecondsLT(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldLT(FieldVideoDurationSeconds, v)) +} + +// VideoDurationSecondsLTE applies the LTE predicate on the "video_duration_seconds" field. +func VideoDurationSecondsLTE(v int) predicate.UsageLog { + return predicate.UsageLog(sql.FieldLTE(FieldVideoDurationSeconds, v)) +} + +// VideoDurationSecondsIsNil applies the IsNil predicate on the "video_duration_seconds" field. +func VideoDurationSecondsIsNil() predicate.UsageLog { + return predicate.UsageLog(sql.FieldIsNull(FieldVideoDurationSeconds)) +} + +// VideoDurationSecondsNotNil applies the NotNil predicate on the "video_duration_seconds" field. +func VideoDurationSecondsNotNil() predicate.UsageLog { + return predicate.UsageLog(sql.FieldNotNull(FieldVideoDurationSeconds)) +} + // CacheTTLOverriddenEQ applies the EQ predicate on the "cache_ttl_overridden" field. func CacheTTLOverriddenEQ(v bool) predicate.UsageLog { return predicate.UsageLog(sql.FieldEQ(FieldCacheTTLOverridden, v)) diff --git a/backend/ent/usagelog_create.go b/backend/ent/usagelog_create.go index 17e800f9ca..3326f72fc0 100644 --- a/backend/ent/usagelog_create.go +++ b/backend/ent/usagelog_create.go @@ -525,6 +525,48 @@ func (_c *UsageLogCreate) SetImageSizeBreakdown(v map[string]int) *UsageLogCreat return _c } +// SetVideoCount sets the "video_count" field. +func (_c *UsageLogCreate) SetVideoCount(v int) *UsageLogCreate { + _c.mutation.SetVideoCount(v) + return _c +} + +// SetNillableVideoCount sets the "video_count" field if the given value is not nil. +func (_c *UsageLogCreate) SetNillableVideoCount(v *int) *UsageLogCreate { + if v != nil { + _c.SetVideoCount(*v) + } + return _c +} + +// SetVideoResolution sets the "video_resolution" field. +func (_c *UsageLogCreate) SetVideoResolution(v string) *UsageLogCreate { + _c.mutation.SetVideoResolution(v) + return _c +} + +// SetNillableVideoResolution sets the "video_resolution" field if the given value is not nil. +func (_c *UsageLogCreate) SetNillableVideoResolution(v *string) *UsageLogCreate { + if v != nil { + _c.SetVideoResolution(*v) + } + return _c +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (_c *UsageLogCreate) SetVideoDurationSeconds(v int) *UsageLogCreate { + _c.mutation.SetVideoDurationSeconds(v) + return _c +} + +// SetNillableVideoDurationSeconds sets the "video_duration_seconds" field if the given value is not nil. +func (_c *UsageLogCreate) SetNillableVideoDurationSeconds(v *int) *UsageLogCreate { + if v != nil { + _c.SetVideoDurationSeconds(*v) + } + return _c +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (_c *UsageLogCreate) SetCacheTTLOverridden(v bool) *UsageLogCreate { _c.mutation.SetCacheTTLOverridden(v) @@ -677,6 +719,10 @@ func (_c *UsageLogCreate) defaults() { v := usagelog.DefaultImageCount _c.mutation.SetImageCount(v) } + if _, ok := _c.mutation.VideoCount(); !ok { + v := usagelog.DefaultVideoCount + _c.mutation.SetVideoCount(v) + } if _, ok := _c.mutation.CacheTTLOverridden(); !ok { v := usagelog.DefaultCacheTTLOverridden _c.mutation.SetCacheTTLOverridden(v) @@ -817,6 +863,14 @@ func (_c *UsageLogCreate) check() error { return &ValidationError{Name: "image_size_source", err: fmt.Errorf(`ent: validator failed for field "UsageLog.image_size_source": %w`, err)} } } + if _, ok := _c.mutation.VideoCount(); !ok { + return &ValidationError{Name: "video_count", err: errors.New(`ent: missing required field "UsageLog.video_count"`)} + } + if v, ok := _c.mutation.VideoResolution(); ok { + if err := usagelog.VideoResolutionValidator(v); err != nil { + return &ValidationError{Name: "video_resolution", err: fmt.Errorf(`ent: validator failed for field "UsageLog.video_resolution": %w`, err)} + } + } if _, ok := _c.mutation.CacheTTLOverridden(); !ok { return &ValidationError{Name: "cache_ttl_overridden", err: errors.New(`ent: missing required field "UsageLog.cache_ttl_overridden"`)} } @@ -995,6 +1049,18 @@ func (_c *UsageLogCreate) createSpec() (*UsageLog, *sqlgraph.CreateSpec) { _spec.SetField(usagelog.FieldImageSizeBreakdown, field.TypeJSON, value) _node.ImageSizeBreakdown = value } + if value, ok := _c.mutation.VideoCount(); ok { + _spec.SetField(usagelog.FieldVideoCount, field.TypeInt, value) + _node.VideoCount = value + } + if value, ok := _c.mutation.VideoResolution(); ok { + _spec.SetField(usagelog.FieldVideoResolution, field.TypeString, value) + _node.VideoResolution = &value + } + if value, ok := _c.mutation.VideoDurationSeconds(); ok { + _spec.SetField(usagelog.FieldVideoDurationSeconds, field.TypeInt, value) + _node.VideoDurationSeconds = &value + } if value, ok := _c.mutation.CacheTTLOverridden(); ok { _spec.SetField(usagelog.FieldCacheTTLOverridden, field.TypeBool, value) _node.CacheTTLOverridden = value @@ -1830,6 +1896,66 @@ func (u *UsageLogUpsert) ClearImageSizeBreakdown() *UsageLogUpsert { return u } +// SetVideoCount sets the "video_count" field. +func (u *UsageLogUpsert) SetVideoCount(v int) *UsageLogUpsert { + u.Set(usagelog.FieldVideoCount, v) + return u +} + +// UpdateVideoCount sets the "video_count" field to the value that was provided on create. +func (u *UsageLogUpsert) UpdateVideoCount() *UsageLogUpsert { + u.SetExcluded(usagelog.FieldVideoCount) + return u +} + +// AddVideoCount adds v to the "video_count" field. +func (u *UsageLogUpsert) AddVideoCount(v int) *UsageLogUpsert { + u.Add(usagelog.FieldVideoCount, v) + return u +} + +// SetVideoResolution sets the "video_resolution" field. +func (u *UsageLogUpsert) SetVideoResolution(v string) *UsageLogUpsert { + u.Set(usagelog.FieldVideoResolution, v) + return u +} + +// UpdateVideoResolution sets the "video_resolution" field to the value that was provided on create. +func (u *UsageLogUpsert) UpdateVideoResolution() *UsageLogUpsert { + u.SetExcluded(usagelog.FieldVideoResolution) + return u +} + +// ClearVideoResolution clears the value of the "video_resolution" field. +func (u *UsageLogUpsert) ClearVideoResolution() *UsageLogUpsert { + u.SetNull(usagelog.FieldVideoResolution) + return u +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (u *UsageLogUpsert) SetVideoDurationSeconds(v int) *UsageLogUpsert { + u.Set(usagelog.FieldVideoDurationSeconds, v) + return u +} + +// UpdateVideoDurationSeconds sets the "video_duration_seconds" field to the value that was provided on create. +func (u *UsageLogUpsert) UpdateVideoDurationSeconds() *UsageLogUpsert { + u.SetExcluded(usagelog.FieldVideoDurationSeconds) + return u +} + +// AddVideoDurationSeconds adds v to the "video_duration_seconds" field. +func (u *UsageLogUpsert) AddVideoDurationSeconds(v int) *UsageLogUpsert { + u.Add(usagelog.FieldVideoDurationSeconds, v) + return u +} + +// ClearVideoDurationSeconds clears the value of the "video_duration_seconds" field. +func (u *UsageLogUpsert) ClearVideoDurationSeconds() *UsageLogUpsert { + u.SetNull(usagelog.FieldVideoDurationSeconds) + return u +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (u *UsageLogUpsert) SetCacheTTLOverridden(v bool) *UsageLogUpsert { u.Set(usagelog.FieldCacheTTLOverridden, v) @@ -2692,6 +2818,76 @@ func (u *UsageLogUpsertOne) ClearImageSizeBreakdown() *UsageLogUpsertOne { }) } +// SetVideoCount sets the "video_count" field. +func (u *UsageLogUpsertOne) SetVideoCount(v int) *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.SetVideoCount(v) + }) +} + +// AddVideoCount adds v to the "video_count" field. +func (u *UsageLogUpsertOne) AddVideoCount(v int) *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.AddVideoCount(v) + }) +} + +// UpdateVideoCount sets the "video_count" field to the value that was provided on create. +func (u *UsageLogUpsertOne) UpdateVideoCount() *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.UpdateVideoCount() + }) +} + +// SetVideoResolution sets the "video_resolution" field. +func (u *UsageLogUpsertOne) SetVideoResolution(v string) *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.SetVideoResolution(v) + }) +} + +// UpdateVideoResolution sets the "video_resolution" field to the value that was provided on create. +func (u *UsageLogUpsertOne) UpdateVideoResolution() *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.UpdateVideoResolution() + }) +} + +// ClearVideoResolution clears the value of the "video_resolution" field. +func (u *UsageLogUpsertOne) ClearVideoResolution() *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.ClearVideoResolution() + }) +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (u *UsageLogUpsertOne) SetVideoDurationSeconds(v int) *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.SetVideoDurationSeconds(v) + }) +} + +// AddVideoDurationSeconds adds v to the "video_duration_seconds" field. +func (u *UsageLogUpsertOne) AddVideoDurationSeconds(v int) *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.AddVideoDurationSeconds(v) + }) +} + +// UpdateVideoDurationSeconds sets the "video_duration_seconds" field to the value that was provided on create. +func (u *UsageLogUpsertOne) UpdateVideoDurationSeconds() *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.UpdateVideoDurationSeconds() + }) +} + +// ClearVideoDurationSeconds clears the value of the "video_duration_seconds" field. +func (u *UsageLogUpsertOne) ClearVideoDurationSeconds() *UsageLogUpsertOne { + return u.Update(func(s *UsageLogUpsert) { + s.ClearVideoDurationSeconds() + }) +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (u *UsageLogUpsertOne) SetCacheTTLOverridden(v bool) *UsageLogUpsertOne { return u.Update(func(s *UsageLogUpsert) { @@ -3722,6 +3918,76 @@ func (u *UsageLogUpsertBulk) ClearImageSizeBreakdown() *UsageLogUpsertBulk { }) } +// SetVideoCount sets the "video_count" field. +func (u *UsageLogUpsertBulk) SetVideoCount(v int) *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.SetVideoCount(v) + }) +} + +// AddVideoCount adds v to the "video_count" field. +func (u *UsageLogUpsertBulk) AddVideoCount(v int) *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.AddVideoCount(v) + }) +} + +// UpdateVideoCount sets the "video_count" field to the value that was provided on create. +func (u *UsageLogUpsertBulk) UpdateVideoCount() *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.UpdateVideoCount() + }) +} + +// SetVideoResolution sets the "video_resolution" field. +func (u *UsageLogUpsertBulk) SetVideoResolution(v string) *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.SetVideoResolution(v) + }) +} + +// UpdateVideoResolution sets the "video_resolution" field to the value that was provided on create. +func (u *UsageLogUpsertBulk) UpdateVideoResolution() *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.UpdateVideoResolution() + }) +} + +// ClearVideoResolution clears the value of the "video_resolution" field. +func (u *UsageLogUpsertBulk) ClearVideoResolution() *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.ClearVideoResolution() + }) +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (u *UsageLogUpsertBulk) SetVideoDurationSeconds(v int) *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.SetVideoDurationSeconds(v) + }) +} + +// AddVideoDurationSeconds adds v to the "video_duration_seconds" field. +func (u *UsageLogUpsertBulk) AddVideoDurationSeconds(v int) *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.AddVideoDurationSeconds(v) + }) +} + +// UpdateVideoDurationSeconds sets the "video_duration_seconds" field to the value that was provided on create. +func (u *UsageLogUpsertBulk) UpdateVideoDurationSeconds() *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.UpdateVideoDurationSeconds() + }) +} + +// ClearVideoDurationSeconds clears the value of the "video_duration_seconds" field. +func (u *UsageLogUpsertBulk) ClearVideoDurationSeconds() *UsageLogUpsertBulk { + return u.Update(func(s *UsageLogUpsert) { + s.ClearVideoDurationSeconds() + }) +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (u *UsageLogUpsertBulk) SetCacheTTLOverridden(v bool) *UsageLogUpsertBulk { return u.Update(func(s *UsageLogUpsert) { diff --git a/backend/ent/usagelog_update.go b/backend/ent/usagelog_update.go index e8fa003c63..00a65ccff1 100644 --- a/backend/ent/usagelog_update.go +++ b/backend/ent/usagelog_update.go @@ -811,6 +811,74 @@ func (_u *UsageLogUpdate) ClearImageSizeBreakdown() *UsageLogUpdate { return _u } +// SetVideoCount sets the "video_count" field. +func (_u *UsageLogUpdate) SetVideoCount(v int) *UsageLogUpdate { + _u.mutation.ResetVideoCount() + _u.mutation.SetVideoCount(v) + return _u +} + +// SetNillableVideoCount sets the "video_count" field if the given value is not nil. +func (_u *UsageLogUpdate) SetNillableVideoCount(v *int) *UsageLogUpdate { + if v != nil { + _u.SetVideoCount(*v) + } + return _u +} + +// AddVideoCount adds value to the "video_count" field. +func (_u *UsageLogUpdate) AddVideoCount(v int) *UsageLogUpdate { + _u.mutation.AddVideoCount(v) + return _u +} + +// SetVideoResolution sets the "video_resolution" field. +func (_u *UsageLogUpdate) SetVideoResolution(v string) *UsageLogUpdate { + _u.mutation.SetVideoResolution(v) + return _u +} + +// SetNillableVideoResolution sets the "video_resolution" field if the given value is not nil. +func (_u *UsageLogUpdate) SetNillableVideoResolution(v *string) *UsageLogUpdate { + if v != nil { + _u.SetVideoResolution(*v) + } + return _u +} + +// ClearVideoResolution clears the value of the "video_resolution" field. +func (_u *UsageLogUpdate) ClearVideoResolution() *UsageLogUpdate { + _u.mutation.ClearVideoResolution() + return _u +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (_u *UsageLogUpdate) SetVideoDurationSeconds(v int) *UsageLogUpdate { + _u.mutation.ResetVideoDurationSeconds() + _u.mutation.SetVideoDurationSeconds(v) + return _u +} + +// SetNillableVideoDurationSeconds sets the "video_duration_seconds" field if the given value is not nil. +func (_u *UsageLogUpdate) SetNillableVideoDurationSeconds(v *int) *UsageLogUpdate { + if v != nil { + _u.SetVideoDurationSeconds(*v) + } + return _u +} + +// AddVideoDurationSeconds adds value to the "video_duration_seconds" field. +func (_u *UsageLogUpdate) AddVideoDurationSeconds(v int) *UsageLogUpdate { + _u.mutation.AddVideoDurationSeconds(v) + return _u +} + +// ClearVideoDurationSeconds clears the value of the "video_duration_seconds" field. +func (_u *UsageLogUpdate) ClearVideoDurationSeconds() *UsageLogUpdate { + _u.mutation.ClearVideoDurationSeconds() + return _u +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (_u *UsageLogUpdate) SetCacheTTLOverridden(v bool) *UsageLogUpdate { _u.mutation.SetCacheTTLOverridden(v) @@ -979,6 +1047,11 @@ func (_u *UsageLogUpdate) check() error { return &ValidationError{Name: "image_size_source", err: fmt.Errorf(`ent: validator failed for field "UsageLog.image_size_source": %w`, err)} } } + if v, ok := _u.mutation.VideoResolution(); ok { + if err := usagelog.VideoResolutionValidator(v); err != nil { + return &ValidationError{Name: "video_resolution", err: fmt.Errorf(`ent: validator failed for field "UsageLog.video_resolution": %w`, err)} + } + } if _u.mutation.UserCleared() && len(_u.mutation.UserIDs()) > 0 { return errors.New(`ent: clearing a required unique edge "UsageLog.user"`) } @@ -1210,6 +1283,27 @@ func (_u *UsageLogUpdate) sqlSave(ctx context.Context) (_node int, err error) { if _u.mutation.ImageSizeBreakdownCleared() { _spec.ClearField(usagelog.FieldImageSizeBreakdown, field.TypeJSON) } + if value, ok := _u.mutation.VideoCount(); ok { + _spec.SetField(usagelog.FieldVideoCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedVideoCount(); ok { + _spec.AddField(usagelog.FieldVideoCount, field.TypeInt, value) + } + if value, ok := _u.mutation.VideoResolution(); ok { + _spec.SetField(usagelog.FieldVideoResolution, field.TypeString, value) + } + if _u.mutation.VideoResolutionCleared() { + _spec.ClearField(usagelog.FieldVideoResolution, field.TypeString) + } + if value, ok := _u.mutation.VideoDurationSeconds(); ok { + _spec.SetField(usagelog.FieldVideoDurationSeconds, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedVideoDurationSeconds(); ok { + _spec.AddField(usagelog.FieldVideoDurationSeconds, field.TypeInt, value) + } + if _u.mutation.VideoDurationSecondsCleared() { + _spec.ClearField(usagelog.FieldVideoDurationSeconds, field.TypeInt) + } if value, ok := _u.mutation.CacheTTLOverridden(); ok { _spec.SetField(usagelog.FieldCacheTTLOverridden, field.TypeBool, value) } @@ -2157,6 +2251,74 @@ func (_u *UsageLogUpdateOne) ClearImageSizeBreakdown() *UsageLogUpdateOne { return _u } +// SetVideoCount sets the "video_count" field. +func (_u *UsageLogUpdateOne) SetVideoCount(v int) *UsageLogUpdateOne { + _u.mutation.ResetVideoCount() + _u.mutation.SetVideoCount(v) + return _u +} + +// SetNillableVideoCount sets the "video_count" field if the given value is not nil. +func (_u *UsageLogUpdateOne) SetNillableVideoCount(v *int) *UsageLogUpdateOne { + if v != nil { + _u.SetVideoCount(*v) + } + return _u +} + +// AddVideoCount adds value to the "video_count" field. +func (_u *UsageLogUpdateOne) AddVideoCount(v int) *UsageLogUpdateOne { + _u.mutation.AddVideoCount(v) + return _u +} + +// SetVideoResolution sets the "video_resolution" field. +func (_u *UsageLogUpdateOne) SetVideoResolution(v string) *UsageLogUpdateOne { + _u.mutation.SetVideoResolution(v) + return _u +} + +// SetNillableVideoResolution sets the "video_resolution" field if the given value is not nil. +func (_u *UsageLogUpdateOne) SetNillableVideoResolution(v *string) *UsageLogUpdateOne { + if v != nil { + _u.SetVideoResolution(*v) + } + return _u +} + +// ClearVideoResolution clears the value of the "video_resolution" field. +func (_u *UsageLogUpdateOne) ClearVideoResolution() *UsageLogUpdateOne { + _u.mutation.ClearVideoResolution() + return _u +} + +// SetVideoDurationSeconds sets the "video_duration_seconds" field. +func (_u *UsageLogUpdateOne) SetVideoDurationSeconds(v int) *UsageLogUpdateOne { + _u.mutation.ResetVideoDurationSeconds() + _u.mutation.SetVideoDurationSeconds(v) + return _u +} + +// SetNillableVideoDurationSeconds sets the "video_duration_seconds" field if the given value is not nil. +func (_u *UsageLogUpdateOne) SetNillableVideoDurationSeconds(v *int) *UsageLogUpdateOne { + if v != nil { + _u.SetVideoDurationSeconds(*v) + } + return _u +} + +// AddVideoDurationSeconds adds value to the "video_duration_seconds" field. +func (_u *UsageLogUpdateOne) AddVideoDurationSeconds(v int) *UsageLogUpdateOne { + _u.mutation.AddVideoDurationSeconds(v) + return _u +} + +// ClearVideoDurationSeconds clears the value of the "video_duration_seconds" field. +func (_u *UsageLogUpdateOne) ClearVideoDurationSeconds() *UsageLogUpdateOne { + _u.mutation.ClearVideoDurationSeconds() + return _u +} + // SetCacheTTLOverridden sets the "cache_ttl_overridden" field. func (_u *UsageLogUpdateOne) SetCacheTTLOverridden(v bool) *UsageLogUpdateOne { _u.mutation.SetCacheTTLOverridden(v) @@ -2338,6 +2500,11 @@ func (_u *UsageLogUpdateOne) check() error { return &ValidationError{Name: "image_size_source", err: fmt.Errorf(`ent: validator failed for field "UsageLog.image_size_source": %w`, err)} } } + if v, ok := _u.mutation.VideoResolution(); ok { + if err := usagelog.VideoResolutionValidator(v); err != nil { + return &ValidationError{Name: "video_resolution", err: fmt.Errorf(`ent: validator failed for field "UsageLog.video_resolution": %w`, err)} + } + } if _u.mutation.UserCleared() && len(_u.mutation.UserIDs()) > 0 { return errors.New(`ent: clearing a required unique edge "UsageLog.user"`) } @@ -2586,6 +2753,27 @@ func (_u *UsageLogUpdateOne) sqlSave(ctx context.Context) (_node *UsageLog, err if _u.mutation.ImageSizeBreakdownCleared() { _spec.ClearField(usagelog.FieldImageSizeBreakdown, field.TypeJSON) } + if value, ok := _u.mutation.VideoCount(); ok { + _spec.SetField(usagelog.FieldVideoCount, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedVideoCount(); ok { + _spec.AddField(usagelog.FieldVideoCount, field.TypeInt, value) + } + if value, ok := _u.mutation.VideoResolution(); ok { + _spec.SetField(usagelog.FieldVideoResolution, field.TypeString, value) + } + if _u.mutation.VideoResolutionCleared() { + _spec.ClearField(usagelog.FieldVideoResolution, field.TypeString) + } + if value, ok := _u.mutation.VideoDurationSeconds(); ok { + _spec.SetField(usagelog.FieldVideoDurationSeconds, field.TypeInt, value) + } + if value, ok := _u.mutation.AddedVideoDurationSeconds(); ok { + _spec.AddField(usagelog.FieldVideoDurationSeconds, field.TypeInt, value) + } + if _u.mutation.VideoDurationSecondsCleared() { + _spec.ClearField(usagelog.FieldVideoDurationSeconds, field.TypeInt) + } if value, ok := _u.mutation.CacheTTLOverridden(); ok { _spec.SetField(usagelog.FieldCacheTTLOverridden, field.TypeBool, value) } diff --git a/backend/ent/user.go b/backend/ent/user.go index 486f2f64d9..299a8d627f 100644 --- a/backend/ent/user.go +++ b/backend/ent/user.go @@ -31,6 +31,8 @@ type User struct { Role string `json:"role,omitempty"` // Balance holds the value of the "balance" field. Balance float64 `json:"balance,omitempty"` + // FrozenBalance holds the value of the "frozen_balance" field. + FrozenBalance float64 `json:"frozen_balance,omitempty"` // Concurrency holds the value of the "concurrency" field. Concurrency int `json:"concurrency,omitempty"` // Status holds the value of the "status" field. @@ -237,7 +239,7 @@ func (*User) scanValues(columns []string) ([]any, error) { switch columns[i] { case user.FieldTotpEnabled, user.FieldBalanceNotifyEnabled: values[i] = new(sql.NullBool) - case user.FieldBalance, user.FieldBalanceNotifyThreshold, user.FieldTotalRecharged: + case user.FieldBalance, user.FieldFrozenBalance, user.FieldBalanceNotifyThreshold, user.FieldTotalRecharged: values[i] = new(sql.NullFloat64) case user.FieldID, user.FieldConcurrency, user.FieldRpmLimit: values[i] = new(sql.NullInt64) @@ -309,6 +311,12 @@ func (_m *User) assignValues(columns []string, values []any) error { } else if value.Valid { _m.Balance = value.Float64 } + case user.FieldFrozenBalance: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field frozen_balance", values[i]) + } else if value.Valid { + _m.FrozenBalance = value.Float64 + } case user.FieldConcurrency: if value, ok := values[i].(*sql.NullInt64); !ok { return fmt.Errorf("unexpected type %T for field concurrency", values[i]) @@ -539,6 +547,9 @@ func (_m *User) String() string { builder.WriteString("balance=") builder.WriteString(fmt.Sprintf("%v", _m.Balance)) builder.WriteString(", ") + builder.WriteString("frozen_balance=") + builder.WriteString(fmt.Sprintf("%v", _m.FrozenBalance)) + builder.WriteString(", ") builder.WriteString("concurrency=") builder.WriteString(fmt.Sprintf("%v", _m.Concurrency)) builder.WriteString(", ") diff --git a/backend/ent/user/user.go b/backend/ent/user/user.go index ff40445bda..ae1a84494d 100644 --- a/backend/ent/user/user.go +++ b/backend/ent/user/user.go @@ -29,6 +29,8 @@ const ( FieldRole = "role" // FieldBalance holds the string denoting the balance field in the database. FieldBalance = "balance" + // FieldFrozenBalance holds the string denoting the frozen_balance field in the database. + FieldFrozenBalance = "frozen_balance" // FieldConcurrency holds the string denoting the concurrency field in the database. FieldConcurrency = "concurrency" // FieldStatus holds the string denoting the status field in the database. @@ -199,6 +201,7 @@ var Columns = []string{ FieldPasswordHash, FieldRole, FieldBalance, + FieldFrozenBalance, FieldConcurrency, FieldStatus, FieldUsername, @@ -257,6 +260,8 @@ var ( RoleValidator func(string) error // DefaultBalance holds the default value on creation for the "balance" field. DefaultBalance float64 + // DefaultFrozenBalance holds the default value on creation for the "frozen_balance" field. + DefaultFrozenBalance float64 // DefaultConcurrency holds the default value on creation for the "concurrency" field. DefaultConcurrency int // DefaultStatus holds the default value on creation for the "status" field. @@ -330,6 +335,11 @@ func ByBalance(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldBalance, opts...).ToFunc() } +// ByFrozenBalance orders the results by the frozen_balance field. +func ByFrozenBalance(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldFrozenBalance, opts...).ToFunc() +} + // ByConcurrency orders the results by the concurrency field. func ByConcurrency(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldConcurrency, opts...).ToFunc() diff --git a/backend/ent/user/where.go b/backend/ent/user/where.go index a18cf49767..c2a71f6172 100644 --- a/backend/ent/user/where.go +++ b/backend/ent/user/where.go @@ -90,6 +90,11 @@ func Balance(v float64) predicate.User { return predicate.User(sql.FieldEQ(FieldBalance, v)) } +// FrozenBalance applies equality check predicate on the "frozen_balance" field. It's identical to FrozenBalanceEQ. +func FrozenBalance(v float64) predicate.User { + return predicate.User(sql.FieldEQ(FieldFrozenBalance, v)) +} + // Concurrency applies equality check predicate on the "concurrency" field. It's identical to ConcurrencyEQ. func Concurrency(v int) predicate.User { return predicate.User(sql.FieldEQ(FieldConcurrency, v)) @@ -535,6 +540,46 @@ func BalanceLTE(v float64) predicate.User { return predicate.User(sql.FieldLTE(FieldBalance, v)) } +// FrozenBalanceEQ applies the EQ predicate on the "frozen_balance" field. +func FrozenBalanceEQ(v float64) predicate.User { + return predicate.User(sql.FieldEQ(FieldFrozenBalance, v)) +} + +// FrozenBalanceNEQ applies the NEQ predicate on the "frozen_balance" field. +func FrozenBalanceNEQ(v float64) predicate.User { + return predicate.User(sql.FieldNEQ(FieldFrozenBalance, v)) +} + +// FrozenBalanceIn applies the In predicate on the "frozen_balance" field. +func FrozenBalanceIn(vs ...float64) predicate.User { + return predicate.User(sql.FieldIn(FieldFrozenBalance, vs...)) +} + +// FrozenBalanceNotIn applies the NotIn predicate on the "frozen_balance" field. +func FrozenBalanceNotIn(vs ...float64) predicate.User { + return predicate.User(sql.FieldNotIn(FieldFrozenBalance, vs...)) +} + +// FrozenBalanceGT applies the GT predicate on the "frozen_balance" field. +func FrozenBalanceGT(v float64) predicate.User { + return predicate.User(sql.FieldGT(FieldFrozenBalance, v)) +} + +// FrozenBalanceGTE applies the GTE predicate on the "frozen_balance" field. +func FrozenBalanceGTE(v float64) predicate.User { + return predicate.User(sql.FieldGTE(FieldFrozenBalance, v)) +} + +// FrozenBalanceLT applies the LT predicate on the "frozen_balance" field. +func FrozenBalanceLT(v float64) predicate.User { + return predicate.User(sql.FieldLT(FieldFrozenBalance, v)) +} + +// FrozenBalanceLTE applies the LTE predicate on the "frozen_balance" field. +func FrozenBalanceLTE(v float64) predicate.User { + return predicate.User(sql.FieldLTE(FieldFrozenBalance, v)) +} + // ConcurrencyEQ applies the EQ predicate on the "concurrency" field. func ConcurrencyEQ(v int) predicate.User { return predicate.User(sql.FieldEQ(FieldConcurrency, v)) diff --git a/backend/ent/user_create.go b/backend/ent/user_create.go index 92f1bd5e07..b5bdf986a0 100644 --- a/backend/ent/user_create.go +++ b/backend/ent/user_create.go @@ -116,6 +116,20 @@ func (_c *UserCreate) SetNillableBalance(v *float64) *UserCreate { return _c } +// SetFrozenBalance sets the "frozen_balance" field. +func (_c *UserCreate) SetFrozenBalance(v float64) *UserCreate { + _c.mutation.SetFrozenBalance(v) + return _c +} + +// SetNillableFrozenBalance sets the "frozen_balance" field if the given value is not nil. +func (_c *UserCreate) SetNillableFrozenBalance(v *float64) *UserCreate { + if v != nil { + _c.SetFrozenBalance(*v) + } + return _c +} + // SetConcurrency sets the "concurrency" field. func (_c *UserCreate) SetConcurrency(v int) *UserCreate { _c.mutation.SetConcurrency(v) @@ -594,6 +608,10 @@ func (_c *UserCreate) defaults() error { v := user.DefaultBalance _c.mutation.SetBalance(v) } + if _, ok := _c.mutation.FrozenBalance(); !ok { + v := user.DefaultFrozenBalance + _c.mutation.SetFrozenBalance(v) + } if _, ok := _c.mutation.Concurrency(); !ok { v := user.DefaultConcurrency _c.mutation.SetConcurrency(v) @@ -676,6 +694,9 @@ func (_c *UserCreate) check() error { if _, ok := _c.mutation.Balance(); !ok { return &ValidationError{Name: "balance", err: errors.New(`ent: missing required field "User.balance"`)} } + if _, ok := _c.mutation.FrozenBalance(); !ok { + return &ValidationError{Name: "frozen_balance", err: errors.New(`ent: missing required field "User.frozen_balance"`)} + } if _, ok := _c.mutation.Concurrency(); !ok { return &ValidationError{Name: "concurrency", err: errors.New(`ent: missing required field "User.concurrency"`)} } @@ -779,6 +800,10 @@ func (_c *UserCreate) createSpec() (*User, *sqlgraph.CreateSpec) { _spec.SetField(user.FieldBalance, field.TypeFloat64, value) _node.Balance = value } + if value, ok := _c.mutation.FrozenBalance(); ok { + _spec.SetField(user.FieldFrozenBalance, field.TypeFloat64, value) + _node.FrozenBalance = value + } if value, ok := _c.mutation.Concurrency(); ok { _spec.SetField(user.FieldConcurrency, field.TypeInt, value) _node.Concurrency = value @@ -1191,6 +1216,24 @@ func (u *UserUpsert) AddBalance(v float64) *UserUpsert { return u } +// SetFrozenBalance sets the "frozen_balance" field. +func (u *UserUpsert) SetFrozenBalance(v float64) *UserUpsert { + u.Set(user.FieldFrozenBalance, v) + return u +} + +// UpdateFrozenBalance sets the "frozen_balance" field to the value that was provided on create. +func (u *UserUpsert) UpdateFrozenBalance() *UserUpsert { + u.SetExcluded(user.FieldFrozenBalance) + return u +} + +// AddFrozenBalance adds v to the "frozen_balance" field. +func (u *UserUpsert) AddFrozenBalance(v float64) *UserUpsert { + u.Add(user.FieldFrozenBalance, v) + return u +} + // SetConcurrency sets the "concurrency" field. func (u *UserUpsert) SetConcurrency(v int) *UserUpsert { u.Set(user.FieldConcurrency, v) @@ -1580,6 +1623,27 @@ func (u *UserUpsertOne) UpdateBalance() *UserUpsertOne { }) } +// SetFrozenBalance sets the "frozen_balance" field. +func (u *UserUpsertOne) SetFrozenBalance(v float64) *UserUpsertOne { + return u.Update(func(s *UserUpsert) { + s.SetFrozenBalance(v) + }) +} + +// AddFrozenBalance adds v to the "frozen_balance" field. +func (u *UserUpsertOne) AddFrozenBalance(v float64) *UserUpsertOne { + return u.Update(func(s *UserUpsert) { + s.AddFrozenBalance(v) + }) +} + +// UpdateFrozenBalance sets the "frozen_balance" field to the value that was provided on create. +func (u *UserUpsertOne) UpdateFrozenBalance() *UserUpsertOne { + return u.Update(func(s *UserUpsert) { + s.UpdateFrozenBalance() + }) +} + // SetConcurrency sets the "concurrency" field. func (u *UserUpsertOne) SetConcurrency(v int) *UserUpsertOne { return u.Update(func(s *UserUpsert) { @@ -2176,6 +2240,27 @@ func (u *UserUpsertBulk) UpdateBalance() *UserUpsertBulk { }) } +// SetFrozenBalance sets the "frozen_balance" field. +func (u *UserUpsertBulk) SetFrozenBalance(v float64) *UserUpsertBulk { + return u.Update(func(s *UserUpsert) { + s.SetFrozenBalance(v) + }) +} + +// AddFrozenBalance adds v to the "frozen_balance" field. +func (u *UserUpsertBulk) AddFrozenBalance(v float64) *UserUpsertBulk { + return u.Update(func(s *UserUpsert) { + s.AddFrozenBalance(v) + }) +} + +// UpdateFrozenBalance sets the "frozen_balance" field to the value that was provided on create. +func (u *UserUpsertBulk) UpdateFrozenBalance() *UserUpsertBulk { + return u.Update(func(s *UserUpsert) { + s.UpdateFrozenBalance() + }) +} + // SetConcurrency sets the "concurrency" field. func (u *UserUpsertBulk) SetConcurrency(v int) *UserUpsertBulk { return u.Update(func(s *UserUpsert) { diff --git a/backend/ent/user_update.go b/backend/ent/user_update.go index 67d3f8e6bb..6df9b320da 100644 --- a/backend/ent/user_update.go +++ b/backend/ent/user_update.go @@ -129,6 +129,27 @@ func (_u *UserUpdate) AddBalance(v float64) *UserUpdate { return _u } +// SetFrozenBalance sets the "frozen_balance" field. +func (_u *UserUpdate) SetFrozenBalance(v float64) *UserUpdate { + _u.mutation.ResetFrozenBalance() + _u.mutation.SetFrozenBalance(v) + return _u +} + +// SetNillableFrozenBalance sets the "frozen_balance" field if the given value is not nil. +func (_u *UserUpdate) SetNillableFrozenBalance(v *float64) *UserUpdate { + if v != nil { + _u.SetFrozenBalance(*v) + } + return _u +} + +// AddFrozenBalance adds value to the "frozen_balance" field. +func (_u *UserUpdate) AddFrozenBalance(v float64) *UserUpdate { + _u.mutation.AddFrozenBalance(v) + return _u +} + // SetConcurrency sets the "concurrency" field. func (_u *UserUpdate) SetConcurrency(v int) *UserUpdate { _u.mutation.ResetConcurrency() @@ -997,6 +1018,12 @@ func (_u *UserUpdate) sqlSave(ctx context.Context) (_node int, err error) { if value, ok := _u.mutation.AddedBalance(); ok { _spec.AddField(user.FieldBalance, field.TypeFloat64, value) } + if value, ok := _u.mutation.FrozenBalance(); ok { + _spec.SetField(user.FieldFrozenBalance, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedFrozenBalance(); ok { + _spec.AddField(user.FieldFrozenBalance, field.TypeFloat64, value) + } if value, ok := _u.mutation.Concurrency(); ok { _spec.SetField(user.FieldConcurrency, field.TypeInt, value) } @@ -1778,6 +1805,27 @@ func (_u *UserUpdateOne) AddBalance(v float64) *UserUpdateOne { return _u } +// SetFrozenBalance sets the "frozen_balance" field. +func (_u *UserUpdateOne) SetFrozenBalance(v float64) *UserUpdateOne { + _u.mutation.ResetFrozenBalance() + _u.mutation.SetFrozenBalance(v) + return _u +} + +// SetNillableFrozenBalance sets the "frozen_balance" field if the given value is not nil. +func (_u *UserUpdateOne) SetNillableFrozenBalance(v *float64) *UserUpdateOne { + if v != nil { + _u.SetFrozenBalance(*v) + } + return _u +} + +// AddFrozenBalance adds value to the "frozen_balance" field. +func (_u *UserUpdateOne) AddFrozenBalance(v float64) *UserUpdateOne { + _u.mutation.AddFrozenBalance(v) + return _u +} + // SetConcurrency sets the "concurrency" field. func (_u *UserUpdateOne) SetConcurrency(v int) *UserUpdateOne { _u.mutation.ResetConcurrency() @@ -2676,6 +2724,12 @@ func (_u *UserUpdateOne) sqlSave(ctx context.Context) (_node *User, err error) { if value, ok := _u.mutation.AddedBalance(); ok { _spec.AddField(user.FieldBalance, field.TypeFloat64, value) } + if value, ok := _u.mutation.FrozenBalance(); ok { + _spec.SetField(user.FieldFrozenBalance, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedFrozenBalance(); ok { + _spec.AddField(user.FieldFrozenBalance, field.TypeFloat64, value) + } if value, ok := _u.mutation.Concurrency(); ok { _spec.SetField(user.FieldConcurrency, field.TypeInt, value) } diff --git a/backend/go.mod b/backend/go.mod index d4e0af6c13..a06f06437d 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -1,6 +1,6 @@ module github.com/Wei-Shaw/sub2api -go 1.26.4 +go 1.26.5 require ( entgo.io/ent v0.14.5 @@ -8,11 +8,10 @@ require ( github.com/alicebob/miniredis/v2 v2.38.0 github.com/alitto/pond/v2 v2.6.2 github.com/andybalholm/brotli v1.2.0 - github.com/aws/aws-sdk-go-v2 v1.41.3 + github.com/aws/aws-sdk-go-v2 v1.41.5 github.com/aws/aws-sdk-go-v2/config v1.32.10 github.com/aws/aws-sdk-go-v2/credentials v1.19.10 - github.com/aws/aws-sdk-go-v2/service/s3 v1.96.2 - github.com/aws/smithy-go v1.24.2 + github.com/aws/aws-sdk-go-v2/service/s3 v1.97.3 github.com/cespare/xxhash/v2 v2.3.0 github.com/coder/websocket v1.8.14 github.com/dgraph-io/ristretto v0.2.0 @@ -39,6 +38,7 @@ require ( github.com/testcontainers/testcontainers-go/modules/redis v0.40.0 github.com/tidwall/gjson v1.18.0 github.com/tidwall/sjson v1.2.5 + github.com/tiktoken-go/tokenizer v0.8.0 github.com/wechatpay-apiv3/wechatpay-go v0.2.21 github.com/zeromicro/go-zero v1.9.4 go.uber.org/zap v1.24.0 @@ -59,20 +59,21 @@ require ( github.com/Microsoft/go-winio v0.6.2 // indirect github.com/agext/levenshtein v1.2.3 // indirect github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect - github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.5 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.18 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.18 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.18 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4 // indirect - github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.18 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.5 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.10 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.18 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.18 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 // indirect github.com/aws/aws-sdk-go-v2/service/signin v1.0.6 // indirect github.com/aws/aws-sdk-go-v2/service/sso v1.30.11 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.15 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.41.7 // indirect + github.com/aws/smithy-go v1.24.2 // indirect github.com/bmatcuk/doublestar v1.3.4 // indirect github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect github.com/bytedance/sonic v1.9.1 // indirect @@ -156,7 +157,6 @@ require ( github.com/testcontainers/testcontainers-go v0.40.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect - github.com/tiktoken-go/tokenizer v0.8.0 // indirect github.com/tklauser/go-sysconf v0.3.12 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect diff --git a/backend/go.sum b/backend/go.sum index 1574798015..4738443bb9 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -24,34 +24,34 @@ github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwTo github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/apparentlymart/go-textseg/v15 v15.0.0 h1:uYvfpb3DyLSCGWnctWKGj857c6ew1u1fNQOlOtuGxQY= github.com/apparentlymart/go-textseg/v15 v15.0.0/go.mod h1:K8XmNZdhEBkdlyDdvbmmsvpAG721bKi0joRfFdHIWJ4= -github.com/aws/aws-sdk-go-v2 v1.41.3 h1:4kQ/fa22KjDt13QCy1+bYADvdgcxpfH18f0zP542kZA= -github.com/aws/aws-sdk-go-v2 v1.41.3/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o= -github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.5 h1:zWFmPmgw4sveAYi1mRqG+E/g0461cJ5M4bJ8/nc6d3Q= -github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.5/go.mod h1:nVUlMLVV8ycXSb7mSkcNu9e3v/1TJq2RTlrPwhYWr5c= +github.com/aws/aws-sdk-go-v2 v1.41.5 h1:dj5kopbwUsVUVFgO4Fi5BIT3t4WyqIDjGKCangnV/yY= +github.com/aws/aws-sdk-go-v2 v1.41.5/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 h1:eBMB84YGghSocM7PsjmmPffTa+1FBUeNvGvFou6V/4o= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8/go.mod h1:lyw7GFp3qENLh7kwzf7iMzAxDn+NzjXEAGjKS2UOKqI= github.com/aws/aws-sdk-go-v2/config v1.32.10 h1:9DMthfO6XWZYLfzZglAgW5Fyou2nRI5CuV44sTedKBI= github.com/aws/aws-sdk-go-v2/config v1.32.10/go.mod h1:2rUIOnA2JaiqYmSKYmRJlcMWy6qTj1vuRFscppSBMcw= github.com/aws/aws-sdk-go-v2/credentials v1.19.10 h1:EEhmEUFCE1Yhl7vDhNOI5OCL/iKMdkkYFTRpZXNw7m8= github.com/aws/aws-sdk-go-v2/credentials v1.19.10/go.mod h1:RnnlFCAlxQCkN2Q379B67USkBMu1PipEEiibzYN5UTE= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.18 h1:Ii4s+Sq3yDfaMLpjrJsqD6SmG/Wq/P5L/hw2qa78UAY= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.18/go.mod h1:6x81qnY++ovptLE6nWQeWrpXxbnlIex+4H4eYYGcqfc= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.18 h1:F43zk1vemYIqPAwhjTjYIz0irU2EY7sOb/F5eJ3HuyM= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.18/go.mod h1:w1jdlZXrGKaJcNoL+Nnrj+k5wlpGXqnNrKoP22HvAug= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.18 h1:xCeWVjj0ki0l3nruoyP2slHsGArMxeiiaoPN5QZH6YQ= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.18/go.mod h1:r/eLGuGCBw6l36ZRWiw6PaZwPXb6YOj+i/7MizNl5/k= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 h1:Rgg6wvjjtX8bNHcvi9OnXWwcE0a2vGpbwmtICOsvcf4= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21/go.mod h1:A/kJFst/nm//cyqonihbdpQZwiUhhzpqTsdbhDdRF9c= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 h1:PEgGVtPoB6NTpPrBgqSE5hE/o47Ij9qk/SEZFbUOe9A= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21/go.mod h1:p+hz+PRAYlY3zcpJhPwXlLC4C+kqn70WIHwnzAfs6ps= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4 h1:WKuaxf++XKWlHWu9ECbMlha8WOEGm0OUEZqm4K/Gcfk= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4/go.mod h1:ZWy7j6v1vWGmPReu0iSGvRiise4YI5SkR3OHKTZ6Wuc= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.18 h1:eZioDaZGJ0tMM4gzmkNIO2aAoQd+je7Ug7TkvAzlmkU= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.18/go.mod h1:CCXwUKAJdoWr6/NcxZ+zsiPr6oH/Q5aTooRGYieAyj4= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.5 h1:CeY9LUdur+Dxoeldqoun6y4WtJ3RQtzk0JMP2gfUay0= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.5/go.mod h1:AZLZf2fMaahW5s/wMRciu1sYbdsikT/UHwbUjOdEVTc= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.10 h1:fJvQ5mIBVfKtiyx0AHY6HeWcRX5LGANLpq8SVR+Uazs= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.10/go.mod h1:Kzm5e6OmNH8VMkgK9t+ry5jEih4Y8whqs+1hrkxim1I= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.18 h1:LTRCYFlnnKFlKsyIQxKhJuDuA3ZkrDQMRYm6rXiHlLY= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.18/go.mod h1:XhwkgGG6bHSd00nO/mexWTcTjgd6PjuvWQMqSn2UaEk= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.18 h1:/A/xDuZAVD2BpsS2fftFRo/NoEKQJ8YTnJDEHBy2Gtg= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.18/go.mod h1:hWe9b4f+djUQGmyiGEeOnZv69dtMSgpDRIvNMvuvzvY= -github.com/aws/aws-sdk-go-v2/service/s3 v1.96.2 h1:M1A9AjcFwlxTLuf0Faj88L8Iqw0n/AJHjpZTQzMMsSc= -github.com/aws/aws-sdk-go-v2/service/s3 v1.96.2/go.mod h1:KsdTV6Q9WKUZm2mNJnUFmIoXfZux91M3sr/a4REX8e0= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 h1:rWyie/PxDRIdhNf4DzRk0lvjVOqFJuNnO8WwaIRVxzQ= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22/go.mod h1:zd/JsJ4P7oGfUhXn1VyLqaRZwPmZwg44Jf2dS84Dm3Y= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 h1:5EniKhLZe4xzL7a+fU3C2tfUN4nWIqlLesfrjkuPFTY= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7/go.mod h1:x0nZssQ3qZSnIcePWLvcoFisRXJzcTVvYpAAdYX8+GI= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 h1:JRaIgADQS/U6uXDqlPiefP32yXTda7Kqfx+LgspooZM= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13/go.mod h1:CEuVn5WqOMilYl+tbccq8+N2ieCy0gVn3OtRb0vBNNM= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 h1:c31//R3xgIJMSC8S6hEVq+38DcvUlgFY0FM6mSI5oto= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21/go.mod h1:r6+pf23ouCB718FUxaqzZdbpYFyDtehyZcmP5KL9FkA= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 h1:ZlvrNcHSFFWURB8avufQq9gFsheUgjVD9536obIknfM= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21/go.mod h1:cv3TNhVrssKR0O/xxLJVRfd2oazSnZnkUeTf6ctUwfQ= +github.com/aws/aws-sdk-go-v2/service/s3 v1.97.3 h1:HwxWTbTrIHm5qY+CAEur0s/figc3qwvLWsNkF4RPToo= +github.com/aws/aws-sdk-go-v2/service/s3 v1.97.3/go.mod h1:uoA43SdFwacedBfSgfFSjjCvYe8aYBS7EnU5GZ/YKMM= github.com/aws/aws-sdk-go-v2/service/signin v1.0.6 h1:MzORe+J94I+hYu2a6XmV5yC9huoTv8NRcCrUNedDypQ= github.com/aws/aws-sdk-go-v2/service/signin v1.0.6/go.mod h1:hXzcHLARD7GeWnifd8j9RWqtfIgxj4/cAtIVIK7hg8g= github.com/aws/aws-sdk-go-v2/service/sso v1.30.11 h1:7oGD8KPfBOJGXiCoRKrrrQkbvCp8N++u36hrLMPey6o= @@ -166,8 +166,6 @@ github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= -github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE= -github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4= @@ -222,8 +220,6 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= -github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U= -github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM= github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI= @@ -257,8 +253,6 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= -github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec= -github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= @@ -288,8 +282,6 @@ github.com/refraction-networking/utls v1.8.2 h1:j4Q1gJj0xngdeH+Ox/qND11aEfhpgoEv github.com/refraction-networking/utls v1.8.2/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= -github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY= -github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= @@ -322,8 +314,6 @@ github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8= github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY= github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0= github.com/spf13/cast v1.6.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= -github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I= -github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0= github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/viper v1.18.2 h1:LUXCnvUvSM6FXAsj6nnfc8Q2tp1dIgUfY9Kc8GsSOiQ= @@ -419,20 +409,14 @@ go.uber.org/zap v1.24.0/go.mod h1:2kMP+WWQ8aoFoedH3T2sq6iJ2yDWpHbP0f6MQbS9Gkg= golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= golang.org/x/arch v0.3.0 h1:02VY4/ZcO/gBOH6PUaoiptASxtXU10jazRCP865E97k= golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= -golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI= -golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q= golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI= golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8= golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY= golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70= golang.org/x/image v0.39.0 h1:skVYidAEVKgn8lZ602XO75asgXBgLj9G/FE3RbuPFww= golang.org/x/image v0.39.0/go.mod h1:sIbmppfU+xFLPIG0FoVUTvyBMmgng1/XAMhQ2ft0hpA= -golang.org/x/mod v0.34.0 h1:xIHgNUUnW6sYkcM5Jleh05DvLOtwc6RitGHbDk4akRI= -golang.org/x/mod v0.34.0/go.mod h1:ykgH52iCZe79kzLLMhyCUzhMci+nQj+0XkbXpNYtVjY= golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= -golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA= -golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs= golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= @@ -446,22 +430,14 @@ golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI= -golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/term v0.42.0 h1:UiKe+zDFmJobeJ5ggPwOshJIVt6/Ft0rcfrXZDLWAWY= -golang.org/x/term v0.42.0/go.mod h1:Dq/D+snpsbazcBG5+F9Q1n2rXV8Ma+71xEjTRufARgY= golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4= golang.org/x/term v0.43.0/go.mod h1:lrhlHNdQJHO+1qVYiHfFKVuVioJIheAc3fBSMFYEIsk= -golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg= -golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164= golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE= golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= -golang.org/x/tools v0.43.0 h1:12BdW9CeB3Z+J/I/wj34VMl8X+fEXBxVR90JeMX5E7s= -golang.org/x/tools v0.43.0/go.mod h1:uHkMso649BX2cZK6+RpuIPXS3ho2hZo4FVwfoy1vIk0= golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c= golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 99fedb5b1c..df3afb6c7e 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -29,7 +29,7 @@ const ( // DefaultCSPPolicy is the default Content-Security-Policy with nonce support // __CSP_NONCE__ will be replaced with actual nonce at request time by the SecurityHeaders middleware -const DefaultCSPPolicy = "default-src 'self'; script-src 'self' __CSP_NONCE__ https://challenges.cloudflare.com https://static.cloudflareinsights.com https://*.stripe.com https://static.airwallex.com https://checkout.airwallex.com https://static-demo.airwallex.com https://checkout-demo.airwallex.com; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com https://static.airwallex.com https://checkout.airwallex.com https://static-demo.airwallex.com https://checkout-demo.airwallex.com; img-src 'self' data: https:; font-src 'self' data: https://fonts.gstatic.com; connect-src 'self' https:; frame-src https://challenges.cloudflare.com https://*.stripe.com https://checkout.airwallex.com https://checkout-demo.airwallex.com; frame-ancestors 'none'; base-uri 'self'; form-action 'self'" +const DefaultCSPPolicy = "default-src 'self'; script-src 'self' __CSP_NONCE__ https://challenges.cloudflare.com https://static.cloudflareinsights.com https://*.stripe.com https://static.airwallex.com https://checkout.airwallex.com https://static-demo.airwallex.com https://checkout-demo.airwallex.com; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com https://static.airwallex.com https://checkout.airwallex.com https://static-demo.airwallex.com https://checkout-demo.airwallex.com; img-src 'self' data: blob: https:; font-src 'self' data: https://fonts.gstatic.com; connect-src 'self' https:; frame-src https://challenges.cloudflare.com https://*.stripe.com https://checkout.airwallex.com https://checkout-demo.airwallex.com; frame-ancestors 'none'; base-uri 'self'; form-action 'self'" // UMQ(用户消息队列)模式常量 const ( @@ -93,6 +93,7 @@ type Config struct { Gemini GeminiConfig `mapstructure:"gemini"` Update UpdateConfig `mapstructure:"update"` Idempotency IdempotencyConfig `mapstructure:"idempotency"` + BatchImage BatchImageConfig `mapstructure:"batch_image"` } type LogConfig struct { @@ -175,6 +176,56 @@ type IdempotencyConfig struct { CleanupBatchSize int `mapstructure:"cleanup_batch_size"` } +type BatchImageConfig struct { + Enabled bool `mapstructure:"enabled"` + MaxItemsPerJobDefault int `mapstructure:"max_items_per_job_default"` + MaxItemsPerJobTrial int `mapstructure:"max_items_per_job_trial"` + MaxOutputImagesPerJob int `mapstructure:"max_output_images_per_job"` + MaxOutputImagesPerItem int `mapstructure:"max_output_images_per_item"` + MaxPromptCharsPerItem int `mapstructure:"max_prompt_chars_per_item"` + MaxReferenceImagesPerJob int `mapstructure:"max_reference_images_per_job"` + MaxReferenceInlineBytesPerJob int `mapstructure:"max_reference_inline_bytes_per_job"` + DefaultResponseMimeType string `mapstructure:"default_response_mime_type"` + DefaultImageSize string `mapstructure:"default_image_size"` + MaxDownloadItemsZip int `mapstructure:"max_download_items_zip"` + MaxDownloadBytesPerRequest int64 `mapstructure:"max_download_bytes_per_request"` + MaxDownloadDurationSeconds int `mapstructure:"max_download_duration_seconds"` + MaxDownloadConcurrencyPerUser int `mapstructure:"max_download_concurrency_per_user"` + InputRetentionAfterTerminalHours int `mapstructure:"input_retention_after_terminal_hours"` + OutputRetentionAfterTerminalHours int `mapstructure:"output_retention_after_terminal_hours"` + OutputRetentionMaxDays int `mapstructure:"output_retention_max_days"` + CleanupIntervalMinutes int `mapstructure:"cleanup_interval_minutes"` + CleanupBatchSize int `mapstructure:"cleanup_batch_size"` + QueueEnabled bool `mapstructure:"queue_enabled"` + QueueReadyKey string `mapstructure:"queue_ready_key"` + QueueDelayedKey string `mapstructure:"queue_delayed_key"` + QueueActiveKey string `mapstructure:"queue_active_key"` + InflightKeyPrefix string `mapstructure:"inflight_key_prefix"` + LockKeyPrefix string `mapstructure:"lock_key_prefix"` + IdempotencyKeyPrefix string `mapstructure:"idempotency_key_prefix"` + InflightTTLSeconds int `mapstructure:"inflight_ttl_seconds"` + JobLockTTLSeconds int `mapstructure:"job_lock_ttl_seconds"` + DefaultRequeueDelaySeconds int `mapstructure:"default_requeue_delay_seconds"` + ErrorRetryDelaySeconds int `mapstructure:"error_retry_delay_seconds"` + LockConflictDelaySeconds int `mapstructure:"lock_conflict_delay_seconds"` + StaleActiveAfterSeconds int `mapstructure:"stale_active_after_seconds"` + DelayedMoverIntervalSeconds int `mapstructure:"delayed_mover_interval_seconds"` + RecoveryIntervalSeconds int `mapstructure:"recovery_interval_seconds"` + DelayedMoveLimit int `mapstructure:"delayed_move_limit"` + RecoverLimit int `mapstructure:"recover_limit"` + VertexEnabled bool `mapstructure:"vertex_enabled"` + VertexProjectID string `mapstructure:"vertex_project_id"` + VertexLocation string `mapstructure:"vertex_location"` + // VertexManagedGCSBucket is a server-owned bucket for batch JSONL input/output. + // Disable Cloud Storage soft delete on this bucket to avoid retaining deleted batch objects. + VertexManagedGCSBucket string `mapstructure:"vertex_managed_gcs_bucket"` + VertexManagedGCSPrefix string `mapstructure:"vertex_managed_gcs_prefix"` + VertexInputRetentionHours int `mapstructure:"vertex_input_retention_hours"` + VertexOutputRetentionHours int `mapstructure:"vertex_output_retention_hours"` + VertexBatchPredictionBaseURL string `mapstructure:"vertex_batch_prediction_base_url"` + VertexGCSBaseURL string `mapstructure:"vertex_gcs_base_url"` +} + type LinuxDoConnectConfig struct { Enabled bool `mapstructure:"enabled"` ClientID string `mapstructure:"client_id"` @@ -1732,6 +1783,53 @@ func setDefaults() { viper.SetDefault("redis.min_idle_conns", 128) viper.SetDefault("redis.enable_tls", false) + // Batch Image queue + viper.SetDefault("batch_image.enabled", false) + viper.SetDefault("batch_image.max_items_per_job_default", 200) + viper.SetDefault("batch_image.max_items_per_job_trial", 50) + viper.SetDefault("batch_image.max_output_images_per_job", 200) + viper.SetDefault("batch_image.max_output_images_per_item", 4) + viper.SetDefault("batch_image.max_prompt_chars_per_item", 8000) + viper.SetDefault("batch_image.max_reference_images_per_job", 1000) + viper.SetDefault("batch_image.max_reference_inline_bytes_per_job", 134217728) + viper.SetDefault("batch_image.default_response_mime_type", "image/png") + viper.SetDefault("batch_image.default_image_size", "1K") + viper.SetDefault("batch_image.max_download_items_zip", 200) + viper.SetDefault("batch_image.max_download_bytes_per_request", 536870912) + viper.SetDefault("batch_image.max_download_duration_seconds", 600) + viper.SetDefault("batch_image.max_download_concurrency_per_user", 1) + viper.SetDefault("batch_image.input_retention_after_terminal_hours", 24) + viper.SetDefault("batch_image.output_retention_after_terminal_hours", 72) + viper.SetDefault("batch_image.output_retention_max_days", 7) + viper.SetDefault("batch_image.cleanup_interval_minutes", 30) + viper.SetDefault("batch_image.cleanup_batch_size", 100) + viper.SetDefault("batch_image.queue_enabled", false) + viper.SetDefault("batch_image.queue_ready_key", "batch_image:queue:ready") + viper.SetDefault("batch_image.queue_delayed_key", "batch_image:queue:delayed") + viper.SetDefault("batch_image.queue_active_key", "batch_image:queue:active") + viper.SetDefault("batch_image.inflight_key_prefix", "batch_image:queue:inflight:") + viper.SetDefault("batch_image.lock_key_prefix", "batch_image:queue:lock:") + viper.SetDefault("batch_image.idempotency_key_prefix", "batch_image:queue:idem:") + viper.SetDefault("batch_image.inflight_ttl_seconds", 604800) + viper.SetDefault("batch_image.job_lock_ttl_seconds", 300) + viper.SetDefault("batch_image.default_requeue_delay_seconds", 30) + viper.SetDefault("batch_image.error_retry_delay_seconds", 60) + viper.SetDefault("batch_image.lock_conflict_delay_seconds", 5) + viper.SetDefault("batch_image.stale_active_after_seconds", 600) + viper.SetDefault("batch_image.delayed_mover_interval_seconds", 5) + viper.SetDefault("batch_image.recovery_interval_seconds", 300) + viper.SetDefault("batch_image.delayed_move_limit", 100) + viper.SetDefault("batch_image.recover_limit", 100) + viper.SetDefault("batch_image.vertex_enabled", false) + viper.SetDefault("batch_image.vertex_project_id", "") + viper.SetDefault("batch_image.vertex_location", "global") + viper.SetDefault("batch_image.vertex_managed_gcs_bucket", "") + viper.SetDefault("batch_image.vertex_managed_gcs_prefix", "batch-image/{env}/{batch_id}") + viper.SetDefault("batch_image.vertex_input_retention_hours", 24) + viper.SetDefault("batch_image.vertex_output_retention_hours", 72) + viper.SetDefault("batch_image.vertex_batch_prediction_base_url", "") + viper.SetDefault("batch_image.vertex_gcs_base_url", "") + // Ops (vNext) viper.SetDefault("ops.enabled", true) viper.SetDefault("ops.use_preaggregated_tables", true) @@ -2333,6 +2431,61 @@ func (c *Config) Validate() error { if c.Redis.MinIdleConns > c.Redis.PoolSize { return fmt.Errorf("redis.min_idle_conns cannot exceed redis.pool_size") } + if c.BatchImage.QueueEnabled { + if strings.TrimSpace(c.BatchImage.QueueReadyKey) == "" { + return fmt.Errorf("batch_image.queue_ready_key must not be empty") + } + if strings.TrimSpace(c.BatchImage.QueueDelayedKey) == "" { + return fmt.Errorf("batch_image.queue_delayed_key must not be empty") + } + if strings.TrimSpace(c.BatchImage.QueueActiveKey) == "" { + return fmt.Errorf("batch_image.queue_active_key must not be empty") + } + if strings.TrimSpace(c.BatchImage.InflightKeyPrefix) == "" { + return fmt.Errorf("batch_image.inflight_key_prefix must not be empty") + } + if strings.TrimSpace(c.BatchImage.LockKeyPrefix) == "" { + return fmt.Errorf("batch_image.lock_key_prefix must not be empty") + } + if c.BatchImage.InflightTTLSeconds <= 0 { + return fmt.Errorf("batch_image.inflight_ttl_seconds must be positive") + } + if c.BatchImage.JobLockTTLSeconds <= 0 { + return fmt.Errorf("batch_image.job_lock_ttl_seconds must be positive") + } + if c.BatchImage.StaleActiveAfterSeconds <= 0 { + return fmt.Errorf("batch_image.stale_active_after_seconds must be positive") + } + if c.BatchImage.DelayedMoveLimit <= 0 { + return fmt.Errorf("batch_image.delayed_move_limit must be positive") + } + if c.BatchImage.RecoverLimit <= 0 { + return fmt.Errorf("batch_image.recover_limit must be positive") + } + } + if c.BatchImage.VertexEnabled { + if strings.TrimSpace(c.BatchImage.VertexManagedGCSBucket) == "" { + return fmt.Errorf("batch_image.vertex_managed_gcs_bucket must not be empty when vertex is enabled") + } + if strings.Contains(c.BatchImage.VertexManagedGCSBucket, "://") { + return fmt.Errorf("batch_image.vertex_managed_gcs_bucket must be a bucket name, not a URI") + } + if strings.TrimSpace(c.BatchImage.VertexLocation) == "" { + return fmt.Errorf("batch_image.vertex_location must not be empty when vertex is enabled") + } + if strings.TrimSpace(c.BatchImage.VertexManagedGCSPrefix) == "" { + return fmt.Errorf("batch_image.vertex_managed_gcs_prefix must not be empty when vertex is enabled") + } + if !strings.Contains(c.BatchImage.VertexManagedGCSPrefix, "{batch_id}") { + return fmt.Errorf("batch_image.vertex_managed_gcs_prefix must contain {batch_id}") + } + if c.BatchImage.VertexInputRetentionHours <= 0 { + return fmt.Errorf("batch_image.vertex_input_retention_hours must be positive") + } + if c.BatchImage.VertexOutputRetentionHours <= 0 { + return fmt.Errorf("batch_image.vertex_output_retention_hours must be positive") + } + } if c.Dashboard.Enabled { if c.Dashboard.StatsFreshTTLSeconds <= 0 { return fmt.Errorf("dashboard_cache.stats_fresh_ttl_seconds must be positive") diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 804155d1a9..32aff543af 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -270,6 +270,14 @@ func TestLoadDefaultIdempotencyConfig(t *testing.T) { } } +func TestLoadDefaultBatchImageQueueDisabled(t *testing.T) { + resetViperWithJWTSecret(t) + + cfg, err := Load() + require.NoError(t, err) + require.False(t, cfg.BatchImage.QueueEnabled) +} + func TestLoadIdempotencyConfigFromEnv(t *testing.T) { resetViperWithJWTSecret(t) t.Setenv("IDEMPOTENCY_OBSERVE_ONLY", "false") diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 8c91245fbf..a4b0773999 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -485,6 +485,8 @@ func (h *AccountHandler) List(c *gin.Context) { search = search[:100] } lite := parseBoolQueryWithDefault(c.Query("lite"), false) + // 调度分需要跨候选池批量打分并读取负载,默认列表不计算;只有前端列可见时才显式开启。 + includeSchedulerScore := parseBoolQueryWithDefault(c.Query("include_scheduler_score"), false) var groupID int64 if groupIDStr := c.Query("group"); groupIDStr != "" { @@ -520,7 +522,7 @@ func (h *AccountHandler) List(c *gin.Context) { var windowCosts map[int64]float64 var activeSessions map[int64]int var rpmCounts map[int64]int - // 仅当前页存在 OpenAI 账号时才计算调度分数,避免为空结果付出池查询开销。 + // 双重门控:用户要看该列,且当前页确实有 OpenAI 账号,才进入昂贵的候选池打分路径。 var schedulerScores map[int64]*AccountSchedulerScore var schedulerGroupScores map[int64][]AccountSchedulerGroupScore pageHasOpenAIAccounts := false @@ -530,7 +532,7 @@ func (h *AccountHandler) List(c *gin.Context) { break } } - if pageHasOpenAIAccounts { + if includeSchedulerScore && pageHasOpenAIAccounts { schedulerFilterPool := h.listAccountSchedulerScoreFilterPool(c.Request.Context(), platform, accountType, status, search, groupID, privacyMode) schedulerScores, schedulerGroupScores = h.buildOpenAIAccountSchedulerScores(c.Request.Context(), accounts, schedulerFilterPool) } diff --git a/backend/internal/handler/admin/account_handler_list_test.go b/backend/internal/handler/admin/account_handler_list_test.go index 29e36ad865..4b1bd7224c 100644 --- a/backend/internal/handler/admin/account_handler_list_test.go +++ b/backend/internal/handler/admin/account_handler_list_test.go @@ -92,7 +92,7 @@ func TestAccountHandlerListReturnsSchedulerScoresPerGroup(t *testing.T) { } rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=20&platform=openai", nil) + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=20&platform=openai&include_scheduler_score=1", nil) router.ServeHTTP(rec, req) require.Equal(t, http.StatusOK, rec.Code) @@ -147,6 +147,43 @@ func TestAccountHandlerListReturnsSchedulerScoresPerGroup(t *testing.T) { require.Greater(t, high.SchedulerScores[0].BaseScore, low.SchedulerScores[0].BaseScore) } +func TestAccountHandlerListSkipsSchedulerScoresByDefault(t *testing.T) { + router, adminSvc := setupAccountListRouter() + now := time.Now().UTC() + adminSvc.accounts = []service.Account{ + { + ID: 110, + Name: "openai-account", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 1, + CreatedAt: now, + UpdatedAt: now, + }, + } + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=20&platform=openai", nil) + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Zero(t, adminSvc.schedulerScoreFilterCalls) + require.Zero(t, adminSvc.openAISchedulerScorePoolCalls) + + var payload struct { + Data struct { + Items []map[string]any `json:"items"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload)) + require.Len(t, payload.Data.Items, 1) + require.NotContains(t, payload.Data.Items[0], "scheduler_score") + require.NotContains(t, payload.Data.Items[0], "scheduler_scores") +} + func TestAccountHandlerListKeepsSchedulerScoreScopedToFilter(t *testing.T) { router, adminSvc := setupAccountListRouter() now := time.Now().UTC() @@ -188,7 +225,7 @@ func TestAccountHandlerListKeepsSchedulerScoreScopedToFilter(t *testing.T) { adminSvc.openAISchedulerScorePoolAccounts = []service.Account{visibleAccount, hiddenGroupPeer} rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=1&platform=openai", nil) + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=1&platform=openai&include_scheduler_score=1", nil) router.ServeHTTP(rec, req) require.Equal(t, http.StatusOK, rec.Code) @@ -246,7 +283,7 @@ func TestAccountHandlerListSchedulerScoreIgnoresPagination(t *testing.T) { adminSvc.accountSchedulerScoreFilterAccounts = []service.Account{visibleAccount, hiddenFilterPeer} rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=1&platform=openai", nil) + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts?page=1&page_size=1&platform=openai&include_scheduler_score=1", nil) router.ServeHTTP(rec, req) require.Equal(t, http.StatusOK, rec.Code) diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index 187925e33a..7a7cbb473e 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -16,6 +16,8 @@ type stubAdminService struct { accounts []service.Account accountSchedulerScoreFilterAccounts []service.Account openAISchedulerScorePoolAccounts []service.Account + schedulerScoreFilterCalls int + openAISchedulerScorePoolCalls int proxies []service.Proxy proxyCounts []service.ProxyWithAccountCount redeems []service.RedeemCode @@ -351,6 +353,7 @@ func (s *stubAdminService) ListAccounts(ctx context.Context, page, pageSize int, } func (s *stubAdminService) ListAccountsForSchedulerScoreFilter(_ context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, error) { + s.schedulerScoreFilterCalls++ if s.accountSchedulerScoreFilterAccounts != nil { return s.accountSchedulerScoreFilterAccounts, nil } @@ -358,6 +361,7 @@ func (s *stubAdminService) ListAccountsForSchedulerScoreFilter(_ context.Context } func (s *stubAdminService) ListOpenAISchedulableAccountsForSchedulerScore(_ context.Context, groupID *int64) ([]service.Account, error) { + s.openAISchedulerScorePoolCalls++ accounts := s.openAISchedulerScorePoolAccounts if accounts == nil { accounts = s.accounts diff --git a/backend/internal/handler/admin/dashboard_handler.go b/backend/internal/handler/admin/dashboard_handler.go index e9fbb630fa..b42b395d33 100644 --- a/backend/internal/handler/admin/dashboard_handler.go +++ b/backend/internal/handler/admin/dashboard_handler.go @@ -675,6 +675,9 @@ func (h *DashboardHandler) GetUserBreakdown(c *gin.Context) { } } + // sort_by 由 repo 层 allowlist 校验;非法值静默回退默认排序(actual_cost)。 + dim.SortBy = strings.TrimSpace(c.Query("sort_by")) + limit := 50 if v := c.Query("limit"); v != "" { if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 200 { diff --git a/backend/internal/handler/admin/dashboard_handler_user_breakdown_test.go b/backend/internal/handler/admin/dashboard_handler_user_breakdown_test.go index b3a05111b3..3065eee3c4 100644 --- a/backend/internal/handler/admin/dashboard_handler_user_breakdown_test.go +++ b/backend/internal/handler/admin/dashboard_handler_user_breakdown_test.go @@ -59,7 +59,21 @@ func TestGetUserBreakdown_GroupIDFilter(t *testing.T) { require.Equal(t, int64(42), repo.capturedDim.GroupID) require.Empty(t, repo.capturedDim.Model) require.Empty(t, repo.capturedDim.Endpoint) - require.Equal(t, 50, repo.capturedLimit) // default limit + require.Equal(t, 50, repo.capturedLimit) // default limit + require.Empty(t, repo.capturedDim.SortBy) // no sort_by => empty (repo falls back to default) +} + +func TestGetUserBreakdown_SortBy(t *testing.T) { + repo := &userBreakdownRepoCapture{} + router := newUserBreakdownRouter(repo) + + req := httptest.NewRequest(http.MethodGet, + "/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&sort_by=total_tokens", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, "total_tokens", repo.capturedDim.SortBy) } func TestGetUserBreakdown_ModelFilter(t *testing.T) { diff --git a/backend/internal/handler/admin/group_handler.go b/backend/internal/handler/admin/group_handler.go index 0a98ad6784..56a0b29ed0 100644 --- a/backend/internal/handler/admin/group_handler.go +++ b/backend/internal/handler/admin/group_handler.go @@ -93,8 +93,13 @@ type CreateGroupRequest struct { MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"` // 图片生成计费配置(antigravity 和 gemini 平台使用,负数表示清除配置) AllowImageGeneration bool `json:"allow_image_generation"` + AllowBatchImageGeneration bool `json:"allow_batch_image_generation"` ImageRateIndependent bool `json:"image_rate_independent"` ImageRateMultiplier *float64 `json:"image_rate_multiplier"` + BatchImageDiscountMultiplier *float64 `json:"batch_image_discount_multiplier"` + BatchImageHoldMultiplier *float64 `json:"batch_image_hold_multiplier"` + VideoRateIndependent bool `json:"video_rate_independent"` + VideoRateMultiplier *float64 `json:"video_rate_multiplier"` PeakRateEnabled bool `json:"peak_rate_enabled"` PeakStart string `json:"peak_start"` PeakEnd string `json:"peak_end"` @@ -102,6 +107,9 @@ type CreateGroupRequest struct { ImagePrice1K *float64 `json:"image_price_1k"` ImagePrice2K *float64 `json:"image_price_2k"` ImagePrice4K *float64 `json:"image_price_4k"` + VideoPrice480P *float64 `json:"video_price_480p"` + VideoPrice720P *float64 `json:"video_price_720p"` + VideoPrice1080P *float64 `json:"video_price_1080p"` ClaudeCodeOnly bool `json:"claude_code_only"` FallbackGroupID *int64 `json:"fallback_group_id"` FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request"` @@ -138,8 +146,13 @@ type UpdateGroupRequest struct { MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"` // 图片生成计费配置(antigravity 和 gemini 平台使用,负数表示清除配置) AllowImageGeneration *bool `json:"allow_image_generation"` + AllowBatchImageGeneration *bool `json:"allow_batch_image_generation"` ImageRateIndependent *bool `json:"image_rate_independent"` ImageRateMultiplier *float64 `json:"image_rate_multiplier"` + BatchImageDiscountMultiplier *float64 `json:"batch_image_discount_multiplier"` + BatchImageHoldMultiplier *float64 `json:"batch_image_hold_multiplier"` + VideoRateIndependent *bool `json:"video_rate_independent"` + VideoRateMultiplier *float64 `json:"video_rate_multiplier"` PeakRateEnabled *bool `json:"peak_rate_enabled"` PeakStart *string `json:"peak_start"` PeakEnd *string `json:"peak_end"` @@ -147,6 +160,9 @@ type UpdateGroupRequest struct { ImagePrice1K *float64 `json:"image_price_1k"` ImagePrice2K *float64 `json:"image_price_2k"` ImagePrice4K *float64 `json:"image_price_4k"` + VideoPrice480P *float64 `json:"video_price_480p"` + VideoPrice720P *float64 `json:"video_price_720p"` + VideoPrice1080P *float64 `json:"video_price_1080p"` ClaudeCodeOnly *bool `json:"claude_code_only"` FallbackGroupID *int64 `json:"fallback_group_id"` FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request"` @@ -301,8 +317,13 @@ func (h *GroupHandler) Create(c *gin.Context) { WeeklyLimitUSD: req.WeeklyLimitUSD.ToServiceInput(), MonthlyLimitUSD: req.MonthlyLimitUSD.ToServiceInput(), AllowImageGeneration: req.AllowImageGeneration, + AllowBatchImageGeneration: req.AllowBatchImageGeneration, ImageRateIndependent: req.ImageRateIndependent, ImageRateMultiplier: req.ImageRateMultiplier, + BatchImageDiscountMultiplier: req.BatchImageDiscountMultiplier, + BatchImageHoldMultiplier: req.BatchImageHoldMultiplier, + VideoRateIndependent: req.VideoRateIndependent, + VideoRateMultiplier: req.VideoRateMultiplier, PeakRateEnabled: req.PeakRateEnabled, PeakStart: req.PeakStart, PeakEnd: req.PeakEnd, @@ -310,6 +331,9 @@ func (h *GroupHandler) Create(c *gin.Context) { ImagePrice1K: req.ImagePrice1K, ImagePrice2K: req.ImagePrice2K, ImagePrice4K: req.ImagePrice4K, + VideoPrice480P: req.VideoPrice480P, + VideoPrice720P: req.VideoPrice720P, + VideoPrice1080P: req.VideoPrice1080P, ClaudeCodeOnly: req.ClaudeCodeOnly, FallbackGroupID: req.FallbackGroupID, FallbackGroupIDOnInvalidRequest: req.FallbackGroupIDOnInvalidRequest, @@ -361,8 +385,13 @@ func (h *GroupHandler) Update(c *gin.Context) { WeeklyLimitUSD: req.WeeklyLimitUSD.ToServiceInput(), MonthlyLimitUSD: req.MonthlyLimitUSD.ToServiceInput(), AllowImageGeneration: req.AllowImageGeneration, + AllowBatchImageGeneration: req.AllowBatchImageGeneration, ImageRateIndependent: req.ImageRateIndependent, ImageRateMultiplier: req.ImageRateMultiplier, + BatchImageDiscountMultiplier: req.BatchImageDiscountMultiplier, + BatchImageHoldMultiplier: req.BatchImageHoldMultiplier, + VideoRateIndependent: req.VideoRateIndependent, + VideoRateMultiplier: req.VideoRateMultiplier, PeakRateEnabled: req.PeakRateEnabled, PeakStart: req.PeakStart, PeakEnd: req.PeakEnd, @@ -370,6 +399,9 @@ func (h *GroupHandler) Update(c *gin.Context) { ImagePrice1K: req.ImagePrice1K, ImagePrice2K: req.ImagePrice2K, ImagePrice4K: req.ImagePrice4K, + VideoPrice480P: req.VideoPrice480P, + VideoPrice720P: req.VideoPrice720P, + VideoPrice1080P: req.VideoPrice1080P, ClaudeCodeOnly: req.ClaudeCodeOnly, FallbackGroupID: req.FallbackGroupID, FallbackGroupIDOnInvalidRequest: req.FallbackGroupIDOnInvalidRequest, diff --git a/backend/internal/handler/admin/setting_handler.go b/backend/internal/handler/admin/setting_handler.go index 529e46c575..1d506f505b 100644 --- a/backend/internal/handler/admin/setting_handler.go +++ b/backend/internal/handler/admin/setting_handler.go @@ -1,21 +1,16 @@ package admin import ( - "context" "crypto/rand" "encoding/hex" "encoding/json" - "errors" "fmt" "log/slog" - "net/http" "regexp" "strings" - "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/handler/dto" "github.com/Wei-Shaw/sub2api/internal/pkg/response" - "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" @@ -415,2556 +410,6 @@ func loginAgreementDocumentsToService(items []dto.LoginAgreementDocument) []serv return result } -// UpdateSettingsRequest 更新设置请求 -type UpdateSettingsRequest struct { - // 注册设置 - RegistrationEnabled bool `json:"registration_enabled"` - EmailVerifyEnabled bool `json:"email_verify_enabled"` - RegistrationEmailSuffixWhitelist []string `json:"registration_email_suffix_whitelist"` - PromoCodeEnabled bool `json:"promo_code_enabled"` - PasswordResetEnabled bool `json:"password_reset_enabled"` - FrontendURL string `json:"frontend_url"` - InvitationCodeEnabled bool `json:"invitation_code_enabled"` - TotpEnabled bool `json:"totp_enabled"` // TOTP 双因素认证 - LoginAgreementEnabled bool `json:"login_agreement_enabled"` - LoginAgreementMode string `json:"login_agreement_mode"` - LoginAgreementUpdatedAt string `json:"login_agreement_updated_at"` - LoginAgreementDocuments []dto.LoginAgreementDocument `json:"login_agreement_documents"` - - // 邮件服务设置 - SMTPHost string `json:"smtp_host"` - SMTPPort int `json:"smtp_port"` - SMTPUsername string `json:"smtp_username"` - SMTPPassword string `json:"smtp_password"` - SMTPFrom string `json:"smtp_from_email"` - SMTPFromName string `json:"smtp_from_name"` - SMTPUseTLS bool `json:"smtp_use_tls"` - - // Cloudflare Turnstile 设置 - TurnstileEnabled bool `json:"turnstile_enabled"` - TurnstileSiteKey string `json:"turnstile_site_key"` - TurnstileSecretKey string `json:"turnstile_secret_key"` - - // API Key IP 访问控制设置 - APIKeyACLTrustForwardedIP *bool `json:"api_key_acl_trust_forwarded_ip"` - - // LinuxDo Connect OAuth 登录 - LinuxDoConnectEnabled bool `json:"linuxdo_connect_enabled"` - LinuxDoConnectClientID string `json:"linuxdo_connect_client_id"` - LinuxDoConnectClientSecret string `json:"linuxdo_connect_client_secret"` - LinuxDoConnectRedirectURL string `json:"linuxdo_connect_redirect_url"` - - // DingTalk Connect OAuth 登录 - DingTalkConnectEnabled bool `json:"dingtalk_connect_enabled"` - DingTalkConnectClientID string `json:"dingtalk_connect_client_id"` - DingTalkConnectClientSecret string `json:"dingtalk_connect_client_secret"` - DingTalkConnectRedirectURL string `json:"dingtalk_connect_redirect_url"` - DingTalkConnectCorpRestrictionPolicy string `json:"dingtalk_connect_corp_restriction_policy"` - DingTalkConnectInternalCorpID string `json:"dingtalk_connect_internal_corp_id"` - DingTalkConnectBypassRegistration bool `json:"dingtalk_connect_bypass_registration"` - DingTalkConnectSyncCorpEmail bool `json:"dingtalk_connect_sync_corp_email"` - DingTalkConnectSyncDisplayName bool `json:"dingtalk_connect_sync_display_name"` - DingTalkConnectSyncDept bool `json:"dingtalk_connect_sync_dept"` - DingTalkConnectSyncCorpEmailAttrKey string `json:"dingtalk_connect_sync_corp_email_attr_key"` - DingTalkConnectSyncDisplayNameAttrKey string `json:"dingtalk_connect_sync_display_name_attr_key"` - DingTalkConnectSyncDeptAttrKey string `json:"dingtalk_connect_sync_dept_attr_key"` - DingTalkConnectSyncCorpEmailAttrName string `json:"dingtalk_connect_sync_corp_email_attr_name"` - DingTalkConnectSyncDisplayNameAttrName string `json:"dingtalk_connect_sync_display_name_attr_name"` - DingTalkConnectSyncDeptAttrName string `json:"dingtalk_connect_sync_dept_attr_name"` - - // WeChat Connect OAuth 登录 - WeChatConnectEnabled bool `json:"wechat_connect_enabled"` - WeChatConnectAppID string `json:"wechat_connect_app_id"` - WeChatConnectAppSecret string `json:"wechat_connect_app_secret"` - WeChatConnectOpenAppID string `json:"wechat_connect_open_app_id"` - WeChatConnectOpenAppSecret string `json:"wechat_connect_open_app_secret"` - WeChatConnectMPAppID string `json:"wechat_connect_mp_app_id"` - WeChatConnectMPAppSecret string `json:"wechat_connect_mp_app_secret"` - WeChatConnectMobileAppID string `json:"wechat_connect_mobile_app_id"` - WeChatConnectMobileAppSecret string `json:"wechat_connect_mobile_app_secret"` - WeChatConnectOpenEnabled bool `json:"wechat_connect_open_enabled"` - WeChatConnectMPEnabled bool `json:"wechat_connect_mp_enabled"` - WeChatConnectMobileEnabled bool `json:"wechat_connect_mobile_enabled"` - WeChatConnectMode string `json:"wechat_connect_mode"` - WeChatConnectScopes string `json:"wechat_connect_scopes"` - WeChatConnectRedirectURL string `json:"wechat_connect_redirect_url"` - WeChatConnectFrontendRedirectURL string `json:"wechat_connect_frontend_redirect_url"` - - // Generic OIDC OAuth 登录 - OIDCConnectEnabled bool `json:"oidc_connect_enabled"` - OIDCConnectProviderName string `json:"oidc_connect_provider_name"` - OIDCConnectClientID string `json:"oidc_connect_client_id"` - OIDCConnectClientSecret string `json:"oidc_connect_client_secret"` - OIDCConnectIssuerURL string `json:"oidc_connect_issuer_url"` - OIDCConnectDiscoveryURL string `json:"oidc_connect_discovery_url"` - OIDCConnectAuthorizeURL string `json:"oidc_connect_authorize_url"` - OIDCConnectTokenURL string `json:"oidc_connect_token_url"` - OIDCConnectUserInfoURL string `json:"oidc_connect_userinfo_url"` - OIDCConnectJWKSURL string `json:"oidc_connect_jwks_url"` - OIDCConnectScopes string `json:"oidc_connect_scopes"` - OIDCConnectRedirectURL string `json:"oidc_connect_redirect_url"` - OIDCConnectFrontendRedirectURL string `json:"oidc_connect_frontend_redirect_url"` - OIDCConnectTokenAuthMethod string `json:"oidc_connect_token_auth_method"` - OIDCConnectUsePKCE *bool `json:"oidc_connect_use_pkce"` - OIDCConnectValidateIDToken *bool `json:"oidc_connect_validate_id_token"` - OIDCConnectAllowedSigningAlgs string `json:"oidc_connect_allowed_signing_algs"` - OIDCConnectClockSkewSeconds int `json:"oidc_connect_clock_skew_seconds"` - OIDCConnectRequireEmailVerified bool `json:"oidc_connect_require_email_verified"` - OIDCConnectUserInfoEmailPath string `json:"oidc_connect_userinfo_email_path"` - OIDCConnectUserInfoIDPath string `json:"oidc_connect_userinfo_id_path"` - OIDCConnectUserInfoUsernamePath string `json:"oidc_connect_userinfo_username_path"` - - GitHubOAuthEnabled bool `json:"github_oauth_enabled"` - GitHubOAuthClientID string `json:"github_oauth_client_id"` - GitHubOAuthClientSecret string `json:"github_oauth_client_secret"` - GitHubOAuthRedirectURL string `json:"github_oauth_redirect_url"` - GitHubOAuthFrontendRedirectURL string `json:"github_oauth_frontend_redirect_url"` - GoogleOAuthEnabled bool `json:"google_oauth_enabled"` - GoogleOAuthClientID string `json:"google_oauth_client_id"` - GoogleOAuthClientSecret string `json:"google_oauth_client_secret"` - GoogleOAuthRedirectURL string `json:"google_oauth_redirect_url"` - GoogleOAuthFrontendRedirectURL string `json:"google_oauth_frontend_redirect_url"` - - // OEM设置 - SiteName string `json:"site_name"` - SiteLogo string `json:"site_logo"` - SiteSubtitle string `json:"site_subtitle"` - APIBaseURL string `json:"api_base_url"` - ContactInfo string `json:"contact_info"` - DocURL string `json:"doc_url"` - HomeContent string `json:"home_content"` - HideCcsImportButton bool `json:"hide_ccs_import_button"` - PurchaseSubscriptionEnabled *bool `json:"purchase_subscription_enabled"` - PurchaseSubscriptionURL *string `json:"purchase_subscription_url"` - TableDefaultPageSize int `json:"table_default_page_size"` - TablePageSizeOptions []int `json:"table_page_size_options"` - CustomMenuItems *[]dto.CustomMenuItem `json:"custom_menu_items"` - CustomEndpoints *[]dto.CustomEndpoint `json:"custom_endpoints"` - - // 默认配置 - DefaultConcurrency int `json:"default_concurrency"` - DefaultBalance float64 `json:"default_balance"` - AffiliateRebateRate *float64 `json:"affiliate_rebate_rate"` - AffiliateRebateFreezeHours *int `json:"affiliate_rebate_freeze_hours"` - AffiliateRebateDurationDays *int `json:"affiliate_rebate_duration_days"` - AffiliateRebatePerInviteeCap *float64 `json:"affiliate_rebate_per_invitee_cap"` - DefaultUserRPMLimit int `json:"default_user_rpm_limit"` - DefaultSubscriptions []dto.DefaultSubscriptionSetting `json:"default_subscriptions"` - AuthSourceDefaultEmailBalance *float64 `json:"auth_source_default_email_balance"` - AuthSourceDefaultEmailConcurrency *int `json:"auth_source_default_email_concurrency"` - AuthSourceDefaultEmailSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_email_subscriptions"` - AuthSourceDefaultEmailGrantOnSignup *bool `json:"auth_source_default_email_grant_on_signup"` - AuthSourceDefaultEmailGrantOnFirstBind *bool `json:"auth_source_default_email_grant_on_first_bind"` - AuthSourceDefaultLinuxDoBalance *float64 `json:"auth_source_default_linuxdo_balance"` - AuthSourceDefaultLinuxDoConcurrency *int `json:"auth_source_default_linuxdo_concurrency"` - AuthSourceDefaultLinuxDoSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_linuxdo_subscriptions"` - AuthSourceDefaultLinuxDoGrantOnSignup *bool `json:"auth_source_default_linuxdo_grant_on_signup"` - AuthSourceDefaultLinuxDoGrantOnFirstBind *bool `json:"auth_source_default_linuxdo_grant_on_first_bind"` - AuthSourceDefaultOIDCBalance *float64 `json:"auth_source_default_oidc_balance"` - AuthSourceDefaultOIDCConcurrency *int `json:"auth_source_default_oidc_concurrency"` - AuthSourceDefaultOIDCSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_oidc_subscriptions"` - AuthSourceDefaultOIDCGrantOnSignup *bool `json:"auth_source_default_oidc_grant_on_signup"` - AuthSourceDefaultOIDCGrantOnFirstBind *bool `json:"auth_source_default_oidc_grant_on_first_bind"` - AuthSourceDefaultWeChatBalance *float64 `json:"auth_source_default_wechat_balance"` - AuthSourceDefaultWeChatConcurrency *int `json:"auth_source_default_wechat_concurrency"` - AuthSourceDefaultWeChatSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_wechat_subscriptions"` - AuthSourceDefaultWeChatGrantOnSignup *bool `json:"auth_source_default_wechat_grant_on_signup"` - AuthSourceDefaultWeChatGrantOnFirstBind *bool `json:"auth_source_default_wechat_grant_on_first_bind"` - AuthSourceDefaultGitHubBalance *float64 `json:"auth_source_default_github_balance"` - AuthSourceDefaultGitHubConcurrency *int `json:"auth_source_default_github_concurrency"` - AuthSourceDefaultGitHubSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_github_subscriptions"` - AuthSourceDefaultGitHubGrantOnSignup *bool `json:"auth_source_default_github_grant_on_signup"` - AuthSourceDefaultGitHubGrantOnFirstBind *bool `json:"auth_source_default_github_grant_on_first_bind"` - AuthSourceDefaultGoogleBalance *float64 `json:"auth_source_default_google_balance"` - AuthSourceDefaultGoogleConcurrency *int `json:"auth_source_default_google_concurrency"` - AuthSourceDefaultGoogleSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_google_subscriptions"` - AuthSourceDefaultGoogleGrantOnSignup *bool `json:"auth_source_default_google_grant_on_signup"` - AuthSourceDefaultGoogleGrantOnFirstBind *bool `json:"auth_source_default_google_grant_on_first_bind"` - AuthSourceDefaultDingTalkBalance *float64 `json:"auth_source_default_dingtalk_balance"` - AuthSourceDefaultDingTalkConcurrency *int `json:"auth_source_default_dingtalk_concurrency"` - AuthSourceDefaultDingTalkSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_dingtalk_subscriptions"` - AuthSourceDefaultDingTalkGrantOnSignup *bool `json:"auth_source_default_dingtalk_grant_on_signup"` - AuthSourceDefaultDingTalkGrantOnFirstBind *bool `json:"auth_source_default_dingtalk_grant_on_first_bind"` - ForceEmailOnThirdPartySignup *bool `json:"force_email_on_third_party_signup"` - - // Model fallback configuration - EnableModelFallback bool `json:"enable_model_fallback"` - FallbackModelAnthropic string `json:"fallback_model_anthropic"` - FallbackModelOpenAI string `json:"fallback_model_openai"` - FallbackModelGemini string `json:"fallback_model_gemini"` - FallbackModelAntigravity string `json:"fallback_model_antigravity"` - - // Identity patch configuration (Claude -> Gemini) - EnableIdentityPatch bool `json:"enable_identity_patch"` - IdentityPatchPrompt string `json:"identity_patch_prompt"` - - // Ops monitoring (vNext) - OpsMonitoringEnabled *bool `json:"ops_monitoring_enabled"` - OpsRealtimeMonitoringEnabled *bool `json:"ops_realtime_monitoring_enabled"` - OpsQueryModeDefault *string `json:"ops_query_mode_default"` - OpsMetricsIntervalSeconds *int `json:"ops_metrics_interval_seconds"` - - MinClaudeCodeVersion string `json:"min_claude_code_version"` - MaxClaudeCodeVersion string `json:"max_claude_code_version"` - - // 分组隔离 - AllowUngroupedKeyScheduling bool `json:"allow_ungrouped_key_scheduling"` - - // Backend Mode - BackendModeEnabled bool `json:"backend_mode_enabled"` - - // Gateway forwarding behavior - EnableFingerprintUnification *bool `json:"enable_fingerprint_unification"` - EnableMetadataPassthrough *bool `json:"enable_metadata_passthrough"` - EnableCCHSigning *bool `json:"enable_cch_signing"` - EnableClaudeOAuthSystemPromptInjection *bool `json:"enable_claude_oauth_system_prompt_injection"` - ClaudeOAuthSystemPrompt *string `json:"claude_oauth_system_prompt"` - ClaudeOAuthSystemPromptBlocks *string `json:"claude_oauth_system_prompt_blocks"` - EnableAnthropicCacheTTL1hInjection *bool `json:"enable_anthropic_cache_ttl_1h_injection"` - RewriteMessageCacheControl *bool `json:"rewrite_message_cache_control"` - EnableClientDatelineNormalization *bool `json:"enable_client_dateline_normalization"` - AntigravityUserAgentVersion *string `json:"antigravity_user_agent_version"` - OpenAICodexUserAgent *string `json:"openai_codex_user_agent"` - - // codex_cli_only 加固(global-only) - MinCodexVersion string `json:"min_codex_version"` - MaxCodexVersion string `json:"max_codex_version"` - CodexCLIOnlyBlacklist string `json:"codex_cli_only_blacklist"` - CodexCLIOnlyWhitelist string `json:"codex_cli_only_whitelist"` - CodexCLIOnlyAllowAppServerClients *bool `json:"codex_cli_only_allow_app_server_clients"` - CodexCLIOnlyEngineFingerprintSignals string `json:"codex_cli_only_engine_fingerprint_signals"` - - // Payment visible method routing - PaymentVisibleMethodAlipaySource *string `json:"payment_visible_method_alipay_source"` - PaymentVisibleMethodWxpaySource *string `json:"payment_visible_method_wxpay_source"` - PaymentVisibleMethodAlipayEnabled *bool `json:"payment_visible_method_alipay_enabled"` - PaymentVisibleMethodWxpayEnabled *bool `json:"payment_visible_method_wxpay_enabled"` - - // OpenAI account scheduling - OpenAIAdvancedSchedulerEnabled *bool `json:"openai_advanced_scheduler_enabled"` - OpenAIAdvancedSchedulerStickyWeightedEnabled *bool `json:"openai_advanced_scheduler_sticky_weighted_enabled"` - OpenAIAdvancedSchedulerSubscriptionPriorityEnabled *bool `json:"openai_advanced_scheduler_subscription_priority_enabled"` - OpenAIAdvancedSchedulerLBTopK *string `json:"openai_advanced_scheduler_lb_top_k"` - OpenAIAdvancedSchedulerWeightPriority *string `json:"openai_advanced_scheduler_weight_priority"` - OpenAIAdvancedSchedulerWeightLoad *string `json:"openai_advanced_scheduler_weight_load"` - OpenAIAdvancedSchedulerWeightQueue *string `json:"openai_advanced_scheduler_weight_queue"` - OpenAIAdvancedSchedulerWeightErrorRate *string `json:"openai_advanced_scheduler_weight_error_rate"` - OpenAIAdvancedSchedulerWeightTTFT *string `json:"openai_advanced_scheduler_weight_ttft"` - OpenAIAdvancedSchedulerWeightReset *string `json:"openai_advanced_scheduler_weight_reset"` - OpenAIAdvancedSchedulerWeightQuotaHeadroom *string `json:"openai_advanced_scheduler_weight_quota_headroom"` - OpenAIAdvancedSchedulerWeightPreviousResponse *string `json:"openai_advanced_scheduler_weight_previous_response"` - OpenAIAdvancedSchedulerWeightSessionSticky *string `json:"openai_advanced_scheduler_weight_session_sticky"` - - // 余额不足提醒 - BalanceLowNotifyEnabled *bool `json:"balance_low_notify_enabled"` - BalanceLowNotifyThreshold *float64 `json:"balance_low_notify_threshold"` - BalanceLowNotifyRechargeURL *string `json:"balance_low_notify_recharge_url"` - SubscriptionExpiryNotifyEnabled *bool `json:"subscription_expiry_notify_enabled"` - AccountQuotaNotifyEnabled *bool `json:"account_quota_notify_enabled"` - AccountQuotaNotifyEmails *[]dto.NotifyEmailEntry `json:"account_quota_notify_emails"` - - // Payment configuration (integrated into settings, full replace) - PaymentEnabled *bool `json:"payment_enabled"` - PaymentMinAmount *float64 `json:"payment_min_amount"` - PaymentMaxAmount *float64 `json:"payment_max_amount"` - PaymentDailyLimit *float64 `json:"payment_daily_limit"` - PaymentOrderTimeoutMin *int `json:"payment_order_timeout_minutes"` - PaymentMaxPendingOrders *int `json:"payment_max_pending_orders"` - PaymentEnabledTypes []string `json:"payment_enabled_types"` - PaymentBalanceDisabled *bool `json:"payment_balance_disabled"` - PaymentBalanceRechargeMultiplier *float64 `json:"payment_balance_recharge_multiplier"` - PaymentSubscriptionUSDToCNYRate *float64 `json:"payment_subscription_usd_to_cny_rate"` - PaymentRechargeFeeRate *float64 `json:"payment_recharge_fee_rate"` - PaymentLoadBalanceStrat *string `json:"payment_load_balance_strategy"` - PaymentProductNamePrefix *string `json:"payment_product_name_prefix"` - PaymentProductNameSuffix *string `json:"payment_product_name_suffix"` - PaymentHelpImageURL *string `json:"payment_help_image_url"` - PaymentHelpText *string `json:"payment_help_text"` - - // Cancel rate limit - PaymentCancelRateLimitEnabled *bool `json:"payment_cancel_rate_limit_enabled"` - PaymentCancelRateLimitMax *int `json:"payment_cancel_rate_limit_max"` - PaymentCancelRateLimitWindow *int `json:"payment_cancel_rate_limit_window"` - PaymentCancelRateLimitUnit *string `json:"payment_cancel_rate_limit_unit"` - PaymentCancelRateLimitMode *string `json:"payment_cancel_rate_limit_window_mode"` - - // Force Alipay mobile clients to use QR code payment instead of mobile redirect - PaymentAlipayForceQRCode *bool `json:"payment_alipay_force_qrcode"` - - // Channel Monitor feature switch - ChannelMonitorEnabled *bool `json:"channel_monitor_enabled"` - ChannelMonitorDefaultIntervalSeconds *int `json:"channel_monitor_default_interval_seconds"` - - // Available Channels feature switch (user-facing) - AvailableChannelsEnabled *bool `json:"available_channels_enabled"` - - // Affiliate (邀请返利) feature switch - AffiliateEnabled *bool `json:"affiliate_enabled"` - - // 风控中心功能开关 - RiskControlEnabled *bool `json:"risk_control_enabled"` - - // cyber 会话屏蔽开关 + TTL - CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"` - CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"` - - // OpenAI fast/flex policy (optional, only updated when provided) - OpenAIFastPolicySettings *dto.OpenAIFastPolicySettings `json:"openai_fast_policy_settings,omitempty"` - - // 系统全局 platform quota 默认值(整体替换语义:nil = 不修改,non-nil = 整体覆盖)。 - DefaultPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"default_platform_quotas"` - - // auth-source 层 platform quota 覆盖(override 语义:nil = 不修改,non-nil = 整体覆盖该 source 的 quota 配置)。 - AuthSourceEmailPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_email_platform_quotas"` - AuthSourceLinuxDoPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_linuxdo_platform_quotas"` - AuthSourceOIDCPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_oidc_platform_quotas"` - AuthSourceWeChatPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_wechat_platform_quotas"` - AuthSourceGitHubPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_github_platform_quotas"` - AuthSourceGooglePlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_google_platform_quotas"` - AuthSourceDingTalkPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_dingtalk_platform_quotas"` - - AllowUserViewErrorRequests *bool `json:"allow_user_view_error_requests"` -} - -// UpdateSettings 更新系统设置 -// PUT /api/v1/admin/settings -func (h *SettingHandler) UpdateSettings(c *gin.Context) { - var req UpdateSettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - previousSettings, err := h.settingService.GetAllSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - previousAuthSourceDefaults, err := h.settingService.GetAuthSourceDefaultSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - // 验证参数 - if req.DefaultConcurrency < 1 { - req.DefaultConcurrency = 1 - } - if req.DefaultBalance < 0 { - req.DefaultBalance = 0 - } - affiliateRebateRate := previousSettings.AffiliateRebateRate - if req.AffiliateRebateRate != nil { - affiliateRebateRate = *req.AffiliateRebateRate - } - if affiliateRebateRate < service.AffiliateRebateRateMin { - affiliateRebateRate = service.AffiliateRebateRateMin - } - if affiliateRebateRate > service.AffiliateRebateRateMax { - affiliateRebateRate = service.AffiliateRebateRateMax - } - affiliateRebateFreezeHours := previousSettings.AffiliateRebateFreezeHours - if req.AffiliateRebateFreezeHours != nil { - affiliateRebateFreezeHours = *req.AffiliateRebateFreezeHours - } - if affiliateRebateFreezeHours < 0 { - affiliateRebateFreezeHours = service.AffiliateRebateFreezeHoursDefault - } - if affiliateRebateFreezeHours > service.AffiliateRebateFreezeHoursMax { - affiliateRebateFreezeHours = service.AffiliateRebateFreezeHoursMax - } - affiliateRebateDurationDays := previousSettings.AffiliateRebateDurationDays - if req.AffiliateRebateDurationDays != nil { - affiliateRebateDurationDays = *req.AffiliateRebateDurationDays - } - if affiliateRebateDurationDays < 0 { - affiliateRebateDurationDays = service.AffiliateRebateDurationDaysDefault - } - if affiliateRebateDurationDays > service.AffiliateRebateDurationDaysMax { - affiliateRebateDurationDays = service.AffiliateRebateDurationDaysMax - } - affiliateRebatePerInviteeCap := previousSettings.AffiliateRebatePerInviteeCap - if req.AffiliateRebatePerInviteeCap != nil { - affiliateRebatePerInviteeCap = *req.AffiliateRebatePerInviteeCap - } - if affiliateRebatePerInviteeCap < 0 { - affiliateRebatePerInviteeCap = service.AffiliateRebatePerInviteeCapDefault - } - // 通用表格配置:兼容旧客户端未传字段时保留当前值。 - if req.TableDefaultPageSize <= 0 { - req.TableDefaultPageSize = previousSettings.TableDefaultPageSize - } - if req.TablePageSizeOptions == nil { - req.TablePageSizeOptions = previousSettings.TablePageSizeOptions - } - req.SMTPHost = strings.TrimSpace(req.SMTPHost) - req.SMTPUsername = strings.TrimSpace(req.SMTPUsername) - req.SMTPPassword = strings.TrimSpace(req.SMTPPassword) - req.SMTPFrom = strings.TrimSpace(req.SMTPFrom) - req.SMTPFromName = strings.TrimSpace(req.SMTPFromName) - if req.SMTPPort <= 0 { - req.SMTPPort = 587 - } - req.DefaultSubscriptions = normalizeDefaultSubscriptions(req.DefaultSubscriptions) - req.AuthSourceDefaultEmailSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultEmailSubscriptions) - req.AuthSourceDefaultLinuxDoSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultLinuxDoSubscriptions) - req.AuthSourceDefaultOIDCSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultOIDCSubscriptions) - req.AuthSourceDefaultWeChatSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultWeChatSubscriptions) - req.AuthSourceDefaultDingTalkSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultDingTalkSubscriptions) - - // SMTP 配置保护:如果请求中 smtp_host 为空但数据库中已有配置,则保留已有 SMTP 配置 - // 防止前端加载设置失败时空表单覆盖已保存的 SMTP 配置 - if req.SMTPHost == "" && previousSettings.SMTPHost != "" { - req.SMTPHost = previousSettings.SMTPHost - req.SMTPPort = previousSettings.SMTPPort - req.SMTPUsername = previousSettings.SMTPUsername - req.SMTPFrom = previousSettings.SMTPFrom - req.SMTPFromName = previousSettings.SMTPFromName - req.SMTPUseTLS = previousSettings.SMTPUseTLS - } - - // Turnstile 参数验证 - if req.TurnstileEnabled { - // 检查必填字段 - if req.TurnstileSiteKey == "" { - response.BadRequest(c, "Turnstile Site Key is required when enabled") - return - } - // 如果未提供 secret key,使用已保存的值(留空保留当前值) - if req.TurnstileSecretKey == "" { - if previousSettings.TurnstileSecretKey == "" { - response.BadRequest(c, "Turnstile Secret Key is required when enabled") - return - } - req.TurnstileSecretKey = previousSettings.TurnstileSecretKey - } - - // 当 site_key 或 secret_key 任一变化时验证(避免配置错误导致无法登录) - siteKeyChanged := previousSettings.TurnstileSiteKey != req.TurnstileSiteKey - secretKeyChanged := previousSettings.TurnstileSecretKey != req.TurnstileSecretKey - if siteKeyChanged || secretKeyChanged { - if err := h.turnstileService.ValidateSecretKey(c.Request.Context(), req.TurnstileSecretKey); err != nil { - response.ErrorFrom(c, err) - return - } - } - } - - // TOTP 双因素认证参数验证 - // 只有手动配置了加密密钥才允许启用 TOTP 功能 - if req.TotpEnabled && !previousSettings.TotpEnabled { - // 尝试启用 TOTP,检查加密密钥是否已手动配置 - if !h.settingService.IsTotpEncryptionKeyConfigured() { - response.BadRequest(c, "Cannot enable TOTP: TOTP_ENCRYPTION_KEY environment variable must be configured first. Generate a key with 'openssl rand -hex 32' and set it in your environment.") - return - } - } - loginAgreementMode := strings.ToLower(strings.TrimSpace(req.LoginAgreementMode)) - if loginAgreementMode == "" { - loginAgreementMode = strings.ToLower(strings.TrimSpace(previousSettings.LoginAgreementMode)) - } - switch loginAgreementMode { - case "", "modal": - loginAgreementMode = "modal" - case "checkbox": - default: - response.BadRequest(c, "Login agreement mode must be modal or checkbox") - return - } - loginAgreementUpdatedAt := strings.TrimSpace(req.LoginAgreementUpdatedAt) - if loginAgreementUpdatedAt == "" { - loginAgreementUpdatedAt = strings.TrimSpace(previousSettings.LoginAgreementUpdatedAt) - } - loginAgreementDocuments := loginAgreementDocumentsToService(req.LoginAgreementDocuments) - if len(loginAgreementDocuments) == 0 { - loginAgreementDocuments = previousSettings.LoginAgreementDocuments - } - for _, doc := range loginAgreementDocuments { - if strings.TrimSpace(doc.Title) == "" { - response.BadRequest(c, "Login agreement document title is required") - return - } - if len(doc.Title) > 80 { - response.BadRequest(c, "Login agreement document title is too long (max 80 characters)") - return - } - if len(doc.ContentMD) > 200*1024 { - response.BadRequest(c, "Login agreement document content is too large (max 200KB)") - return - } - } - if req.LoginAgreementEnabled && len(loginAgreementDocuments) == 0 { - response.BadRequest(c, "Login agreement documents are required when enabled") - return - } - - // LinuxDo Connect 参数验证 - if req.LinuxDoConnectEnabled { - req.LinuxDoConnectClientID = strings.TrimSpace(req.LinuxDoConnectClientID) - req.LinuxDoConnectClientSecret = strings.TrimSpace(req.LinuxDoConnectClientSecret) - req.LinuxDoConnectRedirectURL = strings.TrimSpace(req.LinuxDoConnectRedirectURL) - - if req.LinuxDoConnectClientID == "" { - response.BadRequest(c, "LinuxDo Client ID is required when enabled") - return - } - if req.LinuxDoConnectRedirectURL == "" { - response.BadRequest(c, "LinuxDo Redirect URL is required when enabled") - return - } - if err := config.ValidateAbsoluteHTTPURL(req.LinuxDoConnectRedirectURL); err != nil { - response.BadRequest(c, "LinuxDo Redirect URL must be an absolute http(s) URL") - return - } - - // 如果未提供 client_secret,则保留现有值(如有)。 - if req.LinuxDoConnectClientSecret == "" { - if previousSettings.LinuxDoConnectClientSecret == "" { - response.BadRequest(c, "LinuxDo Client Secret is required when enabled") - return - } - req.LinuxDoConnectClientSecret = previousSettings.LinuxDoConnectClientSecret - } - } - - // DingTalk Connect 参数验证 - // 防御性:任何写入路径上把已废弃的 corp_restriction_policy=whitelist 入参 coerce 为 none, - // 避免任何直连 admin API 的客户端把死值写回 DB(前端 UI 已无此选项)。 - req.DingTalkConnectCorpRestrictionPolicy = service.CoerceDingTalkCorpPolicyForWrite(req.DingTalkConnectCorpRestrictionPolicy) - - if req.DingTalkConnectEnabled { - req.DingTalkConnectClientID = strings.TrimSpace(req.DingTalkConnectClientID) - req.DingTalkConnectClientSecret = strings.TrimSpace(req.DingTalkConnectClientSecret) - req.DingTalkConnectRedirectURL = strings.TrimSpace(req.DingTalkConnectRedirectURL) - req.DingTalkConnectCorpRestrictionPolicy = strings.TrimSpace(req.DingTalkConnectCorpRestrictionPolicy) - req.DingTalkConnectInternalCorpID = strings.TrimSpace(req.DingTalkConnectInternalCorpID) - - if req.DingTalkConnectClientID == "" { - response.BadRequest(c, "DingTalk Client ID is required when enabled") - return - } - if req.DingTalkConnectRedirectURL == "" { - response.BadRequest(c, "DingTalk Redirect URL is required when enabled") - return - } - if err := config.ValidateAbsoluteHTTPURL(req.DingTalkConnectRedirectURL); err != nil { - response.BadRequest(c, "DingTalk Redirect URL must be an absolute http(s) URL") - return - } - - // 如果未提供 client_secret,则保留现有值(如有)。 - if req.DingTalkConnectClientSecret == "" { - if previousSettings.DingTalkConnectClientSecret == "" { - response.BadRequest(c, "DingTalk Client Secret is required when enabled") - return - } - req.DingTalkConnectClientSecret = previousSettings.DingTalkConnectClientSecret - } - - // Corp 策略校验(V1/V4 fail-closed) - dingTalkCfg := config.DingTalkConnectConfig{ - Enabled: true, - DingTalkAppKind: "internal_app", // 硬编码:settings 层仅支持 internal_app - AppType: "internal", // 对于 internal_only 策略的默认值 - CorpRestrictionPolicy: req.DingTalkConnectCorpRestrictionPolicy, - InternalCorpID: req.DingTalkConnectInternalCorpID, - } - // 若未填 corp_restriction_policy,保留已有配置 - if dingTalkCfg.CorpRestrictionPolicy == "" { - dingTalkCfg.CorpRestrictionPolicy = previousSettings.DingTalkConnectCorpRestrictionPolicy - } - // 对于 internal_only 策略,app_type 必须为 internal(V1 校验) - if dingTalkCfg.CorpRestrictionPolicy == "internal_only" { - dingTalkCfg.AppType = "internal" - } else { - dingTalkCfg.AppType = "public" - } - if err := config.ValidateDingTalkConfig(dingTalkCfg); err != nil { - response.ErrorWithDetails(c, http.StatusBadRequest, err.Error(), mapDingTalkValidateError(err), nil) - return - } - - // bypass_registration 仅在 internal_only 模式下有意义;其它策略下强制为 false, - // 防止 admin 在切换 policy 时把 bypass 残留在 DB 中(前端 UI 也已隐藏该开关)。 - if dingTalkCfg.CorpRestrictionPolicy != "internal_only" { - req.DingTalkConnectBypassRegistration = false - // 身份同步三开关同理:仅 internal_only 模式下有意义,其它策略强制 false。 - req.DingTalkConnectSyncCorpEmail = false - req.DingTalkConnectSyncDisplayName = false - req.DingTalkConnectSyncDept = false - } - // 身份同步目标 attr key:trimSpace + 空值 fallback 到默认值 - req.DingTalkConnectSyncCorpEmailAttrKey = strings.TrimSpace(req.DingTalkConnectSyncCorpEmailAttrKey) - if req.DingTalkConnectSyncCorpEmailAttrKey == "" { - req.DingTalkConnectSyncCorpEmailAttrKey = "dingtalk_email" - } - req.DingTalkConnectSyncDisplayNameAttrKey = strings.TrimSpace(req.DingTalkConnectSyncDisplayNameAttrKey) - if req.DingTalkConnectSyncDisplayNameAttrKey == "" { - req.DingTalkConnectSyncDisplayNameAttrKey = "dingtalk_name" - } - req.DingTalkConnectSyncDeptAttrKey = strings.TrimSpace(req.DingTalkConnectSyncDeptAttrKey) - if req.DingTalkConnectSyncDeptAttrKey == "" { - req.DingTalkConnectSyncDeptAttrKey = "dingtalk_department" - } - // 身份同步目标 attr 显示名称:trim + 空值 fallback 到默认中文名 - req.DingTalkConnectSyncCorpEmailAttrName = strings.TrimSpace(req.DingTalkConnectSyncCorpEmailAttrName) - if req.DingTalkConnectSyncCorpEmailAttrName == "" { - req.DingTalkConnectSyncCorpEmailAttrName = "钉钉企业邮箱" - } - req.DingTalkConnectSyncDisplayNameAttrName = strings.TrimSpace(req.DingTalkConnectSyncDisplayNameAttrName) - if req.DingTalkConnectSyncDisplayNameAttrName == "" { - req.DingTalkConnectSyncDisplayNameAttrName = "钉钉姓名" - } - req.DingTalkConnectSyncDeptAttrName = strings.TrimSpace(req.DingTalkConnectSyncDeptAttrName) - if req.DingTalkConnectSyncDeptAttrName == "" { - req.DingTalkConnectSyncDeptAttrName = "钉钉部门" - } - } - - if req.WeChatConnectEnabled { - req.WeChatConnectAppID = strings.TrimSpace(req.WeChatConnectAppID) - req.WeChatConnectAppSecret = strings.TrimSpace(req.WeChatConnectAppSecret) - req.WeChatConnectOpenAppID = strings.TrimSpace(req.WeChatConnectOpenAppID) - req.WeChatConnectOpenAppSecret = strings.TrimSpace(req.WeChatConnectOpenAppSecret) - req.WeChatConnectMPAppID = strings.TrimSpace(req.WeChatConnectMPAppID) - req.WeChatConnectMPAppSecret = strings.TrimSpace(req.WeChatConnectMPAppSecret) - req.WeChatConnectMobileAppID = strings.TrimSpace(req.WeChatConnectMobileAppID) - req.WeChatConnectMobileAppSecret = strings.TrimSpace(req.WeChatConnectMobileAppSecret) - req.WeChatConnectMode = strings.ToLower(strings.TrimSpace(req.WeChatConnectMode)) - req.WeChatConnectScopes = strings.TrimSpace(req.WeChatConnectScopes) - req.WeChatConnectRedirectURL = strings.TrimSpace(req.WeChatConnectRedirectURL) - req.WeChatConnectFrontendRedirectURL = strings.TrimSpace(req.WeChatConnectFrontendRedirectURL) - req.WeChatConnectAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectAppID, previousSettings.WeChatConnectAppID)) - req.WeChatConnectRedirectURL = strings.TrimSpace(firstNonEmpty(req.WeChatConnectRedirectURL, previousSettings.WeChatConnectRedirectURL)) - req.WeChatConnectFrontendRedirectURL = strings.TrimSpace(firstNonEmpty(req.WeChatConnectFrontendRedirectURL, previousSettings.WeChatConnectFrontendRedirectURL)) - if req.WeChatConnectMode == "" { - req.WeChatConnectMode = strings.ToLower(strings.TrimSpace(previousSettings.WeChatConnectMode)) - } - if req.WeChatConnectScopes == "" { - req.WeChatConnectScopes = strings.TrimSpace(previousSettings.WeChatConnectScopes) - } - - if req.WeChatConnectMPEnabled && req.WeChatConnectMobileEnabled { - response.BadRequest(c, "WeChat Official Account and Mobile App cannot be enabled at the same time") - return - } - if req.WeChatConnectMode != "" { - switch req.WeChatConnectMode { - case "open", "mp", "mobile": - default: - response.BadRequest(c, "WeChat mode must be open, mp, or mobile") - return - } - } - if !req.WeChatConnectOpenEnabled && !req.WeChatConnectMPEnabled && !req.WeChatConnectMobileEnabled { - switch req.WeChatConnectMode { - case "mp": - req.WeChatConnectMPEnabled = true - case "mobile": - req.WeChatConnectMobileEnabled = true - default: - req.WeChatConnectOpenEnabled = true - } - } - if req.WeChatConnectMode == "" { - if req.WeChatConnectMPEnabled { - req.WeChatConnectMode = "mp" - } else if req.WeChatConnectMobileEnabled { - req.WeChatConnectMode = "mobile" - } else { - req.WeChatConnectMode = "open" - } - } - - req.WeChatConnectOpenAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectOpenAppID, req.WeChatConnectAppID, previousSettings.WeChatConnectOpenAppID, previousSettings.WeChatConnectAppID)) - req.WeChatConnectMPAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectMPAppID, req.WeChatConnectAppID, previousSettings.WeChatConnectMPAppID, previousSettings.WeChatConnectAppID)) - req.WeChatConnectMobileAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectMobileAppID, req.WeChatConnectAppID, previousSettings.WeChatConnectMobileAppID, previousSettings.WeChatConnectAppID)) - - if req.WeChatConnectOpenAppSecret == "" { - req.WeChatConnectOpenAppSecret = strings.TrimSpace(firstNonEmpty(previousSettings.WeChatConnectOpenAppSecret, previousSettings.WeChatConnectAppSecret, req.WeChatConnectAppSecret)) - } - if req.WeChatConnectMPAppSecret == "" { - req.WeChatConnectMPAppSecret = strings.TrimSpace(firstNonEmpty(previousSettings.WeChatConnectMPAppSecret, previousSettings.WeChatConnectAppSecret, req.WeChatConnectAppSecret)) - } - if req.WeChatConnectMobileAppSecret == "" { - req.WeChatConnectMobileAppSecret = strings.TrimSpace(firstNonEmpty(previousSettings.WeChatConnectMobileAppSecret, previousSettings.WeChatConnectAppSecret, req.WeChatConnectAppSecret)) - } - if req.WeChatConnectAppSecret == "" { - req.WeChatConnectAppSecret = strings.TrimSpace(firstNonEmpty(req.WeChatConnectOpenAppSecret, req.WeChatConnectMPAppSecret, req.WeChatConnectMobileAppSecret, previousSettings.WeChatConnectAppSecret)) - } - - if req.WeChatConnectOpenEnabled { - if req.WeChatConnectOpenAppID == "" { - response.BadRequest(c, "WeChat PC App ID is required when enabled") - return - } - if req.WeChatConnectOpenAppSecret == "" { - response.BadRequest(c, "WeChat PC App Secret is required when enabled") - return - } - } - if req.WeChatConnectMPEnabled { - if req.WeChatConnectMPAppID == "" { - response.BadRequest(c, "WeChat Official Account App ID is required when enabled") - return - } - if req.WeChatConnectMPAppSecret == "" { - response.BadRequest(c, "WeChat Official Account App Secret is required when enabled") - return - } - } - if req.WeChatConnectMobileEnabled { - if req.WeChatConnectMobileAppID == "" { - response.BadRequest(c, "WeChat Mobile App ID is required when enabled") - return - } - if req.WeChatConnectMobileAppSecret == "" { - response.BadRequest(c, "WeChat Mobile App Secret is required when enabled") - return - } - } - - if req.WeChatConnectScopes == "" { - if req.WeChatConnectMPEnabled { - req.WeChatConnectScopes = service.DefaultWeChatConnectScopesForMode("mp") - } else { - req.WeChatConnectScopes = service.DefaultWeChatConnectScopesForMode(req.WeChatConnectMode) - } - } - if req.WeChatConnectOpenEnabled || req.WeChatConnectMPEnabled { - if req.WeChatConnectRedirectURL == "" { - response.BadRequest(c, "WeChat Redirect URL is required when web oauth is enabled") - return - } - if err := config.ValidateAbsoluteHTTPURL(req.WeChatConnectRedirectURL); err != nil { - response.BadRequest(c, "WeChat Redirect URL must be an absolute http(s) URL") - return - } - if req.WeChatConnectFrontendRedirectURL == "" { - req.WeChatConnectFrontendRedirectURL = "/auth/wechat/callback" - } - if err := config.ValidateFrontendRedirectURL(req.WeChatConnectFrontendRedirectURL); err != nil { - response.BadRequest(c, "WeChat Frontend Redirect URL is invalid") - return - } - } - } - - // Generic OIDC 参数验证 - oidcUsePKCE, oidcValidateIDToken, err := h.settingService.OIDCSecurityWriteDefaults(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - if req.OIDCConnectEnabled { - req.OIDCConnectProviderName = strings.TrimSpace(req.OIDCConnectProviderName) - req.OIDCConnectClientID = strings.TrimSpace(req.OIDCConnectClientID) - req.OIDCConnectClientSecret = strings.TrimSpace(req.OIDCConnectClientSecret) - req.OIDCConnectIssuerURL = strings.TrimSpace(req.OIDCConnectIssuerURL) - req.OIDCConnectDiscoveryURL = strings.TrimSpace(req.OIDCConnectDiscoveryURL) - req.OIDCConnectAuthorizeURL = strings.TrimSpace(req.OIDCConnectAuthorizeURL) - req.OIDCConnectTokenURL = strings.TrimSpace(req.OIDCConnectTokenURL) - req.OIDCConnectUserInfoURL = strings.TrimSpace(req.OIDCConnectUserInfoURL) - req.OIDCConnectJWKSURL = strings.TrimSpace(req.OIDCConnectJWKSURL) - req.OIDCConnectScopes = strings.TrimSpace(req.OIDCConnectScopes) - req.OIDCConnectRedirectURL = strings.TrimSpace(req.OIDCConnectRedirectURL) - req.OIDCConnectFrontendRedirectURL = strings.TrimSpace(req.OIDCConnectFrontendRedirectURL) - req.OIDCConnectTokenAuthMethod = strings.ToLower(strings.TrimSpace(req.OIDCConnectTokenAuthMethod)) - req.OIDCConnectAllowedSigningAlgs = strings.TrimSpace(req.OIDCConnectAllowedSigningAlgs) - req.OIDCConnectUserInfoEmailPath = strings.TrimSpace(req.OIDCConnectUserInfoEmailPath) - req.OIDCConnectUserInfoIDPath = strings.TrimSpace(req.OIDCConnectUserInfoIDPath) - req.OIDCConnectUserInfoUsernamePath = strings.TrimSpace(req.OIDCConnectUserInfoUsernamePath) - req.OIDCConnectProviderName = strings.TrimSpace(firstNonEmpty(req.OIDCConnectProviderName, previousSettings.OIDCConnectProviderName, "OIDC")) - req.OIDCConnectClientID = strings.TrimSpace(firstNonEmpty(req.OIDCConnectClientID, previousSettings.OIDCConnectClientID)) - req.OIDCConnectIssuerURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectIssuerURL, previousSettings.OIDCConnectIssuerURL)) - req.OIDCConnectDiscoveryURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectDiscoveryURL, previousSettings.OIDCConnectDiscoveryURL)) - req.OIDCConnectAuthorizeURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectAuthorizeURL, previousSettings.OIDCConnectAuthorizeURL)) - req.OIDCConnectTokenURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectTokenURL, previousSettings.OIDCConnectTokenURL)) - req.OIDCConnectUserInfoURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoURL, previousSettings.OIDCConnectUserInfoURL)) - req.OIDCConnectJWKSURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectJWKSURL, previousSettings.OIDCConnectJWKSURL)) - req.OIDCConnectScopes = strings.TrimSpace(firstNonEmpty(req.OIDCConnectScopes, previousSettings.OIDCConnectScopes, "openid email profile")) - req.OIDCConnectRedirectURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectRedirectURL, previousSettings.OIDCConnectRedirectURL)) - req.OIDCConnectFrontendRedirectURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectFrontendRedirectURL, previousSettings.OIDCConnectFrontendRedirectURL, "/auth/oidc/callback")) - req.OIDCConnectTokenAuthMethod = strings.ToLower(strings.TrimSpace(firstNonEmpty(req.OIDCConnectTokenAuthMethod, previousSettings.OIDCConnectTokenAuthMethod, "client_secret_post"))) - req.OIDCConnectAllowedSigningAlgs = strings.TrimSpace(firstNonEmpty(req.OIDCConnectAllowedSigningAlgs, previousSettings.OIDCConnectAllowedSigningAlgs, "RS256,ES256,PS256")) - req.OIDCConnectUserInfoEmailPath = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoEmailPath, previousSettings.OIDCConnectUserInfoEmailPath)) - req.OIDCConnectUserInfoIDPath = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoIDPath, previousSettings.OIDCConnectUserInfoIDPath)) - req.OIDCConnectUserInfoUsernamePath = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoUsernamePath, previousSettings.OIDCConnectUserInfoUsernamePath)) - if req.OIDCConnectUsePKCE != nil { - oidcUsePKCE = *req.OIDCConnectUsePKCE - } - if req.OIDCConnectValidateIDToken != nil { - oidcValidateIDToken = *req.OIDCConnectValidateIDToken - } - if req.OIDCConnectClockSkewSeconds == 0 { - req.OIDCConnectClockSkewSeconds = previousSettings.OIDCConnectClockSkewSeconds - if req.OIDCConnectClockSkewSeconds == 0 { - req.OIDCConnectClockSkewSeconds = 120 - } - } - - if req.OIDCConnectClientID == "" { - response.BadRequest(c, "OIDC Client ID is required when enabled") - return - } - if req.OIDCConnectIssuerURL == "" { - response.BadRequest(c, "OIDC Issuer URL is required when enabled") - return - } - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectIssuerURL); err != nil { - response.BadRequest(c, "OIDC Issuer URL must be an absolute http(s) URL") - return - } - if req.OIDCConnectDiscoveryURL != "" { - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectDiscoveryURL); err != nil { - response.BadRequest(c, "OIDC Discovery URL must be an absolute http(s) URL") - return - } - } - if req.OIDCConnectAuthorizeURL != "" { - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectAuthorizeURL); err != nil { - response.BadRequest(c, "OIDC Authorize URL must be an absolute http(s) URL") - return - } - } - if req.OIDCConnectTokenURL != "" { - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectTokenURL); err != nil { - response.BadRequest(c, "OIDC Token URL must be an absolute http(s) URL") - return - } - } - if req.OIDCConnectUserInfoURL != "" { - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectUserInfoURL); err != nil { - response.BadRequest(c, "OIDC UserInfo URL must be an absolute http(s) URL") - return - } - } - if req.OIDCConnectRedirectURL == "" { - response.BadRequest(c, "OIDC Redirect URL is required when enabled") - return - } - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectRedirectURL); err != nil { - response.BadRequest(c, "OIDC Redirect URL must be an absolute http(s) URL") - return - } - if req.OIDCConnectFrontendRedirectURL == "" { - response.BadRequest(c, "OIDC Frontend Redirect URL is required when enabled") - return - } - if err := config.ValidateFrontendRedirectURL(req.OIDCConnectFrontendRedirectURL); err != nil { - response.BadRequest(c, "OIDC Frontend Redirect URL is invalid") - return - } - if !scopesContainOpenID(req.OIDCConnectScopes) { - response.BadRequest(c, "OIDC scopes must contain openid") - return - } - switch req.OIDCConnectTokenAuthMethod { - case "", "client_secret_post", "client_secret_basic", "none": - default: - response.BadRequest(c, "OIDC Token Auth Method must be one of client_secret_post/client_secret_basic/none") - return - } - if req.OIDCConnectClockSkewSeconds < 0 || req.OIDCConnectClockSkewSeconds > 600 { - response.BadRequest(c, "OIDC clock skew seconds must be between 0 and 600") - return - } - if oidcValidateIDToken && req.OIDCConnectAllowedSigningAlgs == "" { - response.BadRequest(c, "OIDC Allowed Signing Algs is required when validate_id_token=true") - return - } - if req.OIDCConnectJWKSURL != "" { - if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectJWKSURL); err != nil { - response.BadRequest(c, "OIDC JWKS URL must be an absolute http(s) URL") - return - } - } - if req.OIDCConnectTokenAuthMethod == "" || req.OIDCConnectTokenAuthMethod == "client_secret_post" || req.OIDCConnectTokenAuthMethod == "client_secret_basic" { - if req.OIDCConnectClientSecret == "" { - if previousSettings.OIDCConnectClientSecret == "" { - response.BadRequest(c, "OIDC Client Secret is required when enabled") - return - } - req.OIDCConnectClientSecret = previousSettings.OIDCConnectClientSecret - } - } - } - - // “购买订阅”页面配置验证 - purchaseEnabled := previousSettings.PurchaseSubscriptionEnabled - if req.PurchaseSubscriptionEnabled != nil { - purchaseEnabled = *req.PurchaseSubscriptionEnabled - } - purchaseURL := previousSettings.PurchaseSubscriptionURL - if req.PurchaseSubscriptionURL != nil { - purchaseURL = strings.TrimSpace(*req.PurchaseSubscriptionURL) - } - - // - 启用时要求 URL 合法且非空 - // - 禁用时允许为空;若提供了 URL 也做基本校验,避免误配置 - if purchaseEnabled { - if purchaseURL == "" { - response.BadRequest(c, "Purchase Subscription URL is required when enabled") - return - } - if err := config.ValidateAbsoluteHTTPURL(purchaseURL); err != nil { - response.BadRequest(c, "Purchase Subscription URL must be an absolute http(s) URL") - return - } - } else if purchaseURL != "" { - if err := config.ValidateAbsoluteHTTPURL(purchaseURL); err != nil { - response.BadRequest(c, "Purchase Subscription URL must be an absolute http(s) URL") - return - } - } - - // Frontend URL 验证 - req.FrontendURL = strings.TrimSpace(req.FrontendURL) - if req.FrontendURL != "" { - if err := config.ValidateAbsoluteHTTPURL(req.FrontendURL); err != nil { - response.BadRequest(c, "Frontend URL must be an absolute http(s) URL") - return - } - } - - // 自定义菜单项验证 - const ( - maxCustomMenuItems = 20 - maxMenuItemLabelLen = 50 - maxMenuItemURLLen = 2048 - maxMenuItemIconSVGLen = 10 * 1024 // 10KB - maxMenuItemIDLen = 32 - ) - - customMenuJSON := previousSettings.CustomMenuItems - if req.CustomMenuItems != nil { - items := *req.CustomMenuItems - if len(items) > maxCustomMenuItems { - response.BadRequest(c, "Too many custom menu items (max 20)") - return - } - for i, item := range items { - if strings.TrimSpace(item.Label) == "" { - response.BadRequest(c, "Custom menu item label is required") - return - } - if len(item.Label) > maxMenuItemLabelLen { - response.BadRequest(c, "Custom menu item label is too long (max 50 characters)") - return - } - urlTrimmed := strings.TrimSpace(item.URL) - if strings.HasPrefix(urlTrimmed, "md:") { - // Markdown page mode: URL = "md:" - slug := strings.TrimPrefix(urlTrimmed, "md:") - if slug == "" { - response.BadRequest(c, "Custom menu item markdown slug cannot be empty (use md:slug format)") - return - } - } else { - if urlTrimmed == "" { - response.BadRequest(c, "Custom menu item URL is required (use md:slug for markdown pages)") - return - } - if len(item.URL) > maxMenuItemURLLen { - response.BadRequest(c, "Custom menu item URL is too long (max 2048 characters)") - return - } - if err := config.ValidateAbsoluteHTTPURL(urlTrimmed); err != nil { - response.BadRequest(c, "Custom menu item URL must be an absolute http(s) URL or md:") - return - } - } - if item.Visibility != "user" && item.Visibility != "admin" { - response.BadRequest(c, "Custom menu item visibility must be 'user' or 'admin'") - return - } - if len(item.IconSVG) > maxMenuItemIconSVGLen { - response.BadRequest(c, "Custom menu item icon SVG is too large (max 10KB)") - return - } - // Auto-generate ID if missing - if strings.TrimSpace(item.ID) == "" { - id, err := generateMenuItemID() - if err != nil { - response.Error(c, http.StatusInternalServerError, "Failed to generate menu item ID") - return - } - items[i].ID = id - } else if len(item.ID) > maxMenuItemIDLen { - response.BadRequest(c, "Custom menu item ID is too long (max 32 characters)") - return - } else if !menuItemIDPattern.MatchString(item.ID) { - response.BadRequest(c, "Custom menu item ID contains invalid characters (only a-z, A-Z, 0-9, - and _ are allowed)") - return - } - } - // ID uniqueness check - seen := make(map[string]struct{}, len(items)) - for _, item := range items { - if _, exists := seen[item.ID]; exists { - response.BadRequest(c, "Duplicate custom menu item ID: "+item.ID) - return - } - seen[item.ID] = struct{}{} - } - menuBytes, err := json.Marshal(items) - if err != nil { - response.BadRequest(c, "Failed to serialize custom menu items") - return - } - customMenuJSON = string(menuBytes) - } - - // 自定义端点验证 - const ( - maxCustomEndpoints = 10 - maxEndpointNameLen = 50 - maxEndpointURLLen = 2048 - maxEndpointDescriptionLen = 200 - ) - - customEndpointsJSON := previousSettings.CustomEndpoints - if req.CustomEndpoints != nil { - endpoints := *req.CustomEndpoints - if len(endpoints) > maxCustomEndpoints { - response.BadRequest(c, "Too many custom endpoints (max 10)") - return - } - for _, ep := range endpoints { - if strings.TrimSpace(ep.Name) == "" { - response.BadRequest(c, "Custom endpoint name is required") - return - } - if len(ep.Name) > maxEndpointNameLen { - response.BadRequest(c, "Custom endpoint name is too long (max 50 characters)") - return - } - if strings.TrimSpace(ep.Endpoint) == "" { - response.BadRequest(c, "Custom endpoint URL is required") - return - } - if len(ep.Endpoint) > maxEndpointURLLen { - response.BadRequest(c, "Custom endpoint URL is too long (max 2048 characters)") - return - } - if err := config.ValidateAbsoluteHTTPURL(strings.TrimSpace(ep.Endpoint)); err != nil { - response.BadRequest(c, "Custom endpoint URL must be an absolute http(s) URL") - return - } - if len(ep.Description) > maxEndpointDescriptionLen { - response.BadRequest(c, "Custom endpoint description is too long (max 200 characters)") - return - } - } - endpointBytes, err := json.Marshal(endpoints) - if err != nil { - response.BadRequest(c, "Failed to serialize custom endpoints") - return - } - customEndpointsJSON = string(endpointBytes) - } - - // Ops metrics collector interval validation (seconds). - if req.OpsMetricsIntervalSeconds != nil { - v := *req.OpsMetricsIntervalSeconds - if v < 60 { - v = 60 - } - if v > 3600 { - v = 3600 - } - req.OpsMetricsIntervalSeconds = &v - } - defaultSubscriptions := make([]service.DefaultSubscriptionSetting, 0, len(req.DefaultSubscriptions)) - for _, sub := range req.DefaultSubscriptions { - defaultSubscriptions = append(defaultSubscriptions, service.DefaultSubscriptionSetting{ - GroupID: sub.GroupID, - ValidityDays: sub.ValidityDays, - }) - } - - // 验证最低版本号格式(空字符串=禁用,或合法 semver) - if req.MinClaudeCodeVersion != "" { - if !semverPattern.MatchString(req.MinClaudeCodeVersion) { - response.Error(c, http.StatusBadRequest, "min_claude_code_version must be empty or a valid semver (e.g. 2.1.63)") - return - } - } - - // 验证最高版本号格式(空字符串=禁用,或合法 semver) - if req.MaxClaudeCodeVersion != "" { - if !semverPattern.MatchString(req.MaxClaudeCodeVersion) { - response.Error(c, http.StatusBadRequest, "max_claude_code_version must be empty or a valid semver (e.g. 3.0.0)") - return - } - } - if req.AntigravityUserAgentVersion != nil { - normalized := strings.TrimSpace(*req.AntigravityUserAgentVersion) - req.AntigravityUserAgentVersion = &normalized - if normalized != "" && !semverPattern.MatchString(normalized) { - response.Error(c, http.StatusBadRequest, "antigravity_user_agent_version must be empty or a valid semver (e.g. 1.23.2)") - return - } - } - if req.OpenAICodexUserAgent != nil { - normalized := strings.TrimSpace(*req.OpenAICodexUserAgent) - req.OpenAICodexUserAgent = &normalized - // 仅做长度上限保护,不限制具体格式(运维需要可自由调整 codex 版本号) - if len(normalized) > 512 { - response.Error(c, http.StatusBadRequest, "openai_codex_user_agent must be at most 512 characters") - return - } - } - - // codex_cli_only 加固:最低/最高 Codex 版本(空=禁用,或合法 semver;max>=min) - if req.MinCodexVersion != "" && !semverPattern.MatchString(req.MinCodexVersion) { - response.Error(c, http.StatusBadRequest, "min_codex_version must be empty or a valid semver (e.g. 0.141.0)") - return - } - if req.MaxCodexVersion != "" && !semverPattern.MatchString(req.MaxCodexVersion) { - response.Error(c, http.StatusBadRequest, "max_codex_version must be empty or a valid semver (e.g. 0.200.0)") - return - } - if req.MinCodexVersion != "" && req.MaxCodexVersion != "" && service.CompareVersions(req.MaxCodexVersion, req.MinCodexVersion) < 0 { - response.Error(c, http.StatusBadRequest, "max_codex_version must be greater than or equal to min_codex_version") - return - } - // codex_cli_only 黑/白名单:非空须为合法 []AllowedClientEntry JSON。 - // 黑名单 OR 宽 deny(允许 originator-only);白名单双因子 AND,额外要求每条可命中(非空 originator + ua_contains)。 - if err := service.ValidateCodexClientEntriesJSON(req.CodexCLIOnlyBlacklist); err != nil { - response.Error(c, http.StatusBadRequest, "codex_cli_only_blacklist "+err.Error()) - return - } - if err := service.ValidateCodexWhitelistEntriesJSON(req.CodexCLIOnlyWhitelist); err != nil { - response.Error(c, http.StatusBadRequest, "codex_cli_only_whitelist "+err.Error()) - return - } - if err := service.ValidateEngineFingerprintSignalsJSON(req.CodexCLIOnlyEngineFingerprintSignals); err != nil { - response.Error(c, http.StatusBadRequest, "codex_cli_only_engine_fingerprint_signals "+err.Error()) - return - } - - // 交叉验证:如果同时设置了最低和最高版本号,最高版本号必须 >= 最低版本号 - if req.MinClaudeCodeVersion != "" && req.MaxClaudeCodeVersion != "" { - if service.CompareVersions(req.MaxClaudeCodeVersion, req.MinClaudeCodeVersion) < 0 { - response.Error(c, http.StatusBadRequest, "max_claude_code_version must be greater than or equal to min_claude_code_version") - return - } - } - - // cyber 会话屏蔽 TTL 校验:提供时必须 > 0 - if req.CyberSessionBlockTTLSeconds != nil && *req.CyberSessionBlockTTLSeconds <= 0 { - response.BadRequest(c, "cyber_session_block_ttl_seconds must be > 0") - return - } - - settings := &service.SystemSettings{ - // 系统全局 platform quota 默认值(整体替换语义) - DefaultPlatformQuotas: req.DefaultPlatformQuotas, - - RegistrationEnabled: req.RegistrationEnabled, - EmailVerifyEnabled: req.EmailVerifyEnabled, - RegistrationEmailSuffixWhitelist: req.RegistrationEmailSuffixWhitelist, - PromoCodeEnabled: req.PromoCodeEnabled, - PasswordResetEnabled: req.PasswordResetEnabled, - FrontendURL: req.FrontendURL, - InvitationCodeEnabled: req.InvitationCodeEnabled, - TotpEnabled: req.TotpEnabled, - LoginAgreementEnabled: req.LoginAgreementEnabled, - LoginAgreementMode: loginAgreementMode, - LoginAgreementUpdatedAt: loginAgreementUpdatedAt, - LoginAgreementDocuments: loginAgreementDocuments, - SMTPHost: req.SMTPHost, - SMTPPort: req.SMTPPort, - SMTPUsername: req.SMTPUsername, - SMTPPassword: req.SMTPPassword, - SMTPFrom: req.SMTPFrom, - SMTPFromName: req.SMTPFromName, - SMTPUseTLS: req.SMTPUseTLS, - TurnstileEnabled: req.TurnstileEnabled, - TurnstileSiteKey: req.TurnstileSiteKey, - TurnstileSecretKey: req.TurnstileSecretKey, - APIKeyACLTrustForwardedIP: func() bool { - if req.APIKeyACLTrustForwardedIP != nil { - return *req.APIKeyACLTrustForwardedIP - } - return previousSettings.APIKeyACLTrustForwardedIP - }(), - LinuxDoConnectEnabled: req.LinuxDoConnectEnabled, - LinuxDoConnectClientID: req.LinuxDoConnectClientID, - LinuxDoConnectClientSecret: req.LinuxDoConnectClientSecret, - LinuxDoConnectRedirectURL: req.LinuxDoConnectRedirectURL, - DingTalkConnectEnabled: req.DingTalkConnectEnabled, - DingTalkConnectClientID: req.DingTalkConnectClientID, - DingTalkConnectClientSecret: req.DingTalkConnectClientSecret, - DingTalkConnectRedirectURL: req.DingTalkConnectRedirectURL, - DingTalkConnectCorpRestrictionPolicy: req.DingTalkConnectCorpRestrictionPolicy, - DingTalkConnectInternalCorpID: req.DingTalkConnectInternalCorpID, - DingTalkConnectBypassRegistration: req.DingTalkConnectBypassRegistration, - DingTalkConnectSyncCorpEmail: req.DingTalkConnectSyncCorpEmail, - DingTalkConnectSyncDisplayName: req.DingTalkConnectSyncDisplayName, - DingTalkConnectSyncDept: req.DingTalkConnectSyncDept, - DingTalkConnectSyncCorpEmailAttrKey: req.DingTalkConnectSyncCorpEmailAttrKey, - DingTalkConnectSyncDisplayNameAttrKey: req.DingTalkConnectSyncDisplayNameAttrKey, - DingTalkConnectSyncDeptAttrKey: req.DingTalkConnectSyncDeptAttrKey, - DingTalkConnectSyncCorpEmailAttrName: req.DingTalkConnectSyncCorpEmailAttrName, - DingTalkConnectSyncDisplayNameAttrName: req.DingTalkConnectSyncDisplayNameAttrName, - DingTalkConnectSyncDeptAttrName: req.DingTalkConnectSyncDeptAttrName, - WeChatConnectEnabled: req.WeChatConnectEnabled, - WeChatConnectAppID: req.WeChatConnectAppID, - WeChatConnectAppSecret: req.WeChatConnectAppSecret, - WeChatConnectOpenAppID: req.WeChatConnectOpenAppID, - WeChatConnectOpenAppSecret: req.WeChatConnectOpenAppSecret, - WeChatConnectMPAppID: req.WeChatConnectMPAppID, - WeChatConnectMPAppSecret: req.WeChatConnectMPAppSecret, - WeChatConnectMobileAppID: req.WeChatConnectMobileAppID, - WeChatConnectMobileAppSecret: req.WeChatConnectMobileAppSecret, - WeChatConnectOpenEnabled: req.WeChatConnectOpenEnabled, - WeChatConnectMPEnabled: req.WeChatConnectMPEnabled, - WeChatConnectMobileEnabled: req.WeChatConnectMobileEnabled, - WeChatConnectMode: req.WeChatConnectMode, - WeChatConnectScopes: req.WeChatConnectScopes, - WeChatConnectRedirectURL: req.WeChatConnectRedirectURL, - WeChatConnectFrontendRedirectURL: req.WeChatConnectFrontendRedirectURL, - OIDCConnectEnabled: req.OIDCConnectEnabled, - OIDCConnectProviderName: req.OIDCConnectProviderName, - OIDCConnectClientID: req.OIDCConnectClientID, - OIDCConnectClientSecret: req.OIDCConnectClientSecret, - OIDCConnectIssuerURL: req.OIDCConnectIssuerURL, - OIDCConnectDiscoveryURL: req.OIDCConnectDiscoveryURL, - OIDCConnectAuthorizeURL: req.OIDCConnectAuthorizeURL, - OIDCConnectTokenURL: req.OIDCConnectTokenURL, - OIDCConnectUserInfoURL: req.OIDCConnectUserInfoURL, - OIDCConnectJWKSURL: req.OIDCConnectJWKSURL, - OIDCConnectScopes: req.OIDCConnectScopes, - OIDCConnectRedirectURL: req.OIDCConnectRedirectURL, - OIDCConnectFrontendRedirectURL: req.OIDCConnectFrontendRedirectURL, - OIDCConnectTokenAuthMethod: req.OIDCConnectTokenAuthMethod, - OIDCConnectUsePKCE: oidcUsePKCE, - OIDCConnectValidateIDToken: oidcValidateIDToken, - OIDCConnectAllowedSigningAlgs: req.OIDCConnectAllowedSigningAlgs, - OIDCConnectClockSkewSeconds: req.OIDCConnectClockSkewSeconds, - OIDCConnectRequireEmailVerified: req.OIDCConnectRequireEmailVerified, - OIDCConnectUserInfoEmailPath: req.OIDCConnectUserInfoEmailPath, - OIDCConnectUserInfoIDPath: req.OIDCConnectUserInfoIDPath, - OIDCConnectUserInfoUsernamePath: req.OIDCConnectUserInfoUsernamePath, - GitHubOAuthEnabled: req.GitHubOAuthEnabled, - GitHubOAuthClientID: req.GitHubOAuthClientID, - GitHubOAuthClientSecret: req.GitHubOAuthClientSecret, - GitHubOAuthRedirectURL: req.GitHubOAuthRedirectURL, - GitHubOAuthFrontendRedirectURL: req.GitHubOAuthFrontendRedirectURL, - GoogleOAuthEnabled: req.GoogleOAuthEnabled, - GoogleOAuthClientID: req.GoogleOAuthClientID, - GoogleOAuthClientSecret: req.GoogleOAuthClientSecret, - GoogleOAuthRedirectURL: req.GoogleOAuthRedirectURL, - GoogleOAuthFrontendRedirectURL: req.GoogleOAuthFrontendRedirectURL, - SiteName: req.SiteName, - SiteLogo: req.SiteLogo, - SiteSubtitle: req.SiteSubtitle, - APIBaseURL: req.APIBaseURL, - ContactInfo: req.ContactInfo, - DocURL: req.DocURL, - HomeContent: req.HomeContent, - HideCcsImportButton: req.HideCcsImportButton, - PurchaseSubscriptionEnabled: purchaseEnabled, - PurchaseSubscriptionURL: purchaseURL, - TableDefaultPageSize: req.TableDefaultPageSize, - TablePageSizeOptions: req.TablePageSizeOptions, - CustomMenuItems: customMenuJSON, - CustomEndpoints: customEndpointsJSON, - DefaultConcurrency: req.DefaultConcurrency, - DefaultBalance: req.DefaultBalance, - AffiliateRebateRate: affiliateRebateRate, - AffiliateRebateFreezeHours: affiliateRebateFreezeHours, - AffiliateRebateDurationDays: affiliateRebateDurationDays, - AffiliateRebatePerInviteeCap: affiliateRebatePerInviteeCap, - DefaultUserRPMLimit: req.DefaultUserRPMLimit, - DefaultSubscriptions: defaultSubscriptions, - EnableModelFallback: req.EnableModelFallback, - FallbackModelAnthropic: req.FallbackModelAnthropic, - FallbackModelOpenAI: req.FallbackModelOpenAI, - FallbackModelGemini: req.FallbackModelGemini, - FallbackModelAntigravity: req.FallbackModelAntigravity, - EnableIdentityPatch: req.EnableIdentityPatch, - IdentityPatchPrompt: req.IdentityPatchPrompt, - MinClaudeCodeVersion: req.MinClaudeCodeVersion, - MaxClaudeCodeVersion: req.MaxClaudeCodeVersion, - AllowUngroupedKeyScheduling: req.AllowUngroupedKeyScheduling, - BackendModeEnabled: req.BackendModeEnabled, - AllowUserViewErrorRequests: func() bool { - if req.AllowUserViewErrorRequests != nil { - return *req.AllowUserViewErrorRequests - } - return previousSettings.AllowUserViewErrorRequests - }(), - OpsMonitoringEnabled: func() bool { - if req.OpsMonitoringEnabled != nil { - return *req.OpsMonitoringEnabled - } - return previousSettings.OpsMonitoringEnabled - }(), - OpsRealtimeMonitoringEnabled: func() bool { - if req.OpsRealtimeMonitoringEnabled != nil { - return *req.OpsRealtimeMonitoringEnabled - } - return previousSettings.OpsRealtimeMonitoringEnabled - }(), - OpsQueryModeDefault: func() string { - if req.OpsQueryModeDefault != nil { - return *req.OpsQueryModeDefault - } - return previousSettings.OpsQueryModeDefault - }(), - OpsMetricsIntervalSeconds: func() int { - if req.OpsMetricsIntervalSeconds != nil { - return *req.OpsMetricsIntervalSeconds - } - return previousSettings.OpsMetricsIntervalSeconds - }(), - EnableFingerprintUnification: func() bool { - if req.EnableFingerprintUnification != nil { - return *req.EnableFingerprintUnification - } - return previousSettings.EnableFingerprintUnification - }(), - EnableMetadataPassthrough: func() bool { - if req.EnableMetadataPassthrough != nil { - return *req.EnableMetadataPassthrough - } - return previousSettings.EnableMetadataPassthrough - }(), - EnableCCHSigning: func() bool { - if req.EnableCCHSigning != nil { - return *req.EnableCCHSigning - } - return previousSettings.EnableCCHSigning - }(), - EnableClaudeOAuthSystemPromptInjection: func() bool { - if req.EnableClaudeOAuthSystemPromptInjection != nil { - return *req.EnableClaudeOAuthSystemPromptInjection - } - return previousSettings.EnableClaudeOAuthSystemPromptInjection - }(), - ClaudeOAuthSystemPrompt: func() string { - if req.ClaudeOAuthSystemPrompt != nil { - return *req.ClaudeOAuthSystemPrompt - } - return previousSettings.ClaudeOAuthSystemPrompt - }(), - ClaudeOAuthSystemPromptBlocks: func() string { - if req.ClaudeOAuthSystemPromptBlocks != nil { - return *req.ClaudeOAuthSystemPromptBlocks - } - return previousSettings.ClaudeOAuthSystemPromptBlocks - }(), - EnableAnthropicCacheTTL1hInjection: func() bool { - if req.EnableAnthropicCacheTTL1hInjection != nil { - return *req.EnableAnthropicCacheTTL1hInjection - } - return previousSettings.EnableAnthropicCacheTTL1hInjection - }(), - RewriteMessageCacheControl: func() bool { - if req.RewriteMessageCacheControl != nil { - return *req.RewriteMessageCacheControl - } - return previousSettings.RewriteMessageCacheControl - }(), - EnableClientDatelineNormalization: func() bool { - if req.EnableClientDatelineNormalization != nil { - return *req.EnableClientDatelineNormalization - } - return previousSettings.EnableClientDatelineNormalization - }(), - AntigravityUserAgentVersion: func() string { - if req.AntigravityUserAgentVersion != nil { - return *req.AntigravityUserAgentVersion - } - return previousSettings.AntigravityUserAgentVersion - }(), - OpenAICodexUserAgent: func() string { - if req.OpenAICodexUserAgent != nil { - return *req.OpenAICodexUserAgent - } - return previousSettings.OpenAICodexUserAgent - }(), - MinCodexVersion: strings.TrimSpace(req.MinCodexVersion), - MaxCodexVersion: strings.TrimSpace(req.MaxCodexVersion), - CodexCLIOnlyBlacklist: strings.TrimSpace(req.CodexCLIOnlyBlacklist), - CodexCLIOnlyWhitelist: strings.TrimSpace(req.CodexCLIOnlyWhitelist), - CodexCLIOnlyAllowAppServerClients: func() bool { - if req.CodexCLIOnlyAllowAppServerClients != nil { - return *req.CodexCLIOnlyAllowAppServerClients - } - return previousSettings.CodexCLIOnlyAllowAppServerClients - }(), - CodexCLIOnlyEngineFingerprintSignals: strings.TrimSpace(req.CodexCLIOnlyEngineFingerprintSignals), - PaymentVisibleMethodAlipaySource: func() string { - if req.PaymentVisibleMethodAlipaySource != nil { - return strings.TrimSpace(*req.PaymentVisibleMethodAlipaySource) - } - return previousSettings.PaymentVisibleMethodAlipaySource - }(), - PaymentVisibleMethodWxpaySource: func() string { - if req.PaymentVisibleMethodWxpaySource != nil { - return strings.TrimSpace(*req.PaymentVisibleMethodWxpaySource) - } - return previousSettings.PaymentVisibleMethodWxpaySource - }(), - PaymentVisibleMethodAlipayEnabled: func() bool { - if req.PaymentVisibleMethodAlipayEnabled != nil { - return *req.PaymentVisibleMethodAlipayEnabled - } - return previousSettings.PaymentVisibleMethodAlipayEnabled - }(), - PaymentVisibleMethodWxpayEnabled: func() bool { - if req.PaymentVisibleMethodWxpayEnabled != nil { - return *req.PaymentVisibleMethodWxpayEnabled - } - return previousSettings.PaymentVisibleMethodWxpayEnabled - }(), - OpenAIAdvancedSchedulerEnabled: func() bool { - if req.OpenAIAdvancedSchedulerEnabled != nil { - return *req.OpenAIAdvancedSchedulerEnabled - } - return previousSettings.OpenAIAdvancedSchedulerEnabled - }(), - OpenAIAdvancedSchedulerStickyWeightedEnabled: func() bool { - if req.OpenAIAdvancedSchedulerStickyWeightedEnabled != nil { - return *req.OpenAIAdvancedSchedulerStickyWeightedEnabled - } - return previousSettings.OpenAIAdvancedSchedulerStickyWeightedEnabled - }(), - OpenAIAdvancedSchedulerSubscriptionPriorityEnabled: func() bool { - if req.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled != nil { - return *req.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled - } - return previousSettings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled - }(), - OpenAIAdvancedSchedulerLBTopK: stringSetting(req.OpenAIAdvancedSchedulerLBTopK, previousSettings.OpenAIAdvancedSchedulerLBTopK), - OpenAIAdvancedSchedulerWeightPriority: stringSetting(req.OpenAIAdvancedSchedulerWeightPriority, previousSettings.OpenAIAdvancedSchedulerWeightPriority), - OpenAIAdvancedSchedulerWeightLoad: stringSetting(req.OpenAIAdvancedSchedulerWeightLoad, previousSettings.OpenAIAdvancedSchedulerWeightLoad), - OpenAIAdvancedSchedulerWeightQueue: stringSetting(req.OpenAIAdvancedSchedulerWeightQueue, previousSettings.OpenAIAdvancedSchedulerWeightQueue), - OpenAIAdvancedSchedulerWeightErrorRate: stringSetting(req.OpenAIAdvancedSchedulerWeightErrorRate, previousSettings.OpenAIAdvancedSchedulerWeightErrorRate), - OpenAIAdvancedSchedulerWeightTTFT: stringSetting(req.OpenAIAdvancedSchedulerWeightTTFT, previousSettings.OpenAIAdvancedSchedulerWeightTTFT), - OpenAIAdvancedSchedulerWeightReset: stringSetting(req.OpenAIAdvancedSchedulerWeightReset, previousSettings.OpenAIAdvancedSchedulerWeightReset), - OpenAIAdvancedSchedulerWeightQuotaHeadroom: stringSetting(req.OpenAIAdvancedSchedulerWeightQuotaHeadroom, previousSettings.OpenAIAdvancedSchedulerWeightQuotaHeadroom), - OpenAIAdvancedSchedulerWeightPreviousResponse: stringSetting(req.OpenAIAdvancedSchedulerWeightPreviousResponse, previousSettings.OpenAIAdvancedSchedulerWeightPreviousResponse), - OpenAIAdvancedSchedulerWeightSessionSticky: stringSetting(req.OpenAIAdvancedSchedulerWeightSessionSticky, previousSettings.OpenAIAdvancedSchedulerWeightSessionSticky), - BalanceLowNotifyEnabled: func() bool { - if req.BalanceLowNotifyEnabled != nil { - return *req.BalanceLowNotifyEnabled - } - return previousSettings.BalanceLowNotifyEnabled - }(), - BalanceLowNotifyThreshold: func() float64 { - if req.BalanceLowNotifyThreshold != nil { - return *req.BalanceLowNotifyThreshold - } - return previousSettings.BalanceLowNotifyThreshold - }(), - BalanceLowNotifyRechargeURL: func() string { - if req.BalanceLowNotifyRechargeURL != nil { - return *req.BalanceLowNotifyRechargeURL - } - return previousSettings.BalanceLowNotifyRechargeURL - }(), - SubscriptionExpiryNotifyEnabled: func() bool { - if req.SubscriptionExpiryNotifyEnabled != nil { - return *req.SubscriptionExpiryNotifyEnabled - } - return previousSettings.SubscriptionExpiryNotifyEnabled - }(), - AccountQuotaNotifyEnabled: func() bool { - if req.AccountQuotaNotifyEnabled != nil { - return *req.AccountQuotaNotifyEnabled - } - return previousSettings.AccountQuotaNotifyEnabled - }(), - AccountQuotaNotifyEmails: func() []service.NotifyEmailEntry { - if req.AccountQuotaNotifyEmails != nil { - return dto.NotifyEmailEntriesToService(*req.AccountQuotaNotifyEmails) - } - return previousSettings.AccountQuotaNotifyEmails - }(), - ChannelMonitorEnabled: func() bool { - if req.ChannelMonitorEnabled != nil { - return *req.ChannelMonitorEnabled - } - return previousSettings.ChannelMonitorEnabled - }(), - ChannelMonitorDefaultIntervalSeconds: func() int { - if req.ChannelMonitorDefaultIntervalSeconds != nil { - return *req.ChannelMonitorDefaultIntervalSeconds - } - return previousSettings.ChannelMonitorDefaultIntervalSeconds - }(), - AvailableChannelsEnabled: func() bool { - if req.AvailableChannelsEnabled != nil { - return *req.AvailableChannelsEnabled - } - return previousSettings.AvailableChannelsEnabled - }(), - AffiliateEnabled: func() bool { - if req.AffiliateEnabled != nil { - return *req.AffiliateEnabled - } - return previousSettings.AffiliateEnabled - }(), - RiskControlEnabled: func() bool { - if req.RiskControlEnabled != nil { - return *req.RiskControlEnabled - } - return previousSettings.RiskControlEnabled - }(), - CyberSessionBlockEnabled: func() bool { - if req.CyberSessionBlockEnabled != nil { - return *req.CyberSessionBlockEnabled - } - return previousSettings.CyberSessionBlockEnabled - }(), - CyberSessionBlockTTLSeconds: func() int { - if req.CyberSessionBlockTTLSeconds != nil { - return *req.CyberSessionBlockTTLSeconds - } - return previousSettings.CyberSessionBlockTTLSeconds - }(), - } - - // req.AuthSourceXxxPlatformQuotas 为 nil 表示本次请求未包含该 source 的 quota 配置(保留 previousAuthSourceDefaults 中的值); - // non-nil(含 empty map)表示整体覆盖:empty map = 清空该 source 的所有 quota 配置。 - authSourceDefaults := &service.AuthSourceDefaultSettings{ - Email: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultEmailBalance, previousAuthSourceDefaults.Email.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultEmailConcurrency, previousAuthSourceDefaults.Email.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultEmailSubscriptions, previousAuthSourceDefaults.Email.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultEmailGrantOnSignup, previousAuthSourceDefaults.Email.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultEmailGrantOnFirstBind, previousAuthSourceDefaults.Email.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceEmailPlatformQuotas, previousAuthSourceDefaults.Email.PlatformQuotas), - }, - LinuxDo: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultLinuxDoBalance, previousAuthSourceDefaults.LinuxDo.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultLinuxDoConcurrency, previousAuthSourceDefaults.LinuxDo.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultLinuxDoSubscriptions, previousAuthSourceDefaults.LinuxDo.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultLinuxDoGrantOnSignup, previousAuthSourceDefaults.LinuxDo.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultLinuxDoGrantOnFirstBind, previousAuthSourceDefaults.LinuxDo.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceLinuxDoPlatformQuotas, previousAuthSourceDefaults.LinuxDo.PlatformQuotas), - }, - OIDC: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultOIDCBalance, previousAuthSourceDefaults.OIDC.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultOIDCConcurrency, previousAuthSourceDefaults.OIDC.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultOIDCSubscriptions, previousAuthSourceDefaults.OIDC.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultOIDCGrantOnSignup, previousAuthSourceDefaults.OIDC.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultOIDCGrantOnFirstBind, previousAuthSourceDefaults.OIDC.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceOIDCPlatformQuotas, previousAuthSourceDefaults.OIDC.PlatformQuotas), - }, - WeChat: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultWeChatBalance, previousAuthSourceDefaults.WeChat.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultWeChatConcurrency, previousAuthSourceDefaults.WeChat.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultWeChatSubscriptions, previousAuthSourceDefaults.WeChat.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultWeChatGrantOnSignup, previousAuthSourceDefaults.WeChat.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultWeChatGrantOnFirstBind, previousAuthSourceDefaults.WeChat.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceWeChatPlatformQuotas, previousAuthSourceDefaults.WeChat.PlatformQuotas), - }, - GitHub: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultGitHubBalance, previousAuthSourceDefaults.GitHub.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultGitHubConcurrency, previousAuthSourceDefaults.GitHub.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultGitHubSubscriptions, previousAuthSourceDefaults.GitHub.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultGitHubGrantOnSignup, previousAuthSourceDefaults.GitHub.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultGitHubGrantOnFirstBind, previousAuthSourceDefaults.GitHub.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceGitHubPlatformQuotas, previousAuthSourceDefaults.GitHub.PlatformQuotas), - }, - Google: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultGoogleBalance, previousAuthSourceDefaults.Google.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultGoogleConcurrency, previousAuthSourceDefaults.Google.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultGoogleSubscriptions, previousAuthSourceDefaults.Google.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultGoogleGrantOnSignup, previousAuthSourceDefaults.Google.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultGoogleGrantOnFirstBind, previousAuthSourceDefaults.Google.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceGooglePlatformQuotas, previousAuthSourceDefaults.Google.PlatformQuotas), - }, - DingTalk: service.ProviderDefaultGrantSettings{ - Balance: float64ValueOrDefault(req.AuthSourceDefaultDingTalkBalance, previousAuthSourceDefaults.DingTalk.Balance), - Concurrency: intValueOrDefault(req.AuthSourceDefaultDingTalkConcurrency, previousAuthSourceDefaults.DingTalk.Concurrency), - Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultDingTalkSubscriptions, previousAuthSourceDefaults.DingTalk.Subscriptions), - GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultDingTalkGrantOnSignup, previousAuthSourceDefaults.DingTalk.GrantOnSignup), - GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultDingTalkGrantOnFirstBind, previousAuthSourceDefaults.DingTalk.GrantOnFirstBind), - PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceDingTalkPlatformQuotas, previousAuthSourceDefaults.DingTalk.PlatformQuotas), - }, - ForceEmailOnThirdPartySignup: boolValueOrDefault(req.ForceEmailOnThirdPartySignup, previousAuthSourceDefaults.ForceEmailOnThirdPartySignup), - } - if err := h.settingService.UpdateSettingsWithAuthSourceDefaults(c.Request.Context(), settings, authSourceDefaults); err != nil { - response.ErrorFrom(c, err) - return - } - - // Update OpenAI fast policy (stored under dedicated key, only when provided). - if req.OpenAIFastPolicySettings != nil { - if err := h.settingService.SetOpenAIFastPolicySettings(c.Request.Context(), openaiFastPolicySettingsFromDTO(req.OpenAIFastPolicySettings)); err != nil { - response.BadRequest(c, err.Error()) - return - } - } - - // Update payment configuration (integrated into system settings). - // Skip if no payment fields were provided (prevents accidental wipe). - if h.paymentConfigService != nil && hasPaymentFields(req) { - paymentReq := service.UpdatePaymentConfigRequest{ - Enabled: req.PaymentEnabled, - MinAmount: req.PaymentMinAmount, - MaxAmount: req.PaymentMaxAmount, - DailyLimit: req.PaymentDailyLimit, - OrderTimeoutMin: req.PaymentOrderTimeoutMin, - MaxPendingOrders: req.PaymentMaxPendingOrders, - EnabledTypes: req.PaymentEnabledTypes, - BalanceDisabled: req.PaymentBalanceDisabled, - BalanceRechargeMultiplier: req.PaymentBalanceRechargeMultiplier, - SubscriptionUSDToCNYRate: req.PaymentSubscriptionUSDToCNYRate, - RechargeFeeRate: req.PaymentRechargeFeeRate, - LoadBalanceStrategy: req.PaymentLoadBalanceStrat, - ProductNamePrefix: req.PaymentProductNamePrefix, - ProductNameSuffix: req.PaymentProductNameSuffix, - HelpImageURL: req.PaymentHelpImageURL, - HelpText: req.PaymentHelpText, - CancelRateLimitEnabled: req.PaymentCancelRateLimitEnabled, - CancelRateLimitMax: req.PaymentCancelRateLimitMax, - CancelRateLimitWindow: req.PaymentCancelRateLimitWindow, - CancelRateLimitUnit: req.PaymentCancelRateLimitUnit, - CancelRateLimitMode: req.PaymentCancelRateLimitMode, - AlipayForceQRCode: req.PaymentAlipayForceQRCode, - } - if err := h.paymentConfigService.UpdatePaymentConfig(c.Request.Context(), paymentReq); err != nil { - response.ErrorFrom(c, err) - return - } - // Refresh in-memory provider registry so config changes take effect immediately - if h.paymentService != nil { - h.paymentService.RefreshProviders(c.Request.Context()) - } - } - - h.auditSettingsUpdate(c, previousSettings, settings, previousAuthSourceDefaults, authSourceDefaults, req) - - // 重新获取设置返回 - updatedSettings, err := h.settingService.GetAllSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - h.ensureDingTalkSyncAttributes(c.Request.Context(), updatedSettings) - updatedAuthSourceDefaults, err := h.settingService.GetAuthSourceDefaultSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - updatedDefaultSubscriptions := make([]dto.DefaultSubscriptionSetting, 0, len(updatedSettings.DefaultSubscriptions)) - for _, sub := range updatedSettings.DefaultSubscriptions { - updatedDefaultSubscriptions = append(updatedDefaultSubscriptions, dto.DefaultSubscriptionSetting{ - GroupID: sub.GroupID, - ValidityDays: sub.ValidityDays, - }) - } - - // Reload payment config for response - var updatedPaymentCfg *service.PaymentConfig - if h.paymentConfigService != nil { - updatedPaymentCfg, _ = h.paymentConfigService.GetPaymentConfig(c.Request.Context()) - } - if updatedPaymentCfg == nil { - updatedPaymentCfg = &service.PaymentConfig{} - } - - payload := dto.SystemSettings{ - RegistrationEnabled: updatedSettings.RegistrationEnabled, - EmailVerifyEnabled: updatedSettings.EmailVerifyEnabled, - RegistrationEmailSuffixWhitelist: updatedSettings.RegistrationEmailSuffixWhitelist, - PromoCodeEnabled: updatedSettings.PromoCodeEnabled, - PasswordResetEnabled: updatedSettings.PasswordResetEnabled, - FrontendURL: updatedSettings.FrontendURL, - InvitationCodeEnabled: updatedSettings.InvitationCodeEnabled, - TotpEnabled: updatedSettings.TotpEnabled, - TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(), - LoginAgreementEnabled: updatedSettings.LoginAgreementEnabled, - LoginAgreementMode: updatedSettings.LoginAgreementMode, - LoginAgreementUpdatedAt: updatedSettings.LoginAgreementUpdatedAt, - LoginAgreementDocuments: loginAgreementDocumentsToDTO(updatedSettings.LoginAgreementDocuments), - SMTPHost: updatedSettings.SMTPHost, - SMTPPort: updatedSettings.SMTPPort, - SMTPUsername: updatedSettings.SMTPUsername, - SMTPPasswordConfigured: updatedSettings.SMTPPasswordConfigured, - SMTPFrom: updatedSettings.SMTPFrom, - SMTPFromName: updatedSettings.SMTPFromName, - SMTPUseTLS: updatedSettings.SMTPUseTLS, - TurnstileEnabled: updatedSettings.TurnstileEnabled, - TurnstileSiteKey: updatedSettings.TurnstileSiteKey, - TurnstileSecretKeyConfigured: updatedSettings.TurnstileSecretKeyConfigured, - APIKeyACLTrustForwardedIP: updatedSettings.APIKeyACLTrustForwardedIP, - LinuxDoConnectEnabled: updatedSettings.LinuxDoConnectEnabled, - LinuxDoConnectClientID: updatedSettings.LinuxDoConnectClientID, - LinuxDoConnectClientSecretConfigured: updatedSettings.LinuxDoConnectClientSecretConfigured, - LinuxDoConnectRedirectURL: updatedSettings.LinuxDoConnectRedirectURL, - DingTalkConnectEnabled: updatedSettings.DingTalkConnectEnabled, - DingTalkConnectClientID: updatedSettings.DingTalkConnectClientID, - DingTalkConnectClientSecretConfigured: updatedSettings.DingTalkConnectClientSecretConfigured, - DingTalkConnectRedirectURL: updatedSettings.DingTalkConnectRedirectURL, - DingTalkConnectCorpRestrictionPolicy: updatedSettings.DingTalkConnectCorpRestrictionPolicy, - DingTalkConnectInternalCorpID: updatedSettings.DingTalkConnectInternalCorpID, - DingTalkConnectBypassRegistration: updatedSettings.DingTalkConnectBypassRegistration, - DingTalkConnectSyncCorpEmail: updatedSettings.DingTalkConnectSyncCorpEmail, - DingTalkConnectSyncDisplayName: updatedSettings.DingTalkConnectSyncDisplayName, - DingTalkConnectSyncDept: updatedSettings.DingTalkConnectSyncDept, - DingTalkConnectSyncCorpEmailAttrKey: updatedSettings.DingTalkConnectSyncCorpEmailAttrKey, - DingTalkConnectSyncDisplayNameAttrKey: updatedSettings.DingTalkConnectSyncDisplayNameAttrKey, - DingTalkConnectSyncDeptAttrKey: updatedSettings.DingTalkConnectSyncDeptAttrKey, - DingTalkConnectSyncCorpEmailAttrName: updatedSettings.DingTalkConnectSyncCorpEmailAttrName, - DingTalkConnectSyncDisplayNameAttrName: updatedSettings.DingTalkConnectSyncDisplayNameAttrName, - DingTalkConnectSyncDeptAttrName: updatedSettings.DingTalkConnectSyncDeptAttrName, - WeChatConnectEnabled: updatedSettings.WeChatConnectEnabled, - WeChatConnectAppID: updatedSettings.WeChatConnectAppID, - WeChatConnectAppSecretConfigured: updatedSettings.WeChatConnectAppSecretConfigured, - WeChatConnectOpenAppID: updatedSettings.WeChatConnectOpenAppID, - WeChatConnectOpenAppSecretConfigured: updatedSettings.WeChatConnectOpenAppSecretConfigured, - WeChatConnectMPAppID: updatedSettings.WeChatConnectMPAppID, - WeChatConnectMPAppSecretConfigured: updatedSettings.WeChatConnectMPAppSecretConfigured, - WeChatConnectMobileAppID: updatedSettings.WeChatConnectMobileAppID, - WeChatConnectMobileAppSecretConfigured: updatedSettings.WeChatConnectMobileAppSecretConfigured, - WeChatConnectOpenEnabled: updatedSettings.WeChatConnectOpenEnabled, - WeChatConnectMPEnabled: updatedSettings.WeChatConnectMPEnabled, - WeChatConnectMobileEnabled: updatedSettings.WeChatConnectMobileEnabled, - WeChatConnectMode: updatedSettings.WeChatConnectMode, - WeChatConnectScopes: updatedSettings.WeChatConnectScopes, - WeChatConnectRedirectURL: updatedSettings.WeChatConnectRedirectURL, - WeChatConnectFrontendRedirectURL: updatedSettings.WeChatConnectFrontendRedirectURL, - OIDCConnectEnabled: updatedSettings.OIDCConnectEnabled, - OIDCConnectProviderName: updatedSettings.OIDCConnectProviderName, - OIDCConnectClientID: updatedSettings.OIDCConnectClientID, - OIDCConnectClientSecretConfigured: updatedSettings.OIDCConnectClientSecretConfigured, - OIDCConnectIssuerURL: updatedSettings.OIDCConnectIssuerURL, - OIDCConnectDiscoveryURL: updatedSettings.OIDCConnectDiscoveryURL, - OIDCConnectAuthorizeURL: updatedSettings.OIDCConnectAuthorizeURL, - OIDCConnectTokenURL: updatedSettings.OIDCConnectTokenURL, - OIDCConnectUserInfoURL: updatedSettings.OIDCConnectUserInfoURL, - OIDCConnectJWKSURL: updatedSettings.OIDCConnectJWKSURL, - OIDCConnectScopes: updatedSettings.OIDCConnectScopes, - OIDCConnectRedirectURL: updatedSettings.OIDCConnectRedirectURL, - OIDCConnectFrontendRedirectURL: updatedSettings.OIDCConnectFrontendRedirectURL, - OIDCConnectTokenAuthMethod: updatedSettings.OIDCConnectTokenAuthMethod, - OIDCConnectUsePKCE: updatedSettings.OIDCConnectUsePKCE, - OIDCConnectValidateIDToken: updatedSettings.OIDCConnectValidateIDToken, - OIDCConnectAllowedSigningAlgs: updatedSettings.OIDCConnectAllowedSigningAlgs, - OIDCConnectClockSkewSeconds: updatedSettings.OIDCConnectClockSkewSeconds, - OIDCConnectRequireEmailVerified: updatedSettings.OIDCConnectRequireEmailVerified, - OIDCConnectUserInfoEmailPath: updatedSettings.OIDCConnectUserInfoEmailPath, - OIDCConnectUserInfoIDPath: updatedSettings.OIDCConnectUserInfoIDPath, - OIDCConnectUserInfoUsernamePath: updatedSettings.OIDCConnectUserInfoUsernamePath, - GitHubOAuthEnabled: updatedSettings.GitHubOAuthEnabled, - GitHubOAuthClientID: updatedSettings.GitHubOAuthClientID, - GitHubOAuthClientSecretConfigured: updatedSettings.GitHubOAuthClientSecretConfigured, - GitHubOAuthRedirectURL: updatedSettings.GitHubOAuthRedirectURL, - GitHubOAuthFrontendRedirectURL: updatedSettings.GitHubOAuthFrontendRedirectURL, - GoogleOAuthEnabled: updatedSettings.GoogleOAuthEnabled, - GoogleOAuthClientID: updatedSettings.GoogleOAuthClientID, - GoogleOAuthClientSecretConfigured: updatedSettings.GoogleOAuthClientSecretConfigured, - GoogleOAuthRedirectURL: updatedSettings.GoogleOAuthRedirectURL, - GoogleOAuthFrontendRedirectURL: updatedSettings.GoogleOAuthFrontendRedirectURL, - SiteName: updatedSettings.SiteName, - SiteLogo: updatedSettings.SiteLogo, - SiteSubtitle: updatedSettings.SiteSubtitle, - APIBaseURL: updatedSettings.APIBaseURL, - ContactInfo: updatedSettings.ContactInfo, - DocURL: updatedSettings.DocURL, - HomeContent: updatedSettings.HomeContent, - HideCcsImportButton: updatedSettings.HideCcsImportButton, - PurchaseSubscriptionEnabled: updatedSettings.PurchaseSubscriptionEnabled, - PurchaseSubscriptionURL: updatedSettings.PurchaseSubscriptionURL, - TableDefaultPageSize: updatedSettings.TableDefaultPageSize, - TablePageSizeOptions: updatedSettings.TablePageSizeOptions, - CustomMenuItems: dto.ParseCustomMenuItems(updatedSettings.CustomMenuItems), - CustomEndpoints: dto.ParseCustomEndpoints(updatedSettings.CustomEndpoints), - DefaultConcurrency: updatedSettings.DefaultConcurrency, - DefaultBalance: updatedSettings.DefaultBalance, - AffiliateRebateRate: updatedSettings.AffiliateRebateRate, - AffiliateRebateFreezeHours: updatedSettings.AffiliateRebateFreezeHours, - AffiliateRebateDurationDays: updatedSettings.AffiliateRebateDurationDays, - AffiliateRebatePerInviteeCap: updatedSettings.AffiliateRebatePerInviteeCap, - DefaultUserRPMLimit: updatedSettings.DefaultUserRPMLimit, - DefaultSubscriptions: updatedDefaultSubscriptions, - EnableModelFallback: updatedSettings.EnableModelFallback, - FallbackModelAnthropic: updatedSettings.FallbackModelAnthropic, - FallbackModelOpenAI: updatedSettings.FallbackModelOpenAI, - FallbackModelGemini: updatedSettings.FallbackModelGemini, - FallbackModelAntigravity: updatedSettings.FallbackModelAntigravity, - EnableIdentityPatch: updatedSettings.EnableIdentityPatch, - IdentityPatchPrompt: updatedSettings.IdentityPatchPrompt, - OpsMonitoringEnabled: updatedSettings.OpsMonitoringEnabled, - OpsRealtimeMonitoringEnabled: updatedSettings.OpsRealtimeMonitoringEnabled, - OpsQueryModeDefault: updatedSettings.OpsQueryModeDefault, - OpsMetricsIntervalSeconds: updatedSettings.OpsMetricsIntervalSeconds, - MinClaudeCodeVersion: updatedSettings.MinClaudeCodeVersion, - MaxClaudeCodeVersion: updatedSettings.MaxClaudeCodeVersion, - AllowUngroupedKeyScheduling: updatedSettings.AllowUngroupedKeyScheduling, - BackendModeEnabled: updatedSettings.BackendModeEnabled, - EnableFingerprintUnification: updatedSettings.EnableFingerprintUnification, - EnableMetadataPassthrough: updatedSettings.EnableMetadataPassthrough, - EnableCCHSigning: updatedSettings.EnableCCHSigning, - EnableClaudeOAuthSystemPromptInjection: updatedSettings.EnableClaudeOAuthSystemPromptInjection, - ClaudeOAuthSystemPrompt: updatedSettings.ClaudeOAuthSystemPrompt, - ClaudeOAuthSystemPromptBlocks: updatedSettings.ClaudeOAuthSystemPromptBlocks, - EnableAnthropicCacheTTL1hInjection: updatedSettings.EnableAnthropicCacheTTL1hInjection, - RewriteMessageCacheControl: updatedSettings.RewriteMessageCacheControl, - EnableClientDatelineNormalization: updatedSettings.EnableClientDatelineNormalization, - AntigravityUserAgentVersion: updatedSettings.AntigravityUserAgentVersion, - OpenAICodexUserAgent: updatedSettings.OpenAICodexUserAgent, - MinCodexVersion: updatedSettings.MinCodexVersion, - MaxCodexVersion: updatedSettings.MaxCodexVersion, - CodexCLIOnlyBlacklist: updatedSettings.CodexCLIOnlyBlacklist, - CodexCLIOnlyWhitelist: updatedSettings.CodexCLIOnlyWhitelist, - CodexCLIOnlyAllowAppServerClients: updatedSettings.CodexCLIOnlyAllowAppServerClients, - CodexCLIOnlyEngineFingerprintSignals: updatedSettings.CodexCLIOnlyEngineFingerprintSignals, - PaymentVisibleMethodAlipaySource: updatedSettings.PaymentVisibleMethodAlipaySource, - PaymentVisibleMethodWxpaySource: updatedSettings.PaymentVisibleMethodWxpaySource, - PaymentVisibleMethodAlipayEnabled: updatedSettings.PaymentVisibleMethodAlipayEnabled, - PaymentVisibleMethodWxpayEnabled: updatedSettings.PaymentVisibleMethodWxpayEnabled, - OpenAIAdvancedSchedulerEnabled: updatedSettings.OpenAIAdvancedSchedulerEnabled, - OpenAIAdvancedSchedulerStickyWeightedEnabled: updatedSettings.OpenAIAdvancedSchedulerStickyWeightedEnabled, - OpenAIAdvancedSchedulerSubscriptionPriorityEnabled: updatedSettings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled, - OpenAIAdvancedSchedulerLBTopK: updatedSettings.OpenAIAdvancedSchedulerLBTopK, - OpenAIAdvancedSchedulerWeightPriority: updatedSettings.OpenAIAdvancedSchedulerWeightPriority, - OpenAIAdvancedSchedulerWeightLoad: updatedSettings.OpenAIAdvancedSchedulerWeightLoad, - OpenAIAdvancedSchedulerWeightQueue: updatedSettings.OpenAIAdvancedSchedulerWeightQueue, - OpenAIAdvancedSchedulerWeightErrorRate: updatedSettings.OpenAIAdvancedSchedulerWeightErrorRate, - OpenAIAdvancedSchedulerWeightTTFT: updatedSettings.OpenAIAdvancedSchedulerWeightTTFT, - OpenAIAdvancedSchedulerWeightReset: updatedSettings.OpenAIAdvancedSchedulerWeightReset, - OpenAIAdvancedSchedulerWeightQuotaHeadroom: updatedSettings.OpenAIAdvancedSchedulerWeightQuotaHeadroom, - OpenAIAdvancedSchedulerWeightPreviousResponse: updatedSettings.OpenAIAdvancedSchedulerWeightPreviousResponse, - OpenAIAdvancedSchedulerWeightSessionSticky: updatedSettings.OpenAIAdvancedSchedulerWeightSessionSticky, - OpenAIAdvancedSchedulerEffectiveLBTopK: updatedSettings.OpenAIAdvancedSchedulerEffectiveLBTopK, - OpenAIAdvancedSchedulerEffectiveWeightPriority: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightPriority, - OpenAIAdvancedSchedulerEffectiveWeightLoad: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightLoad, - OpenAIAdvancedSchedulerEffectiveWeightQueue: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightQueue, - OpenAIAdvancedSchedulerEffectiveWeightErrorRate: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightErrorRate, - OpenAIAdvancedSchedulerEffectiveWeightTTFT: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightTTFT, - OpenAIAdvancedSchedulerEffectiveWeightReset: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightReset, - OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom, - OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse, - OpenAIAdvancedSchedulerEffectiveWeightSessionSticky: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightSessionSticky, - BalanceLowNotifyEnabled: updatedSettings.BalanceLowNotifyEnabled, - BalanceLowNotifyThreshold: updatedSettings.BalanceLowNotifyThreshold, - BalanceLowNotifyRechargeURL: updatedSettings.BalanceLowNotifyRechargeURL, - SubscriptionExpiryNotifyEnabled: updatedSettings.SubscriptionExpiryNotifyEnabled, - AccountQuotaNotifyEnabled: updatedSettings.AccountQuotaNotifyEnabled, - AccountQuotaNotifyEmails: dto.NotifyEmailEntriesFromService(updatedSettings.AccountQuotaNotifyEmails), - PaymentEnabled: updatedPaymentCfg.Enabled, - PaymentMinAmount: updatedPaymentCfg.MinAmount, - PaymentMaxAmount: updatedPaymentCfg.MaxAmount, - PaymentDailyLimit: updatedPaymentCfg.DailyLimit, - PaymentOrderTimeoutMin: updatedPaymentCfg.OrderTimeoutMin, - PaymentMaxPendingOrders: updatedPaymentCfg.MaxPendingOrders, - PaymentEnabledTypes: updatedPaymentCfg.EnabledTypes, - PaymentBalanceDisabled: updatedPaymentCfg.BalanceDisabled, - PaymentBalanceRechargeMultiplier: updatedPaymentCfg.BalanceRechargeMultiplier, - PaymentSubscriptionUSDToCNYRate: updatedPaymentCfg.SubscriptionUSDToCNYRate, - PaymentRechargeFeeRate: updatedPaymentCfg.RechargeFeeRate, - PaymentLoadBalanceStrat: updatedPaymentCfg.LoadBalanceStrategy, - PaymentProductNamePrefix: updatedPaymentCfg.ProductNamePrefix, - PaymentProductNameSuffix: updatedPaymentCfg.ProductNameSuffix, - PaymentHelpImageURL: updatedPaymentCfg.HelpImageURL, - PaymentHelpText: updatedPaymentCfg.HelpText, - PaymentCancelRateLimitEnabled: updatedPaymentCfg.CancelRateLimitEnabled, - PaymentCancelRateLimitMax: updatedPaymentCfg.CancelRateLimitMax, - PaymentCancelRateLimitWindow: updatedPaymentCfg.CancelRateLimitWindow, - PaymentCancelRateLimitUnit: updatedPaymentCfg.CancelRateLimitUnit, - PaymentCancelRateLimitMode: updatedPaymentCfg.CancelRateLimitMode, - PaymentAlipayForceQRCode: updatedPaymentCfg.AlipayForceQRCode, - - ChannelMonitorEnabled: updatedSettings.ChannelMonitorEnabled, - ChannelMonitorDefaultIntervalSeconds: updatedSettings.ChannelMonitorDefaultIntervalSeconds, - - AvailableChannelsEnabled: updatedSettings.AvailableChannelsEnabled, - - AffiliateEnabled: updatedSettings.AffiliateEnabled, - - RiskControlEnabled: updatedSettings.RiskControlEnabled, - CyberSessionBlockEnabled: updatedSettings.CyberSessionBlockEnabled, - CyberSessionBlockTTLSeconds: updatedSettings.CyberSessionBlockTTLSeconds, - AllowUserViewErrorRequests: updatedSettings.AllowUserViewErrorRequests, - } - if fastPolicy, err := h.settingService.GetOpenAIFastPolicySettings(c.Request.Context()); err != nil { - slog.Error("openai_fast_policy_settings_get_failed", "error", err) - } else if fastPolicy != nil { - payload.OpenAIFastPolicySettings = openaiFastPolicySettingsToDTO(fastPolicy) - } - - // Default platform quotas(JSON map)—— 与 GetSettings 一致,避免保存后响应缺失该字段 - if platformQuotas, err := h.settingService.GetDefaultPlatformQuotas(c.Request.Context()); err != nil { - slog.Error("default_platform_quotas_get_failed", "error", err) - } else { - payload.DefaultPlatformQuotas = platformQuotas - } - response.Success(c, systemSettingsResponseData(payload, updatedAuthSourceDefaults)) -} - -// hasPaymentFields returns true if any payment-related field was explicitly provided. -// mapDingTalkValidateError maps ValidateDingTalkConfig errors to machine-readable reason codes. -func mapDingTalkValidateError(err error) string { - switch { - case errors.Is(err, config.ErrDingTalkV1AppTypeMismatch): - return "dingtalk_apptype_mismatch" - case errors.Is(err, config.ErrDingTalkV4InvalidAppKind): - return "dingtalk_app_kind_invalid" - default: - return "dingtalk_corp_config_invalid" - } -} - -func hasPaymentFields(req UpdateSettingsRequest) bool { - return req.PaymentEnabled != nil || req.PaymentMinAmount != nil || - req.PaymentMaxAmount != nil || req.PaymentDailyLimit != nil || - req.PaymentOrderTimeoutMin != nil || req.PaymentMaxPendingOrders != nil || - req.PaymentEnabledTypes != nil || req.PaymentBalanceDisabled != nil || - req.PaymentBalanceRechargeMultiplier != nil || req.PaymentSubscriptionUSDToCNYRate != nil || - req.PaymentRechargeFeeRate != nil || - req.PaymentLoadBalanceStrat != nil || req.PaymentProductNamePrefix != nil || - req.PaymentProductNameSuffix != nil || req.PaymentHelpImageURL != nil || - req.PaymentHelpText != nil || req.PaymentCancelRateLimitEnabled != nil || - req.PaymentCancelRateLimitMax != nil || req.PaymentCancelRateLimitWindow != nil || - req.PaymentCancelRateLimitUnit != nil || req.PaymentCancelRateLimitMode != nil || - req.PaymentAlipayForceQRCode != nil -} - -func (h *SettingHandler) auditSettingsUpdate(c *gin.Context, before *service.SystemSettings, after *service.SystemSettings, beforeAuthSourceDefaults *service.AuthSourceDefaultSettings, afterAuthSourceDefaults *service.AuthSourceDefaultSettings, req UpdateSettingsRequest) { - if before == nil || after == nil { - return - } - - changed := diffSettings(before, after, beforeAuthSourceDefaults, afterAuthSourceDefaults, req) - if len(changed) == 0 { - return - } - - subject, _ := middleware.GetAuthSubjectFromContext(c) - role, _ := middleware.GetUserRoleFromContext(c) - slog.Info("settings updated", - "audit", true, - "user_id", subject.UserID, - "role", role, - "changed", changed, - ) -} - -func diffSettings(before *service.SystemSettings, after *service.SystemSettings, beforeAuthSourceDefaults *service.AuthSourceDefaultSettings, afterAuthSourceDefaults *service.AuthSourceDefaultSettings, req UpdateSettingsRequest) []string { - changed := make([]string, 0, 20) - if before.RegistrationEnabled != after.RegistrationEnabled { - changed = append(changed, "registration_enabled") - } - if before.EmailVerifyEnabled != after.EmailVerifyEnabled { - changed = append(changed, "email_verify_enabled") - } - if !equalStringSlice(before.RegistrationEmailSuffixWhitelist, after.RegistrationEmailSuffixWhitelist) { - changed = append(changed, "registration_email_suffix_whitelist") - } - if before.PromoCodeEnabled != after.PromoCodeEnabled { - changed = append(changed, "promo_code_enabled") - } - if before.InvitationCodeEnabled != after.InvitationCodeEnabled { - changed = append(changed, "invitation_code_enabled") - } - if before.PasswordResetEnabled != after.PasswordResetEnabled { - changed = append(changed, "password_reset_enabled") - } - if before.FrontendURL != after.FrontendURL { - changed = append(changed, "frontend_url") - } - if before.TotpEnabled != after.TotpEnabled { - changed = append(changed, "totp_enabled") - } - if before.LoginAgreementEnabled != after.LoginAgreementEnabled { - changed = append(changed, "login_agreement_enabled") - } - if before.LoginAgreementMode != after.LoginAgreementMode { - changed = append(changed, "login_agreement_mode") - } - if before.LoginAgreementUpdatedAt != after.LoginAgreementUpdatedAt { - changed = append(changed, "login_agreement_updated_at") - } - if !equalLoginAgreementDocuments(before.LoginAgreementDocuments, after.LoginAgreementDocuments) { - changed = append(changed, "login_agreement_documents") - } - if before.SMTPHost != after.SMTPHost { - changed = append(changed, "smtp_host") - } - if before.SMTPPort != after.SMTPPort { - changed = append(changed, "smtp_port") - } - if before.SMTPUsername != after.SMTPUsername { - changed = append(changed, "smtp_username") - } - if req.SMTPPassword != "" { - changed = append(changed, "smtp_password") - } - if before.SMTPFrom != after.SMTPFrom { - changed = append(changed, "smtp_from_email") - } - if before.SMTPFromName != after.SMTPFromName { - changed = append(changed, "smtp_from_name") - } - if before.SMTPUseTLS != after.SMTPUseTLS { - changed = append(changed, "smtp_use_tls") - } - if before.TurnstileEnabled != after.TurnstileEnabled { - changed = append(changed, "turnstile_enabled") - } - if before.TurnstileSiteKey != after.TurnstileSiteKey { - changed = append(changed, "turnstile_site_key") - } - if req.TurnstileSecretKey != "" { - changed = append(changed, "turnstile_secret_key") - } - if before.APIKeyACLTrustForwardedIP != after.APIKeyACLTrustForwardedIP { - changed = append(changed, "api_key_acl_trust_forwarded_ip") - } - if before.LinuxDoConnectEnabled != after.LinuxDoConnectEnabled { - changed = append(changed, "linuxdo_connect_enabled") - } - if before.LinuxDoConnectClientID != after.LinuxDoConnectClientID { - changed = append(changed, "linuxdo_connect_client_id") - } - if req.LinuxDoConnectClientSecret != "" { - changed = append(changed, "linuxdo_connect_client_secret") - } - if before.LinuxDoConnectRedirectURL != after.LinuxDoConnectRedirectURL { - changed = append(changed, "linuxdo_connect_redirect_url") - } - if before.DingTalkConnectEnabled != after.DingTalkConnectEnabled { - changed = append(changed, "dingtalk_connect_enabled") - } - if before.DingTalkConnectClientID != after.DingTalkConnectClientID { - changed = append(changed, "dingtalk_connect_client_id") - } - if req.DingTalkConnectClientSecret != "" { - changed = append(changed, "dingtalk_connect_client_secret") - } - if before.DingTalkConnectRedirectURL != after.DingTalkConnectRedirectURL { - changed = append(changed, "dingtalk_connect_redirect_url") - } - if before.DingTalkConnectCorpRestrictionPolicy != after.DingTalkConnectCorpRestrictionPolicy { - changed = append(changed, "dingtalk_connect_corp_restriction_policy") - } - if before.DingTalkConnectInternalCorpID != after.DingTalkConnectInternalCorpID { - changed = append(changed, "dingtalk_connect_internal_corp_id") - } - if before.DingTalkConnectBypassRegistration != after.DingTalkConnectBypassRegistration { - changed = append(changed, "dingtalk_connect_bypass_registration") - } - if before.DingTalkConnectSyncCorpEmail != after.DingTalkConnectSyncCorpEmail { - changed = append(changed, "dingtalk_connect_sync_corp_email") - } - if before.DingTalkConnectSyncDisplayName != after.DingTalkConnectSyncDisplayName { - changed = append(changed, "dingtalk_connect_sync_display_name") - } - if before.DingTalkConnectSyncDept != after.DingTalkConnectSyncDept { - changed = append(changed, "dingtalk_connect_sync_dept") - } - if before.DingTalkConnectSyncCorpEmailAttrKey != after.DingTalkConnectSyncCorpEmailAttrKey { - changed = append(changed, "dingtalk_connect_sync_corp_email_attr_key") - } - if before.DingTalkConnectSyncDisplayNameAttrKey != after.DingTalkConnectSyncDisplayNameAttrKey { - changed = append(changed, "dingtalk_connect_sync_display_name_attr_key") - } - if before.DingTalkConnectSyncDeptAttrKey != after.DingTalkConnectSyncDeptAttrKey { - changed = append(changed, "dingtalk_connect_sync_dept_attr_key") - } - if before.WeChatConnectEnabled != after.WeChatConnectEnabled { - changed = append(changed, "wechat_connect_enabled") - } - if before.WeChatConnectAppID != after.WeChatConnectAppID { - changed = append(changed, "wechat_connect_app_id") - } - if req.WeChatConnectAppSecret != "" { - changed = append(changed, "wechat_connect_app_secret") - } - if before.WeChatConnectOpenAppID != after.WeChatConnectOpenAppID { - changed = append(changed, "wechat_connect_open_app_id") - } - if req.WeChatConnectOpenAppSecret != "" { - changed = append(changed, "wechat_connect_open_app_secret") - } - if before.WeChatConnectMPAppID != after.WeChatConnectMPAppID { - changed = append(changed, "wechat_connect_mp_app_id") - } - if req.WeChatConnectMPAppSecret != "" { - changed = append(changed, "wechat_connect_mp_app_secret") - } - if before.WeChatConnectMobileAppID != after.WeChatConnectMobileAppID { - changed = append(changed, "wechat_connect_mobile_app_id") - } - if req.WeChatConnectMobileAppSecret != "" { - changed = append(changed, "wechat_connect_mobile_app_secret") - } - if before.WeChatConnectOpenEnabled != after.WeChatConnectOpenEnabled { - changed = append(changed, "wechat_connect_open_enabled") - } - if before.WeChatConnectMPEnabled != after.WeChatConnectMPEnabled { - changed = append(changed, "wechat_connect_mp_enabled") - } - if before.WeChatConnectMobileEnabled != after.WeChatConnectMobileEnabled { - changed = append(changed, "wechat_connect_mobile_enabled") - } - if before.WeChatConnectMode != after.WeChatConnectMode { - changed = append(changed, "wechat_connect_mode") - } - if before.WeChatConnectScopes != after.WeChatConnectScopes { - changed = append(changed, "wechat_connect_scopes") - } - if before.WeChatConnectRedirectURL != after.WeChatConnectRedirectURL { - changed = append(changed, "wechat_connect_redirect_url") - } - if before.WeChatConnectFrontendRedirectURL != after.WeChatConnectFrontendRedirectURL { - changed = append(changed, "wechat_connect_frontend_redirect_url") - } - if before.OIDCConnectEnabled != after.OIDCConnectEnabled { - changed = append(changed, "oidc_connect_enabled") - } - if before.OIDCConnectProviderName != after.OIDCConnectProviderName { - changed = append(changed, "oidc_connect_provider_name") - } - if before.OIDCConnectClientID != after.OIDCConnectClientID { - changed = append(changed, "oidc_connect_client_id") - } - if req.OIDCConnectClientSecret != "" { - changed = append(changed, "oidc_connect_client_secret") - } - if before.OIDCConnectIssuerURL != after.OIDCConnectIssuerURL { - changed = append(changed, "oidc_connect_issuer_url") - } - if before.OIDCConnectDiscoveryURL != after.OIDCConnectDiscoveryURL { - changed = append(changed, "oidc_connect_discovery_url") - } - if before.OIDCConnectAuthorizeURL != after.OIDCConnectAuthorizeURL { - changed = append(changed, "oidc_connect_authorize_url") - } - if before.OIDCConnectTokenURL != after.OIDCConnectTokenURL { - changed = append(changed, "oidc_connect_token_url") - } - if before.OIDCConnectUserInfoURL != after.OIDCConnectUserInfoURL { - changed = append(changed, "oidc_connect_userinfo_url") - } - if before.OIDCConnectJWKSURL != after.OIDCConnectJWKSURL { - changed = append(changed, "oidc_connect_jwks_url") - } - if before.OIDCConnectScopes != after.OIDCConnectScopes { - changed = append(changed, "oidc_connect_scopes") - } - if before.OIDCConnectRedirectURL != after.OIDCConnectRedirectURL { - changed = append(changed, "oidc_connect_redirect_url") - } - if before.OIDCConnectFrontendRedirectURL != after.OIDCConnectFrontendRedirectURL { - changed = append(changed, "oidc_connect_frontend_redirect_url") - } - if before.OIDCConnectTokenAuthMethod != after.OIDCConnectTokenAuthMethod { - changed = append(changed, "oidc_connect_token_auth_method") - } - if before.OIDCConnectUsePKCE != after.OIDCConnectUsePKCE { - changed = append(changed, "oidc_connect_use_pkce") - } - if before.OIDCConnectValidateIDToken != after.OIDCConnectValidateIDToken { - changed = append(changed, "oidc_connect_validate_id_token") - } - if before.OIDCConnectAllowedSigningAlgs != after.OIDCConnectAllowedSigningAlgs { - changed = append(changed, "oidc_connect_allowed_signing_algs") - } - if before.OIDCConnectClockSkewSeconds != after.OIDCConnectClockSkewSeconds { - changed = append(changed, "oidc_connect_clock_skew_seconds") - } - if before.OIDCConnectRequireEmailVerified != after.OIDCConnectRequireEmailVerified { - changed = append(changed, "oidc_connect_require_email_verified") - } - if before.OIDCConnectUserInfoEmailPath != after.OIDCConnectUserInfoEmailPath { - changed = append(changed, "oidc_connect_userinfo_email_path") - } - if before.OIDCConnectUserInfoIDPath != after.OIDCConnectUserInfoIDPath { - changed = append(changed, "oidc_connect_userinfo_id_path") - } - if before.OIDCConnectUserInfoUsernamePath != after.OIDCConnectUserInfoUsernamePath { - changed = append(changed, "oidc_connect_userinfo_username_path") - } - if before.SiteName != after.SiteName { - changed = append(changed, "site_name") - } - if before.SiteLogo != after.SiteLogo { - changed = append(changed, "site_logo") - } - if before.SiteSubtitle != after.SiteSubtitle { - changed = append(changed, "site_subtitle") - } - if before.APIBaseURL != after.APIBaseURL { - changed = append(changed, "api_base_url") - } - if before.ContactInfo != after.ContactInfo { - changed = append(changed, "contact_info") - } - if before.DocURL != after.DocURL { - changed = append(changed, "doc_url") - } - if before.HomeContent != after.HomeContent { - changed = append(changed, "home_content") - } - if before.HideCcsImportButton != after.HideCcsImportButton { - changed = append(changed, "hide_ccs_import_button") - } - if before.DefaultConcurrency != after.DefaultConcurrency { - changed = append(changed, "default_concurrency") - } - if before.DefaultBalance != after.DefaultBalance { - changed = append(changed, "default_balance") - } - if before.AffiliateRebateRate != after.AffiliateRebateRate { - changed = append(changed, "affiliate_rebate_rate") - } - if before.AffiliateRebateFreezeHours != after.AffiliateRebateFreezeHours { - changed = append(changed, "affiliate_rebate_freeze_hours") - } - if before.AffiliateRebateDurationDays != after.AffiliateRebateDurationDays { - changed = append(changed, "affiliate_rebate_duration_days") - } - if before.AffiliateRebatePerInviteeCap != after.AffiliateRebatePerInviteeCap { - changed = append(changed, "affiliate_rebate_per_invitee_cap") - } - if !equalDefaultSubscriptions(before.DefaultSubscriptions, after.DefaultSubscriptions) { - changed = append(changed, "default_subscriptions") - } - if before.EnableModelFallback != after.EnableModelFallback { - changed = append(changed, "enable_model_fallback") - } - if before.FallbackModelAnthropic != after.FallbackModelAnthropic { - changed = append(changed, "fallback_model_anthropic") - } - if before.FallbackModelOpenAI != after.FallbackModelOpenAI { - changed = append(changed, "fallback_model_openai") - } - if before.FallbackModelGemini != after.FallbackModelGemini { - changed = append(changed, "fallback_model_gemini") - } - if before.FallbackModelAntigravity != after.FallbackModelAntigravity { - changed = append(changed, "fallback_model_antigravity") - } - if before.EnableIdentityPatch != after.EnableIdentityPatch { - changed = append(changed, "enable_identity_patch") - } - if before.IdentityPatchPrompt != after.IdentityPatchPrompt { - changed = append(changed, "identity_patch_prompt") - } - if before.OpsMonitoringEnabled != after.OpsMonitoringEnabled { - changed = append(changed, "ops_monitoring_enabled") - } - if before.OpsRealtimeMonitoringEnabled != after.OpsRealtimeMonitoringEnabled { - changed = append(changed, "ops_realtime_monitoring_enabled") - } - if before.OpsQueryModeDefault != after.OpsQueryModeDefault { - changed = append(changed, "ops_query_mode_default") - } - if before.OpsMetricsIntervalSeconds != after.OpsMetricsIntervalSeconds { - changed = append(changed, "ops_metrics_interval_seconds") - } - if before.MinClaudeCodeVersion != after.MinClaudeCodeVersion { - changed = append(changed, "min_claude_code_version") - } - if before.MaxClaudeCodeVersion != after.MaxClaudeCodeVersion { - changed = append(changed, "max_claude_code_version") - } - if before.MinCodexVersion != after.MinCodexVersion { - changed = append(changed, "min_codex_version") - } - if before.MaxCodexVersion != after.MaxCodexVersion { - changed = append(changed, "max_codex_version") - } - if before.CodexCLIOnlyAllowAppServerClients != after.CodexCLIOnlyAllowAppServerClients { - changed = append(changed, "codex_cli_only_allow_app_server_clients") - } - if before.CodexCLIOnlyEngineFingerprintSignals != after.CodexCLIOnlyEngineFingerprintSignals { - changed = append(changed, "codex_cli_only_engine_fingerprint_signals") - } - if before.CodexCLIOnlyBlacklist != after.CodexCLIOnlyBlacklist { - changed = append(changed, "codex_cli_only_blacklist") - } - if before.CodexCLIOnlyWhitelist != after.CodexCLIOnlyWhitelist { - changed = append(changed, "codex_cli_only_whitelist") - } - if before.AllowUngroupedKeyScheduling != after.AllowUngroupedKeyScheduling { - changed = append(changed, "allow_ungrouped_key_scheduling") - } - if before.BackendModeEnabled != after.BackendModeEnabled { - changed = append(changed, "backend_mode_enabled") - } - if before.PurchaseSubscriptionEnabled != after.PurchaseSubscriptionEnabled { - changed = append(changed, "purchase_subscription_enabled") - } - if before.PurchaseSubscriptionURL != after.PurchaseSubscriptionURL { - changed = append(changed, "purchase_subscription_url") - } - if before.TableDefaultPageSize != after.TableDefaultPageSize { - changed = append(changed, "table_default_page_size") - } - if !equalIntSlice(before.TablePageSizeOptions, after.TablePageSizeOptions) { - changed = append(changed, "table_page_size_options") - } - if before.CustomMenuItems != after.CustomMenuItems { - changed = append(changed, "custom_menu_items") - } - if before.CustomEndpoints != after.CustomEndpoints { - changed = append(changed, "custom_endpoints") - } - if before.EnableFingerprintUnification != after.EnableFingerprintUnification { - changed = append(changed, "enable_fingerprint_unification") - } - if before.EnableMetadataPassthrough != after.EnableMetadataPassthrough { - changed = append(changed, "enable_metadata_passthrough") - } - if before.EnableCCHSigning != after.EnableCCHSigning { - changed = append(changed, "enable_cch_signing") - } - if before.EnableClaudeOAuthSystemPromptInjection != after.EnableClaudeOAuthSystemPromptInjection { - changed = append(changed, "enable_claude_oauth_system_prompt_injection") - } - if before.ClaudeOAuthSystemPrompt != after.ClaudeOAuthSystemPrompt { - changed = append(changed, "claude_oauth_system_prompt") - } - if before.ClaudeOAuthSystemPromptBlocks != after.ClaudeOAuthSystemPromptBlocks { - changed = append(changed, "claude_oauth_system_prompt_blocks") - } - if before.EnableAnthropicCacheTTL1hInjection != after.EnableAnthropicCacheTTL1hInjection { - changed = append(changed, "enable_anthropic_cache_ttl_1h_injection") - } - if before.RewriteMessageCacheControl != after.RewriteMessageCacheControl { - changed = append(changed, "rewrite_message_cache_control") - } - if before.EnableClientDatelineNormalization != after.EnableClientDatelineNormalization { - changed = append(changed, "enable_client_dateline_normalization") - } - if before.AntigravityUserAgentVersion != after.AntigravityUserAgentVersion { - changed = append(changed, "antigravity_user_agent_version") - } - if before.OpenAICodexUserAgent != after.OpenAICodexUserAgent { - changed = append(changed, "openai_codex_user_agent") - } - if before.PaymentVisibleMethodAlipaySource != after.PaymentVisibleMethodAlipaySource { - changed = append(changed, "payment_visible_method_alipay_source") - } - if before.PaymentVisibleMethodWxpaySource != after.PaymentVisibleMethodWxpaySource { - changed = append(changed, "payment_visible_method_wxpay_source") - } - if before.PaymentVisibleMethodAlipayEnabled != after.PaymentVisibleMethodAlipayEnabled { - changed = append(changed, "payment_visible_method_alipay_enabled") - } - if before.PaymentVisibleMethodWxpayEnabled != after.PaymentVisibleMethodWxpayEnabled { - changed = append(changed, "payment_visible_method_wxpay_enabled") - } - if before.OpenAIAdvancedSchedulerEnabled != after.OpenAIAdvancedSchedulerEnabled { - changed = append(changed, "openai_advanced_scheduler_enabled") - } - if before.OpenAIAdvancedSchedulerStickyWeightedEnabled != after.OpenAIAdvancedSchedulerStickyWeightedEnabled { - changed = append(changed, "openai_advanced_scheduler_sticky_weighted_enabled") - } - if before.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled != after.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled { - changed = append(changed, "openai_advanced_scheduler_subscription_priority_enabled") - } - if before.OpenAIAdvancedSchedulerLBTopK != after.OpenAIAdvancedSchedulerLBTopK { - changed = append(changed, "openai_advanced_scheduler_lb_top_k") - } - if before.OpenAIAdvancedSchedulerWeightPriority != after.OpenAIAdvancedSchedulerWeightPriority { - changed = append(changed, "openai_advanced_scheduler_weight_priority") - } - if before.OpenAIAdvancedSchedulerWeightLoad != after.OpenAIAdvancedSchedulerWeightLoad { - changed = append(changed, "openai_advanced_scheduler_weight_load") - } - if before.OpenAIAdvancedSchedulerWeightQueue != after.OpenAIAdvancedSchedulerWeightQueue { - changed = append(changed, "openai_advanced_scheduler_weight_queue") - } - if before.OpenAIAdvancedSchedulerWeightErrorRate != after.OpenAIAdvancedSchedulerWeightErrorRate { - changed = append(changed, "openai_advanced_scheduler_weight_error_rate") - } - if before.OpenAIAdvancedSchedulerWeightTTFT != after.OpenAIAdvancedSchedulerWeightTTFT { - changed = append(changed, "openai_advanced_scheduler_weight_ttft") - } - if before.OpenAIAdvancedSchedulerWeightReset != after.OpenAIAdvancedSchedulerWeightReset { - changed = append(changed, "openai_advanced_scheduler_weight_reset") - } - if before.OpenAIAdvancedSchedulerWeightQuotaHeadroom != after.OpenAIAdvancedSchedulerWeightQuotaHeadroom { - changed = append(changed, "openai_advanced_scheduler_weight_quota_headroom") - } - if before.OpenAIAdvancedSchedulerWeightPreviousResponse != after.OpenAIAdvancedSchedulerWeightPreviousResponse { - changed = append(changed, "openai_advanced_scheduler_weight_previous_response") - } - if before.OpenAIAdvancedSchedulerWeightSessionSticky != after.OpenAIAdvancedSchedulerWeightSessionSticky { - changed = append(changed, "openai_advanced_scheduler_weight_session_sticky") - } - // 余额、订阅到期与账号限额通知 - if before.BalanceLowNotifyEnabled != after.BalanceLowNotifyEnabled { - changed = append(changed, "balance_low_notify_enabled") - } - if before.BalanceLowNotifyThreshold != after.BalanceLowNotifyThreshold { - changed = append(changed, "balance_low_notify_threshold") - } - if before.BalanceLowNotifyRechargeURL != after.BalanceLowNotifyRechargeURL { - changed = append(changed, "balance_low_notify_recharge_url") - } - if before.SubscriptionExpiryNotifyEnabled != after.SubscriptionExpiryNotifyEnabled { - changed = append(changed, "subscription_expiry_notify_enabled") - } - if before.AccountQuotaNotifyEnabled != after.AccountQuotaNotifyEnabled { - changed = append(changed, "account_quota_notify_enabled") - } - if !equalNotifyEmailEntries(before.AccountQuotaNotifyEmails, after.AccountQuotaNotifyEmails) { - changed = append(changed, "account_quota_notify_emails") - } - if before.ChannelMonitorEnabled != after.ChannelMonitorEnabled { - changed = append(changed, "channel_monitor_enabled") - } - if before.ChannelMonitorDefaultIntervalSeconds != after.ChannelMonitorDefaultIntervalSeconds { - changed = append(changed, "channel_monitor_default_interval_seconds") - } - if before.AvailableChannelsEnabled != after.AvailableChannelsEnabled { - changed = append(changed, "available_channels_enabled") - } - if before.AffiliateEnabled != after.AffiliateEnabled { - changed = append(changed, "affiliate_enabled") - } - if before.RiskControlEnabled != after.RiskControlEnabled { - changed = append(changed, "risk_control_enabled") - } - if before.CyberSessionBlockEnabled != after.CyberSessionBlockEnabled { - changed = append(changed, "cyber_session_block_enabled") - } - if before.CyberSessionBlockTTLSeconds != after.CyberSessionBlockTTLSeconds { - changed = append(changed, "cyber_session_block_ttl_seconds") - } - // Default platform quotas(JSON map,整体比较) - if !equalPlatformQuotaSettings(before.DefaultPlatformQuotas, after.DefaultPlatformQuotas) { - changed = append(changed, service.SettingKeyDefaultPlatformQuotas) - } - changed = appendAuthSourceDefaultChanges(changed, beforeAuthSourceDefaults, afterAuthSourceDefaults) - return changed -} - -func appendAuthSourceDefaultChanges(changed []string, before *service.AuthSourceDefaultSettings, after *service.AuthSourceDefaultSettings) []string { - if before == nil { - before = &service.AuthSourceDefaultSettings{} - } - if after == nil { - after = &service.AuthSourceDefaultSettings{} - } - - type providerDefaultGrantField struct { - name string - before service.ProviderDefaultGrantSettings - after service.ProviderDefaultGrantSettings - } - - fields := []providerDefaultGrantField{ - {name: "email", before: before.Email, after: after.Email}, - {name: "linuxdo", before: before.LinuxDo, after: after.LinuxDo}, - {name: "oidc", before: before.OIDC, after: after.OIDC}, - {name: "wechat", before: before.WeChat, after: after.WeChat}, - {name: "github", before: before.GitHub, after: after.GitHub}, - {name: "google", before: before.Google, after: after.Google}, - {name: "dingtalk", before: before.DingTalk, after: after.DingTalk}, - } - for _, field := range fields { - if field.before.Balance != field.after.Balance { - changed = append(changed, "auth_source_default_"+field.name+"_balance") - } - if field.before.Concurrency != field.after.Concurrency { - changed = append(changed, "auth_source_default_"+field.name+"_concurrency") - } - if !equalDefaultSubscriptions(field.before.Subscriptions, field.after.Subscriptions) { - changed = append(changed, "auth_source_default_"+field.name+"_subscriptions") - } - if field.before.GrantOnSignup != field.after.GrantOnSignup { - changed = append(changed, "auth_source_default_"+field.name+"_grant_on_signup") - } - if field.before.GrantOnFirstBind != field.after.GrantOnFirstBind { - changed = append(changed, "auth_source_default_"+field.name+"_grant_on_first_bind") - } - // Platform quotas diff:整体替换语义,发单个 JSON key。 - if !equalPlatformQuotaSettings(field.before.PlatformQuotas, field.after.PlatformQuotas) { - changed = append(changed, service.SettingKeyAuthSourcePlatformQuotas(field.name)) - } - } - if before.ForceEmailOnThirdPartySignup != after.ForceEmailOnThirdPartySignup { - changed = append(changed, "force_email_on_third_party_signup") - } - return changed -} - -func normalizeDefaultSubscriptions(input []dto.DefaultSubscriptionSetting) []dto.DefaultSubscriptionSetting { - if len(input) == 0 { - return nil - } - normalized := make([]dto.DefaultSubscriptionSetting, 0, len(input)) - for _, item := range input { - if item.GroupID <= 0 || item.ValidityDays <= 0 { - continue - } - if item.ValidityDays > service.MaxValidityDays { - item.ValidityDays = service.MaxValidityDays - } - normalized = append(normalized, item) - } - return normalized -} - -func normalizeOptionalDefaultSubscriptions(input *[]dto.DefaultSubscriptionSetting) *[]dto.DefaultSubscriptionSetting { - if input == nil { - return nil - } - normalized := normalizeDefaultSubscriptions(*input) - return &normalized -} - -func float64ValueOrDefault(value *float64, fallback float64) float64 { - if value == nil { - return fallback - } - return *value -} - -func intValueOrDefault(value *int, fallback int) int { - if value == nil { - return fallback - } - return *value -} - -func boolValueOrDefault(value *bool, fallback bool) bool { - if value == nil { - return fallback - } - return *value -} - -func defaultSubscriptionsValueOrDefault(input *[]dto.DefaultSubscriptionSetting, fallback []service.DefaultSubscriptionSetting) []service.DefaultSubscriptionSetting { - if input == nil { - return fallback - } - result := make([]service.DefaultSubscriptionSetting, 0, len(*input)) - for _, item := range *input { - result = append(result, service.DefaultSubscriptionSetting{ - GroupID: item.GroupID, - ValidityDays: item.ValidityDays, - }) - } - return result -} - -// platformQuotasValueOrDefault 处理 auth-source platform quota 的 nil 语义: -// nil = 请求未包含该字段(保留 fallback),non-nil(含 empty map)= 整体覆盖。 -// 注意:JSON null 与字段省略等价——两者均反序列化为 nil map,因此都保留旧值; -// 若要清空某 source 的所有 quota 配置,须显式发空对象 {}。 -func platformQuotasValueOrDefault(value, fallback map[string]*service.DefaultPlatformQuotaSetting) map[string]*service.DefaultPlatformQuotaSetting { - if value == nil { - return fallback - } - return value -} - func systemSettingsResponseData(settings dto.SystemSettings, authSourceDefaults *service.AuthSourceDefaultSettings) map[string]any { data := make(map[string]any) raw, err := json.Marshal(settings) @@ -3021,937 +466,3 @@ func systemSettingsResponseData(settings dto.SystemSettings, authSourceDefaults return data } - -func equalStringSlice(a, b []string) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i] != b[i] { - return false - } - } - return true -} - -func equalDefaultSubscriptions(a, b []service.DefaultSubscriptionSetting) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i].GroupID != b[i].GroupID || a[i].ValidityDays != b[i].ValidityDays { - return false - } - } - return true -} - -func equalLoginAgreementDocuments(a, b []service.LoginAgreementDocument) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i].ID != b[i].ID || a[i].Title != b[i].Title || a[i].ContentMD != b[i].ContentMD { - return false - } - } - return true -} - -func equalIntSlice(a, b []int) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i] != b[i] { - return false - } - } - return true -} - -func equalNotifyEmailEntries(a, b []service.NotifyEmailEntry) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i].Email != b[i].Email || a[i].Verified != b[i].Verified || a[i].Disabled != b[i].Disabled { - return false - } - } - return true -} - -// TestSMTPRequest 测试SMTP连接请求 -type TestSMTPRequest struct { - SMTPHost string `json:"smtp_host"` - SMTPPort int `json:"smtp_port"` - SMTPUsername string `json:"smtp_username"` - SMTPPassword string `json:"smtp_password"` - SMTPUseTLS bool `json:"smtp_use_tls"` -} - -// TestSMTPConnection 测试SMTP连接 -// POST /api/v1/admin/settings/test-smtp -func (h *SettingHandler) TestSMTPConnection(c *gin.Context) { - var req TestSMTPRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - req.SMTPHost = strings.TrimSpace(req.SMTPHost) - req.SMTPUsername = strings.TrimSpace(req.SMTPUsername) - - var savedConfig *service.SMTPConfig - if cfg, err := h.emailService.GetSMTPConfig(c.Request.Context()); err == nil && cfg != nil { - savedConfig = cfg - } - - if req.SMTPHost == "" && savedConfig != nil { - req.SMTPHost = savedConfig.Host - } - if req.SMTPPort <= 0 { - if savedConfig != nil && savedConfig.Port > 0 { - req.SMTPPort = savedConfig.Port - } else { - req.SMTPPort = 587 - } - } - if req.SMTPUsername == "" && savedConfig != nil { - req.SMTPUsername = savedConfig.Username - } - password := strings.TrimSpace(req.SMTPPassword) - if password == "" && savedConfig != nil { - password = savedConfig.Password - } - if req.SMTPHost == "" { - response.BadRequest(c, "SMTP host is required") - return - } - - config := &service.SMTPConfig{ - Host: req.SMTPHost, - Port: req.SMTPPort, - Username: req.SMTPUsername, - Password: password, - UseTLS: req.SMTPUseTLS, - } - - err := h.emailService.TestSMTPConnectionWithConfig(config) - if err != nil { - response.BadRequest(c, "SMTP connection test failed: "+err.Error()) - return - } - - response.Success(c, gin.H{"message": "SMTP connection successful"}) -} - -// SendTestEmailRequest 发送测试邮件请求 -type SendTestEmailRequest struct { - Email string `json:"email" binding:"required,email"` - SMTPHost string `json:"smtp_host"` - SMTPPort int `json:"smtp_port"` - SMTPUsername string `json:"smtp_username"` - SMTPPassword string `json:"smtp_password"` - SMTPFrom string `json:"smtp_from_email"` - SMTPFromName string `json:"smtp_from_name"` - SMTPUseTLS bool `json:"smtp_use_tls"` -} - -// SendTestEmail 发送测试邮件 -// POST /api/v1/admin/settings/send-test-email -func (h *SettingHandler) SendTestEmail(c *gin.Context) { - var req SendTestEmailRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - req.SMTPHost = strings.TrimSpace(req.SMTPHost) - req.SMTPUsername = strings.TrimSpace(req.SMTPUsername) - req.SMTPFrom = strings.TrimSpace(req.SMTPFrom) - req.SMTPFromName = strings.TrimSpace(req.SMTPFromName) - - var savedConfig *service.SMTPConfig - if cfg, err := h.emailService.GetSMTPConfig(c.Request.Context()); err == nil && cfg != nil { - savedConfig = cfg - } - - if req.SMTPHost == "" && savedConfig != nil { - req.SMTPHost = savedConfig.Host - } - if req.SMTPPort <= 0 { - if savedConfig != nil && savedConfig.Port > 0 { - req.SMTPPort = savedConfig.Port - } else { - req.SMTPPort = 587 - } - } - if req.SMTPUsername == "" && savedConfig != nil { - req.SMTPUsername = savedConfig.Username - } - password := strings.TrimSpace(req.SMTPPassword) - if password == "" && savedConfig != nil { - password = savedConfig.Password - } - if req.SMTPFrom == "" && savedConfig != nil { - req.SMTPFrom = savedConfig.From - } - if req.SMTPFromName == "" && savedConfig != nil { - req.SMTPFromName = savedConfig.FromName - } - if req.SMTPHost == "" { - response.BadRequest(c, "SMTP host is required") - return - } - - config := &service.SMTPConfig{ - Host: req.SMTPHost, - Port: req.SMTPPort, - Username: req.SMTPUsername, - Password: password, - From: req.SMTPFrom, - FromName: req.SMTPFromName, - UseTLS: req.SMTPUseTLS, - } - - siteName := h.settingService.GetSiteName(c.Request.Context()) - subject := "[" + siteName + "] Test Email" - body := ` - - - - - - - -
-
-

` + siteName + `

-
-
-
✓
-

Email Configuration Successful!

-

This is a test email to verify your SMTP settings are working correctly.

-
- -
- - -` - - if err := h.emailService.SendEmailWithConfig(config, req.Email, subject, body); err != nil { - response.BadRequest(c, "Failed to send test email: "+err.Error()) - return - } - - response.Success(c, gin.H{"message": "Test email sent successfully"}) -} - -// GetAdminAPIKey 获取管理员 API Key 状态 -// GET /api/v1/admin/settings/admin-api-key -func (h *SettingHandler) GetAdminAPIKey(c *gin.Context) { - maskedKey, exists, err := h.settingService.GetAdminAPIKeyStatus(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, gin.H{ - "exists": exists, - "masked_key": maskedKey, - }) -} - -// RegenerateAdminAPIKey 生成/重新生成管理员 API Key -// POST /api/v1/admin/settings/admin-api-key/regenerate -func (h *SettingHandler) RegenerateAdminAPIKey(c *gin.Context) { - key, err := h.settingService.GenerateAdminAPIKey(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, gin.H{ - "key": key, // 完整 key 只在生成时返回一次 - }) -} - -// DeleteAdminAPIKey 删除管理员 API Key -// DELETE /api/v1/admin/settings/admin-api-key -func (h *SettingHandler) DeleteAdminAPIKey(c *gin.Context) { - if err := h.settingService.DeleteAdminAPIKey(c.Request.Context()); err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, gin.H{"message": "Admin API key deleted"}) -} - -// GetOverloadCooldownSettings 获取529过载冷却配置 -// GET /api/v1/admin/settings/overload-cooldown -func (h *SettingHandler) GetOverloadCooldownSettings(c *gin.Context) { - settings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, dto.OverloadCooldownSettings{ - Enabled: settings.Enabled, - CooldownMinutes: settings.CooldownMinutes, - }) -} - -// UpdateOverloadCooldownSettingsRequest 更新529过载冷却配置请求 -type UpdateOverloadCooldownSettingsRequest struct { - Enabled bool `json:"enabled"` - CooldownMinutes int `json:"cooldown_minutes"` -} - -// UpdateOverloadCooldownSettings 更新529过载冷却配置 -// PUT /api/v1/admin/settings/overload-cooldown -func (h *SettingHandler) UpdateOverloadCooldownSettings(c *gin.Context) { - var req UpdateOverloadCooldownSettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - settings := &service.OverloadCooldownSettings{ - Enabled: req.Enabled, - CooldownMinutes: req.CooldownMinutes, - } - - if err := h.settingService.SetOverloadCooldownSettings(c.Request.Context(), settings); err != nil { - response.BadRequest(c, err.Error()) - return - } - - updatedSettings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, dto.OverloadCooldownSettings{ - Enabled: updatedSettings.Enabled, - CooldownMinutes: updatedSettings.CooldownMinutes, - }) -} - -// GetRateLimit429CooldownSettings 获取429默认回避配置 -// GET /api/v1/admin/settings/rate-limit-429-cooldown -func (h *SettingHandler) GetRateLimit429CooldownSettings(c *gin.Context) { - settings, err := h.settingService.GetRateLimit429CooldownSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, dto.RateLimit429CooldownSettings{ - Enabled: settings.Enabled, - CooldownSeconds: settings.CooldownSeconds, - }) -} - -// UpdateRateLimit429CooldownSettingsRequest 更新429默认回避配置请求 -type UpdateRateLimit429CooldownSettingsRequest struct { - Enabled bool `json:"enabled"` - CooldownSeconds int `json:"cooldown_seconds"` -} - -// UpdateRateLimit429CooldownSettings 更新429默认回避配置 -// PUT /api/v1/admin/settings/rate-limit-429-cooldown -func (h *SettingHandler) UpdateRateLimit429CooldownSettings(c *gin.Context) { - var req UpdateRateLimit429CooldownSettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - settings := &service.RateLimit429CooldownSettings{ - Enabled: req.Enabled, - CooldownSeconds: req.CooldownSeconds, - } - - if err := h.settingService.SetRateLimit429CooldownSettings(c.Request.Context(), settings); err != nil { - response.BadRequest(c, err.Error()) - return - } - - updatedSettings, err := h.settingService.GetRateLimit429CooldownSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, dto.RateLimit429CooldownSettings{ - Enabled: updatedSettings.Enabled, - CooldownSeconds: updatedSettings.CooldownSeconds, - }) -} - -// GetStreamTimeoutSettings 获取流超时处理配置 -// GET /api/v1/admin/settings/stream-timeout -func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) { - settings, err := h.settingService.GetStreamTimeoutSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, dto.StreamTimeoutSettings{ - Enabled: settings.Enabled, - Action: settings.Action, - TempUnschedMinutes: settings.TempUnschedMinutes, - ThresholdCount: settings.ThresholdCount, - ThresholdWindowMinutes: settings.ThresholdWindowMinutes, - }) -} - -// GetRectifierSettings 获取请求整流器配置 -// GET /api/v1/admin/settings/rectifier -func (h *SettingHandler) GetRectifierSettings(c *gin.Context) { - settings, err := h.settingService.GetRectifierSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - patterns := settings.APIKeySignaturePatterns - if patterns == nil { - patterns = []string{} - } - response.Success(c, dto.RectifierSettings{ - Enabled: settings.Enabled, - ThinkingSignatureEnabled: settings.ThinkingSignatureEnabled, - ThinkingBudgetEnabled: settings.ThinkingBudgetEnabled, - APIKeySignatureEnabled: settings.APIKeySignatureEnabled, - APIKeySignaturePatterns: patterns, - }) -} - -// UpdateRectifierSettingsRequest 更新整流器配置请求 -type UpdateRectifierSettingsRequest struct { - Enabled bool `json:"enabled"` - ThinkingSignatureEnabled bool `json:"thinking_signature_enabled"` - ThinkingBudgetEnabled bool `json:"thinking_budget_enabled"` - APIKeySignatureEnabled bool `json:"apikey_signature_enabled"` - APIKeySignaturePatterns []string `json:"apikey_signature_patterns"` -} - -// UpdateRectifierSettings 更新请求整流器配置 -// PUT /api/v1/admin/settings/rectifier -func (h *SettingHandler) UpdateRectifierSettings(c *gin.Context) { - var req UpdateRectifierSettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - // 校验并清理自定义匹配关键词 - const maxPatterns = 50 - const maxPatternLen = 500 - if len(req.APIKeySignaturePatterns) > maxPatterns { - response.BadRequest(c, "Too many signature patterns (max 50)") - return - } - var cleanedPatterns []string - for _, p := range req.APIKeySignaturePatterns { - p = strings.TrimSpace(p) - if p == "" { - continue - } - if len(p) > maxPatternLen { - response.BadRequest(c, "Signature pattern too long (max 500 characters)") - return - } - cleanedPatterns = append(cleanedPatterns, p) - } - - settings := &service.RectifierSettings{ - Enabled: req.Enabled, - ThinkingSignatureEnabled: req.ThinkingSignatureEnabled, - ThinkingBudgetEnabled: req.ThinkingBudgetEnabled, - APIKeySignatureEnabled: req.APIKeySignatureEnabled, - APIKeySignaturePatterns: cleanedPatterns, - } - - if err := h.settingService.SetRectifierSettings(c.Request.Context(), settings); err != nil { - response.BadRequest(c, err.Error()) - return - } - - // 重新获取设置返回 - updatedSettings, err := h.settingService.GetRectifierSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - updatedPatterns := updatedSettings.APIKeySignaturePatterns - if updatedPatterns == nil { - updatedPatterns = []string{} - } - response.Success(c, dto.RectifierSettings{ - Enabled: updatedSettings.Enabled, - ThinkingSignatureEnabled: updatedSettings.ThinkingSignatureEnabled, - ThinkingBudgetEnabled: updatedSettings.ThinkingBudgetEnabled, - APIKeySignatureEnabled: updatedSettings.APIKeySignatureEnabled, - APIKeySignaturePatterns: updatedPatterns, - }) -} - -// GetBetaPolicySettings 获取 Beta 策略配置 -// GET /api/v1/admin/settings/beta-policy -func (h *SettingHandler) GetBetaPolicySettings(c *gin.Context) { - settings, err := h.settingService.GetBetaPolicySettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - rules := make([]dto.BetaPolicyRule, len(settings.Rules)) - for i, r := range settings.Rules { - rules[i] = dto.BetaPolicyRule(r) - } - response.Success(c, dto.BetaPolicySettings{Rules: rules}) -} - -// UpdateBetaPolicySettingsRequest 更新 Beta 策略配置请求 -type UpdateBetaPolicySettingsRequest struct { - Rules []dto.BetaPolicyRule `json:"rules"` -} - -// UpdateBetaPolicySettings 更新 Beta 策略配置 -// PUT /api/v1/admin/settings/beta-policy -func (h *SettingHandler) UpdateBetaPolicySettings(c *gin.Context) { - var req UpdateBetaPolicySettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - rules := make([]service.BetaPolicyRule, len(req.Rules)) - for i, r := range req.Rules { - rules[i] = service.BetaPolicyRule(r) - } - - settings := &service.BetaPolicySettings{Rules: rules} - if err := h.settingService.SetBetaPolicySettings(c.Request.Context(), settings); err != nil { - response.BadRequest(c, err.Error()) - return - } - - // Re-fetch to return updated settings - updated, err := h.settingService.GetBetaPolicySettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - outRules := make([]dto.BetaPolicyRule, len(updated.Rules)) - for i, r := range updated.Rules { - outRules[i] = dto.BetaPolicyRule(r) - } - response.Success(c, dto.BetaPolicySettings{Rules: outRules}) -} - -// UpdateStreamTimeoutSettingsRequest 更新流超时配置请求 -type UpdateStreamTimeoutSettingsRequest struct { - Enabled bool `json:"enabled"` - Action string `json:"action"` - TempUnschedMinutes int `json:"temp_unsched_minutes"` - ThresholdCount int `json:"threshold_count"` - ThresholdWindowMinutes int `json:"threshold_window_minutes"` -} - -// UpdateStreamTimeoutSettings 更新流超时处理配置 -// PUT /api/v1/admin/settings/stream-timeout -func (h *SettingHandler) UpdateStreamTimeoutSettings(c *gin.Context) { - var req UpdateStreamTimeoutSettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - settings := &service.StreamTimeoutSettings{ - Enabled: req.Enabled, - Action: req.Action, - TempUnschedMinutes: req.TempUnschedMinutes, - ThresholdCount: req.ThresholdCount, - ThresholdWindowMinutes: req.ThresholdWindowMinutes, - } - - if err := h.settingService.SetStreamTimeoutSettings(c.Request.Context(), settings); err != nil { - response.BadRequest(c, err.Error()) - return - } - - // 重新获取设置返回 - updatedSettings, err := h.settingService.GetStreamTimeoutSettings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, dto.StreamTimeoutSettings{ - Enabled: updatedSettings.Enabled, - Action: updatedSettings.Action, - TempUnschedMinutes: updatedSettings.TempUnschedMinutes, - ThresholdCount: updatedSettings.ThresholdCount, - ThresholdWindowMinutes: updatedSettings.ThresholdWindowMinutes, - }) -} - -// GetWebSearchEmulationConfig 获取 Web Search 模拟配置 -// GET /api/v1/admin/settings/web-search-emulation -func (h *SettingHandler) GetWebSearchEmulationConfig(c *gin.Context) { - cfg, err := h.settingService.GetWebSearchEmulationConfig(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, service.PopulateWebSearchUsage(c.Request.Context(), cfg)) -} - -// UpdateWebSearchEmulationConfig 更新 Web Search 模拟配置 -// PUT /api/v1/admin/settings/web-search-emulation -func (h *SettingHandler) UpdateWebSearchEmulationConfig(c *gin.Context) { - var cfg service.WebSearchEmulationConfig - if err := c.ShouldBindJSON(&cfg); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - if err := h.settingService.SaveWebSearchEmulationConfig(c.Request.Context(), &cfg); err != nil { - response.ErrorFrom(c, err) - return - } - - // Re-read (with sanitized api keys) to return current state - updated, err := h.settingService.GetWebSearchEmulationConfig(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, service.PopulateWebSearchUsage(c.Request.Context(), updated)) -} - -// ResetWebSearchUsage 重置指定 provider 的配额用量 -// POST /api/v1/admin/settings/web-search-emulation/reset-usage -func (h *SettingHandler) ResetWebSearchUsage(c *gin.Context) { - var req struct { - ProviderType string `json:"provider_type"` - } - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - if req.ProviderType == "" { - response.BadRequest(c, "provider_type is required") - return - } - if err := service.ResetWebSearchUsage(c.Request.Context(), req.ProviderType); err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, nil) -} - -// TestWebSearchEmulation 测试 Web Search 搜索 -// POST /api/v1/admin/settings/web-search-emulation/test -func (h *SettingHandler) TestWebSearchEmulation(c *gin.Context) { - var req struct { - Query string `json:"query"` - } - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - if strings.TrimSpace(req.Query) == "" { - req.Query = "搜索今年世界大事件" - } - - result, err := service.TestWebSearch(c.Request.Context(), req.Query) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, result) -} - -// ensureDingTalkSyncAttributes 在保存 settings 后,按 admin 配置的 (attr key, attr name) -// 兜底 upsert 对应 user attribute definition:不存在则创建;存在但 name 不同则更新 name -// (type/options/required 不变)。仅 internal_only + 对应 sync 开关开启时执行。 -// 失败仅记录日志,不阻塞 settings 保存。 -func (h *SettingHandler) ensureDingTalkSyncAttributes(ctx context.Context, settings *service.SystemSettings) { - if h.userAttributeService == nil || settings == nil { - return - } - if settings.DingTalkConnectCorpRestrictionPolicy != "internal_only" { - return - } - if settings.DingTalkConnectSyncDisplayName { - h.ensureUserAttributeDefinition(ctx, settings.DingTalkConnectSyncDisplayNameAttrKey, settings.DingTalkConnectSyncDisplayNameAttrName, "钉钉 internal_only 登录时同步的钉钉姓名", service.AttributeTypeText) - } - if settings.DingTalkConnectSyncCorpEmail { - h.ensureUserAttributeDefinition(ctx, settings.DingTalkConnectSyncCorpEmailAttrKey, settings.DingTalkConnectSyncCorpEmailAttrName, "钉钉 internal_only 登录时同步的企业邮箱", service.AttributeTypeEmail) - } - if settings.DingTalkConnectSyncDept { - h.ensureUserAttributeDefinition(ctx, settings.DingTalkConnectSyncDeptAttrKey, settings.DingTalkConnectSyncDeptAttrName, "钉钉 internal_only 登录时同步的完整部门路径(如:公司/研发部)", service.AttributeTypeText) - } -} - -func (h *SettingHandler) ensureUserAttributeDefinition(ctx context.Context, key, name, description string, attrType service.UserAttributeType) { - key = strings.TrimSpace(key) - if key == "" { - return - } - existing, err := h.userAttributeService.GetDefinitionByKey(ctx, key) - if err == nil && existing != nil { - if strings.TrimSpace(name) != "" && existing.Name != name { - if _, err := h.userAttributeService.UpdateDefinition(ctx, existing.ID, service.UpdateAttributeDefinitionInput{ - Name: &name, - }); err != nil { - slog.Warn("dingtalk: update user attribute definition name failed", "key", key, "err", err.Error()) - return - } - slog.Info("dingtalk: updated user attribute definition name", "key", key, "name", name) - } - return - } - if _, err := h.userAttributeService.CreateDefinition(ctx, service.CreateAttributeDefinitionInput{ - Key: key, - Name: name, - Description: description, - Type: attrType, - Enabled: true, - }); err != nil { - slog.Warn("dingtalk: ensure user attribute definition failed", "key", key, "err", err.Error()) - return - } - slog.Info("dingtalk: created user attribute definition", "key", key, "name", name, "type", attrType) -} - -// ListEmailTemplates returns all editable notification email templates. -// GET /api/v1/admin/settings/email-templates -func (h *SettingHandler) ListEmailTemplates(c *gin.Context) { - if h.notificationEmailService == nil { - response.InternalError(c, "notification email service is not configured") - return - } - events := h.notificationEmailService.ListEventInfos() - templates, err := h.notificationEmailService.ListTemplates(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, dto.EmailTemplateListResponse{ - Events: emailTemplateEventOptionsToDTO(events), - Locales: h.notificationEmailService.SupportedLocales(), - Templates: emailTemplateSummariesToDTO(templates), - Placeholders: emailTemplatePlaceholderUnion(events), - }) -} - -// GetEmailTemplate returns one editable notification email template. -// GET /api/v1/admin/settings/email-templates/:event/:locale -func (h *SettingHandler) GetEmailTemplate(c *gin.Context) { - if h.notificationEmailService == nil { - response.InternalError(c, "notification email service is not configured") - return - } - tmpl, err := h.notificationEmailService.GetTemplate(c.Request.Context(), c.Param("event"), c.Param("locale")) - if err != nil { - response.BadRequest(c, err.Error()) - return - } - response.Success(c, emailTemplateDetailToDTO(tmpl)) -} - -// UpdateEmailTemplate saves an override for one event/locale template. -// PUT /api/v1/admin/settings/email-templates/:event/:locale -func (h *SettingHandler) UpdateEmailTemplate(c *gin.Context) { - if h.notificationEmailService == nil { - response.InternalError(c, "notification email service is not configured") - return - } - var req dto.UpdateEmailTemplateRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - tmpl, err := h.notificationEmailService.UpdateTemplate(c.Request.Context(), c.Param("event"), c.Param("locale"), req.Subject, req.HTML) - if err != nil { - response.BadRequest(c, err.Error()) - return - } - response.Success(c, emailTemplateDetailToDTO(tmpl)) -} - -// RestoreOfficialEmailTemplate removes an override and returns the built-in template. -// POST /api/v1/admin/settings/email-templates/:event/:locale/restore-official -func (h *SettingHandler) RestoreOfficialEmailTemplate(c *gin.Context) { - if h.notificationEmailService == nil { - response.InternalError(c, "notification email service is not configured") - return - } - tmpl, err := h.notificationEmailService.RestoreOfficialTemplate(c.Request.Context(), c.Param("event"), c.Param("locale")) - if err != nil { - response.BadRequest(c, err.Error()) - return - } - response.Success(c, emailTemplateDetailToDTO(tmpl)) -} - -// PreviewEmailTemplate renders a template with safe sample variables without saving it. -// POST /api/v1/admin/settings/email-templates/preview -func (h *SettingHandler) PreviewEmailTemplate(c *gin.Context) { - if h.notificationEmailService == nil { - response.InternalError(c, "notification email service is not configured") - return - } - var req dto.PreviewEmailTemplateRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - preview, err := h.notificationEmailService.PreviewTemplate(c.Request.Context(), service.NotificationEmailPreviewInput{ - Event: req.Event, - Locale: req.Locale, - Subject: req.Subject, - HTML: req.HTML, - Variables: req.Variables, - }) - if err != nil { - response.BadRequest(c, err.Error()) - return - } - response.Success(c, dto.EmailTemplatePreviewResponse{Subject: preview.Subject, HTML: preview.HTML}) -} - -func emailTemplateEventOptionsToDTO(events []service.NotificationEmailEventInfo) []dto.EmailTemplateEventOption { - items := make([]dto.EmailTemplateEventOption, 0, len(events)) - for _, event := range events { - items = append(items, dto.EmailTemplateEventOption{ - Value: event.Event, - Label: event.Label, - Description: event.Description, - Category: event.Category, - Optional: event.Optional, - }) - } - return items -} - -func emailTemplateSummariesToDTO(templates []service.NotificationEmailTemplate) []dto.EmailTemplateSummary { - items := make([]dto.EmailTemplateSummary, 0, len(templates)) - for _, tmpl := range templates { - items = append(items, dto.EmailTemplateSummary{ - Event: tmpl.Event, - Locale: tmpl.Locale, - Subject: tmpl.Subject, - IsCustom: tmpl.IsCustom, - UpdatedAt: emailTemplateUpdatedAt(tmpl), - }) - } - return items -} - -func emailTemplateDetailToDTO(tmpl service.NotificationEmailTemplate) dto.EmailTemplateDetail { - return dto.EmailTemplateDetail{ - Event: tmpl.Event, - Locale: tmpl.Locale, - Subject: tmpl.Subject, - HTML: tmpl.HTML, - IsCustom: tmpl.IsCustom, - UpdatedAt: emailTemplateUpdatedAt(tmpl), - Placeholders: tmpl.Placeholders, - } -} - -func emailTemplateUpdatedAt(tmpl service.NotificationEmailTemplate) string { - if tmpl.UpdatedAt == nil { - return "" - } - return tmpl.UpdatedAt.Format("2006-01-02T15:04:05Z07:00") -} - -func emailTemplatePlaceholderUnion(events []service.NotificationEmailEventInfo) []string { - seen := make(map[string]struct{}) - placeholders := make([]string, 0) - for _, event := range events { - for _, placeholder := range event.Placeholders { - if _, ok := seen[placeholder]; ok { - continue - } - seen[placeholder] = struct{}{} - placeholders = append(placeholders, placeholder) - } - } - return placeholders -} - -// equalNullableFloat compares two *float64 values treating nil as a distinct case. -func equalNullableFloat(a, b *float64) bool { - if a == nil && b == nil { - return true - } - if a == nil || b == nil { - return false - } - return *a == *b -} - -// slotOf returns the *float64 for the given window from a DefaultPlatformQuotaSetting. -func slotOf(s *service.DefaultPlatformQuotaSetting, win string) *float64 { - if s == nil { - return nil - } - switch win { - case "daily": - return s.DailyLimitUSD - case "weekly": - return s.WeeklyLimitUSD - case "monthly": - return s.MonthlyLimitUSD - } - return nil -} - -// equalPlatformQuotaSettings reports whether two platform-quota maps are identical across all allowed slots. -func equalPlatformQuotaSettings(before, after map[string]*service.DefaultPlatformQuotaSetting) bool { - for _, platform := range service.AllowedQuotaPlatforms { - b := before[platform] - a := after[platform] - if !equalNullableFloat(slotOf(b, "daily"), slotOf(a, "daily")) { - return false - } - if !equalNullableFloat(slotOf(b, "weekly"), slotOf(a, "weekly")) { - return false - } - if !equalNullableFloat(slotOf(b, "monthly"), slotOf(a, "monthly")) { - return false - } - } - return true -} - -func stringSetting(value *string, fallback string) string { - if value == nil { - return fallback - } - return *value -} diff --git a/backend/internal/handler/admin/setting_handler_audit.go b/backend/internal/handler/admin/setting_handler_audit.go new file mode 100644 index 0000000000..899030d1b9 --- /dev/null +++ b/backend/internal/handler/admin/setting_handler_audit.go @@ -0,0 +1,758 @@ +package admin + +import ( + "log/slog" + + "github.com/Wei-Shaw/sub2api/internal/handler/dto" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + + "github.com/gin-gonic/gin" +) + +func (h *SettingHandler) auditSettingsUpdate(c *gin.Context, before *service.SystemSettings, after *service.SystemSettings, beforeAuthSourceDefaults *service.AuthSourceDefaultSettings, afterAuthSourceDefaults *service.AuthSourceDefaultSettings, req UpdateSettingsRequest) { + if before == nil || after == nil { + return + } + + changed := diffSettings(before, after, beforeAuthSourceDefaults, afterAuthSourceDefaults, req) + if len(changed) == 0 { + return + } + + subject, _ := middleware.GetAuthSubjectFromContext(c) + role, _ := middleware.GetUserRoleFromContext(c) + slog.Info("settings updated", + "audit", true, + "user_id", subject.UserID, + "role", role, + "changed", changed, + ) +} + +func diffSettings(before *service.SystemSettings, after *service.SystemSettings, beforeAuthSourceDefaults *service.AuthSourceDefaultSettings, afterAuthSourceDefaults *service.AuthSourceDefaultSettings, req UpdateSettingsRequest) []string { + changed := make([]string, 0, 20) + if before.RegistrationEnabled != after.RegistrationEnabled { + changed = append(changed, "registration_enabled") + } + if before.EmailVerifyEnabled != after.EmailVerifyEnabled { + changed = append(changed, "email_verify_enabled") + } + if !equalStringSlice(before.RegistrationEmailSuffixWhitelist, after.RegistrationEmailSuffixWhitelist) { + changed = append(changed, "registration_email_suffix_whitelist") + } + if before.PromoCodeEnabled != after.PromoCodeEnabled { + changed = append(changed, "promo_code_enabled") + } + if before.InvitationCodeEnabled != after.InvitationCodeEnabled { + changed = append(changed, "invitation_code_enabled") + } + if before.PasswordResetEnabled != after.PasswordResetEnabled { + changed = append(changed, "password_reset_enabled") + } + if before.FrontendURL != after.FrontendURL { + changed = append(changed, "frontend_url") + } + if before.TotpEnabled != after.TotpEnabled { + changed = append(changed, "totp_enabled") + } + if before.LoginAgreementEnabled != after.LoginAgreementEnabled { + changed = append(changed, "login_agreement_enabled") + } + if before.LoginAgreementMode != after.LoginAgreementMode { + changed = append(changed, "login_agreement_mode") + } + if before.LoginAgreementUpdatedAt != after.LoginAgreementUpdatedAt { + changed = append(changed, "login_agreement_updated_at") + } + if !equalLoginAgreementDocuments(before.LoginAgreementDocuments, after.LoginAgreementDocuments) { + changed = append(changed, "login_agreement_documents") + } + if before.SMTPHost != after.SMTPHost { + changed = append(changed, "smtp_host") + } + if before.SMTPPort != after.SMTPPort { + changed = append(changed, "smtp_port") + } + if before.SMTPUsername != after.SMTPUsername { + changed = append(changed, "smtp_username") + } + if req.SMTPPassword != "" { + changed = append(changed, "smtp_password") + } + if before.SMTPFrom != after.SMTPFrom { + changed = append(changed, "smtp_from_email") + } + if before.SMTPFromName != after.SMTPFromName { + changed = append(changed, "smtp_from_name") + } + if before.SMTPUseTLS != after.SMTPUseTLS { + changed = append(changed, "smtp_use_tls") + } + if before.TurnstileEnabled != after.TurnstileEnabled { + changed = append(changed, "turnstile_enabled") + } + if before.TurnstileSiteKey != after.TurnstileSiteKey { + changed = append(changed, "turnstile_site_key") + } + if req.TurnstileSecretKey != "" { + changed = append(changed, "turnstile_secret_key") + } + if before.APIKeyACLTrustForwardedIP != after.APIKeyACLTrustForwardedIP { + changed = append(changed, "api_key_acl_trust_forwarded_ip") + } + if before.LinuxDoConnectEnabled != after.LinuxDoConnectEnabled { + changed = append(changed, "linuxdo_connect_enabled") + } + if before.LinuxDoConnectClientID != after.LinuxDoConnectClientID { + changed = append(changed, "linuxdo_connect_client_id") + } + if req.LinuxDoConnectClientSecret != "" { + changed = append(changed, "linuxdo_connect_client_secret") + } + if before.LinuxDoConnectRedirectURL != after.LinuxDoConnectRedirectURL { + changed = append(changed, "linuxdo_connect_redirect_url") + } + if before.DingTalkConnectEnabled != after.DingTalkConnectEnabled { + changed = append(changed, "dingtalk_connect_enabled") + } + if before.DingTalkConnectClientID != after.DingTalkConnectClientID { + changed = append(changed, "dingtalk_connect_client_id") + } + if req.DingTalkConnectClientSecret != "" { + changed = append(changed, "dingtalk_connect_client_secret") + } + if before.DingTalkConnectRedirectURL != after.DingTalkConnectRedirectURL { + changed = append(changed, "dingtalk_connect_redirect_url") + } + if before.DingTalkConnectCorpRestrictionPolicy != after.DingTalkConnectCorpRestrictionPolicy { + changed = append(changed, "dingtalk_connect_corp_restriction_policy") + } + if before.DingTalkConnectInternalCorpID != after.DingTalkConnectInternalCorpID { + changed = append(changed, "dingtalk_connect_internal_corp_id") + } + if before.DingTalkConnectBypassRegistration != after.DingTalkConnectBypassRegistration { + changed = append(changed, "dingtalk_connect_bypass_registration") + } + if before.DingTalkConnectSyncCorpEmail != after.DingTalkConnectSyncCorpEmail { + changed = append(changed, "dingtalk_connect_sync_corp_email") + } + if before.DingTalkConnectSyncDisplayName != after.DingTalkConnectSyncDisplayName { + changed = append(changed, "dingtalk_connect_sync_display_name") + } + if before.DingTalkConnectSyncDept != after.DingTalkConnectSyncDept { + changed = append(changed, "dingtalk_connect_sync_dept") + } + if before.DingTalkConnectSyncCorpEmailAttrKey != after.DingTalkConnectSyncCorpEmailAttrKey { + changed = append(changed, "dingtalk_connect_sync_corp_email_attr_key") + } + if before.DingTalkConnectSyncDisplayNameAttrKey != after.DingTalkConnectSyncDisplayNameAttrKey { + changed = append(changed, "dingtalk_connect_sync_display_name_attr_key") + } + if before.DingTalkConnectSyncDeptAttrKey != after.DingTalkConnectSyncDeptAttrKey { + changed = append(changed, "dingtalk_connect_sync_dept_attr_key") + } + if before.WeChatConnectEnabled != after.WeChatConnectEnabled { + changed = append(changed, "wechat_connect_enabled") + } + if before.WeChatConnectAppID != after.WeChatConnectAppID { + changed = append(changed, "wechat_connect_app_id") + } + if req.WeChatConnectAppSecret != "" { + changed = append(changed, "wechat_connect_app_secret") + } + if before.WeChatConnectOpenAppID != after.WeChatConnectOpenAppID { + changed = append(changed, "wechat_connect_open_app_id") + } + if req.WeChatConnectOpenAppSecret != "" { + changed = append(changed, "wechat_connect_open_app_secret") + } + if before.WeChatConnectMPAppID != after.WeChatConnectMPAppID { + changed = append(changed, "wechat_connect_mp_app_id") + } + if req.WeChatConnectMPAppSecret != "" { + changed = append(changed, "wechat_connect_mp_app_secret") + } + if before.WeChatConnectMobileAppID != after.WeChatConnectMobileAppID { + changed = append(changed, "wechat_connect_mobile_app_id") + } + if req.WeChatConnectMobileAppSecret != "" { + changed = append(changed, "wechat_connect_mobile_app_secret") + } + if before.WeChatConnectOpenEnabled != after.WeChatConnectOpenEnabled { + changed = append(changed, "wechat_connect_open_enabled") + } + if before.WeChatConnectMPEnabled != after.WeChatConnectMPEnabled { + changed = append(changed, "wechat_connect_mp_enabled") + } + if before.WeChatConnectMobileEnabled != after.WeChatConnectMobileEnabled { + changed = append(changed, "wechat_connect_mobile_enabled") + } + if before.WeChatConnectMode != after.WeChatConnectMode { + changed = append(changed, "wechat_connect_mode") + } + if before.WeChatConnectScopes != after.WeChatConnectScopes { + changed = append(changed, "wechat_connect_scopes") + } + if before.WeChatConnectRedirectURL != after.WeChatConnectRedirectURL { + changed = append(changed, "wechat_connect_redirect_url") + } + if before.WeChatConnectFrontendRedirectURL != after.WeChatConnectFrontendRedirectURL { + changed = append(changed, "wechat_connect_frontend_redirect_url") + } + if before.OIDCConnectEnabled != after.OIDCConnectEnabled { + changed = append(changed, "oidc_connect_enabled") + } + if before.OIDCConnectProviderName != after.OIDCConnectProviderName { + changed = append(changed, "oidc_connect_provider_name") + } + if before.OIDCConnectClientID != after.OIDCConnectClientID { + changed = append(changed, "oidc_connect_client_id") + } + if req.OIDCConnectClientSecret != "" { + changed = append(changed, "oidc_connect_client_secret") + } + if before.OIDCConnectIssuerURL != after.OIDCConnectIssuerURL { + changed = append(changed, "oidc_connect_issuer_url") + } + if before.OIDCConnectDiscoveryURL != after.OIDCConnectDiscoveryURL { + changed = append(changed, "oidc_connect_discovery_url") + } + if before.OIDCConnectAuthorizeURL != after.OIDCConnectAuthorizeURL { + changed = append(changed, "oidc_connect_authorize_url") + } + if before.OIDCConnectTokenURL != after.OIDCConnectTokenURL { + changed = append(changed, "oidc_connect_token_url") + } + if before.OIDCConnectUserInfoURL != after.OIDCConnectUserInfoURL { + changed = append(changed, "oidc_connect_userinfo_url") + } + if before.OIDCConnectJWKSURL != after.OIDCConnectJWKSURL { + changed = append(changed, "oidc_connect_jwks_url") + } + if before.OIDCConnectScopes != after.OIDCConnectScopes { + changed = append(changed, "oidc_connect_scopes") + } + if before.OIDCConnectRedirectURL != after.OIDCConnectRedirectURL { + changed = append(changed, "oidc_connect_redirect_url") + } + if before.OIDCConnectFrontendRedirectURL != after.OIDCConnectFrontendRedirectURL { + changed = append(changed, "oidc_connect_frontend_redirect_url") + } + if before.OIDCConnectTokenAuthMethod != after.OIDCConnectTokenAuthMethod { + changed = append(changed, "oidc_connect_token_auth_method") + } + if before.OIDCConnectUsePKCE != after.OIDCConnectUsePKCE { + changed = append(changed, "oidc_connect_use_pkce") + } + if before.OIDCConnectValidateIDToken != after.OIDCConnectValidateIDToken { + changed = append(changed, "oidc_connect_validate_id_token") + } + if before.OIDCConnectAllowedSigningAlgs != after.OIDCConnectAllowedSigningAlgs { + changed = append(changed, "oidc_connect_allowed_signing_algs") + } + if before.OIDCConnectClockSkewSeconds != after.OIDCConnectClockSkewSeconds { + changed = append(changed, "oidc_connect_clock_skew_seconds") + } + if before.OIDCConnectRequireEmailVerified != after.OIDCConnectRequireEmailVerified { + changed = append(changed, "oidc_connect_require_email_verified") + } + if before.OIDCConnectUserInfoEmailPath != after.OIDCConnectUserInfoEmailPath { + changed = append(changed, "oidc_connect_userinfo_email_path") + } + if before.OIDCConnectUserInfoIDPath != after.OIDCConnectUserInfoIDPath { + changed = append(changed, "oidc_connect_userinfo_id_path") + } + if before.OIDCConnectUserInfoUsernamePath != after.OIDCConnectUserInfoUsernamePath { + changed = append(changed, "oidc_connect_userinfo_username_path") + } + if before.SiteName != after.SiteName { + changed = append(changed, "site_name") + } + if before.SiteLogo != after.SiteLogo { + changed = append(changed, "site_logo") + } + if before.SiteSubtitle != after.SiteSubtitle { + changed = append(changed, "site_subtitle") + } + if before.APIBaseURL != after.APIBaseURL { + changed = append(changed, "api_base_url") + } + if before.ContactInfo != after.ContactInfo { + changed = append(changed, "contact_info") + } + if before.DocURL != after.DocURL { + changed = append(changed, "doc_url") + } + if before.HomeContent != after.HomeContent { + changed = append(changed, "home_content") + } + if before.HideCcsImportButton != after.HideCcsImportButton { + changed = append(changed, "hide_ccs_import_button") + } + if before.DefaultConcurrency != after.DefaultConcurrency { + changed = append(changed, "default_concurrency") + } + if before.DefaultBalance != after.DefaultBalance { + changed = append(changed, "default_balance") + } + if before.AffiliateRebateRate != after.AffiliateRebateRate { + changed = append(changed, "affiliate_rebate_rate") + } + if before.AffiliateRebateFreezeHours != after.AffiliateRebateFreezeHours { + changed = append(changed, "affiliate_rebate_freeze_hours") + } + if before.AffiliateRebateDurationDays != after.AffiliateRebateDurationDays { + changed = append(changed, "affiliate_rebate_duration_days") + } + if before.AffiliateRebatePerInviteeCap != after.AffiliateRebatePerInviteeCap { + changed = append(changed, "affiliate_rebate_per_invitee_cap") + } + if !equalDefaultSubscriptions(before.DefaultSubscriptions, after.DefaultSubscriptions) { + changed = append(changed, "default_subscriptions") + } + if before.EnableModelFallback != after.EnableModelFallback { + changed = append(changed, "enable_model_fallback") + } + if before.FallbackModelAnthropic != after.FallbackModelAnthropic { + changed = append(changed, "fallback_model_anthropic") + } + if before.FallbackModelOpenAI != after.FallbackModelOpenAI { + changed = append(changed, "fallback_model_openai") + } + if before.FallbackModelGemini != after.FallbackModelGemini { + changed = append(changed, "fallback_model_gemini") + } + if before.FallbackModelAntigravity != after.FallbackModelAntigravity { + changed = append(changed, "fallback_model_antigravity") + } + if before.EnableIdentityPatch != after.EnableIdentityPatch { + changed = append(changed, "enable_identity_patch") + } + if before.IdentityPatchPrompt != after.IdentityPatchPrompt { + changed = append(changed, "identity_patch_prompt") + } + if before.OpsMonitoringEnabled != after.OpsMonitoringEnabled { + changed = append(changed, "ops_monitoring_enabled") + } + if before.OpsRealtimeMonitoringEnabled != after.OpsRealtimeMonitoringEnabled { + changed = append(changed, "ops_realtime_monitoring_enabled") + } + if before.OpsQueryModeDefault != after.OpsQueryModeDefault { + changed = append(changed, "ops_query_mode_default") + } + if before.OpsMetricsIntervalSeconds != after.OpsMetricsIntervalSeconds { + changed = append(changed, "ops_metrics_interval_seconds") + } + if before.MinClaudeCodeVersion != after.MinClaudeCodeVersion { + changed = append(changed, "min_claude_code_version") + } + if before.MaxClaudeCodeVersion != after.MaxClaudeCodeVersion { + changed = append(changed, "max_claude_code_version") + } + if before.MinCodexVersion != after.MinCodexVersion { + changed = append(changed, "min_codex_version") + } + if before.MaxCodexVersion != after.MaxCodexVersion { + changed = append(changed, "max_codex_version") + } + if before.CodexCLIOnlyAllowAppServerClients != after.CodexCLIOnlyAllowAppServerClients { + changed = append(changed, "codex_cli_only_allow_app_server_clients") + } + if before.CodexCLIOnlyEngineFingerprintSignals != after.CodexCLIOnlyEngineFingerprintSignals { + changed = append(changed, "codex_cli_only_engine_fingerprint_signals") + } + if before.CodexCLIOnlyBlacklist != after.CodexCLIOnlyBlacklist { + changed = append(changed, "codex_cli_only_blacklist") + } + if before.CodexCLIOnlyWhitelist != after.CodexCLIOnlyWhitelist { + changed = append(changed, "codex_cli_only_whitelist") + } + if before.AllowUngroupedKeyScheduling != after.AllowUngroupedKeyScheduling { + changed = append(changed, "allow_ungrouped_key_scheduling") + } + if before.BackendModeEnabled != after.BackendModeEnabled { + changed = append(changed, "backend_mode_enabled") + } + if before.PurchaseSubscriptionEnabled != after.PurchaseSubscriptionEnabled { + changed = append(changed, "purchase_subscription_enabled") + } + if before.PurchaseSubscriptionURL != after.PurchaseSubscriptionURL { + changed = append(changed, "purchase_subscription_url") + } + if before.TableDefaultPageSize != after.TableDefaultPageSize { + changed = append(changed, "table_default_page_size") + } + if !equalIntSlice(before.TablePageSizeOptions, after.TablePageSizeOptions) { + changed = append(changed, "table_page_size_options") + } + if before.CustomMenuItems != after.CustomMenuItems { + changed = append(changed, "custom_menu_items") + } + if before.CustomEndpoints != after.CustomEndpoints { + changed = append(changed, "custom_endpoints") + } + if before.EnableFingerprintUnification != after.EnableFingerprintUnification { + changed = append(changed, "enable_fingerprint_unification") + } + if before.EnableMetadataPassthrough != after.EnableMetadataPassthrough { + changed = append(changed, "enable_metadata_passthrough") + } + if before.EnableCCHSigning != after.EnableCCHSigning { + changed = append(changed, "enable_cch_signing") + } + if before.EnableClaudeOAuthSystemPromptInjection != after.EnableClaudeOAuthSystemPromptInjection { + changed = append(changed, "enable_claude_oauth_system_prompt_injection") + } + if before.ClaudeOAuthSystemPrompt != after.ClaudeOAuthSystemPrompt { + changed = append(changed, "claude_oauth_system_prompt") + } + if before.ClaudeOAuthSystemPromptBlocks != after.ClaudeOAuthSystemPromptBlocks { + changed = append(changed, "claude_oauth_system_prompt_blocks") + } + if before.EnableAnthropicCacheTTL1hInjection != after.EnableAnthropicCacheTTL1hInjection { + changed = append(changed, "enable_anthropic_cache_ttl_1h_injection") + } + if before.RewriteMessageCacheControl != after.RewriteMessageCacheControl { + changed = append(changed, "rewrite_message_cache_control") + } + if before.EnableClientDatelineNormalization != after.EnableClientDatelineNormalization { + changed = append(changed, "enable_client_dateline_normalization") + } + if before.AntigravityUserAgentVersion != after.AntigravityUserAgentVersion { + changed = append(changed, "antigravity_user_agent_version") + } + if before.OpenAICodexUserAgent != after.OpenAICodexUserAgent { + changed = append(changed, "openai_codex_user_agent") + } + if before.PaymentVisibleMethodAlipaySource != after.PaymentVisibleMethodAlipaySource { + changed = append(changed, "payment_visible_method_alipay_source") + } + if before.PaymentVisibleMethodWxpaySource != after.PaymentVisibleMethodWxpaySource { + changed = append(changed, "payment_visible_method_wxpay_source") + } + if before.PaymentVisibleMethodAlipayEnabled != after.PaymentVisibleMethodAlipayEnabled { + changed = append(changed, "payment_visible_method_alipay_enabled") + } + if before.PaymentVisibleMethodWxpayEnabled != after.PaymentVisibleMethodWxpayEnabled { + changed = append(changed, "payment_visible_method_wxpay_enabled") + } + if before.OpenAIAdvancedSchedulerEnabled != after.OpenAIAdvancedSchedulerEnabled { + changed = append(changed, "openai_advanced_scheduler_enabled") + } + if before.OpenAIAdvancedSchedulerStickyWeightedEnabled != after.OpenAIAdvancedSchedulerStickyWeightedEnabled { + changed = append(changed, "openai_advanced_scheduler_sticky_weighted_enabled") + } + if before.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled != after.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled { + changed = append(changed, "openai_advanced_scheduler_subscription_priority_enabled") + } + if before.OpenAIAdvancedSchedulerLBTopK != after.OpenAIAdvancedSchedulerLBTopK { + changed = append(changed, "openai_advanced_scheduler_lb_top_k") + } + if before.OpenAIAdvancedSchedulerWeightPriority != after.OpenAIAdvancedSchedulerWeightPriority { + changed = append(changed, "openai_advanced_scheduler_weight_priority") + } + if before.OpenAIAdvancedSchedulerWeightLoad != after.OpenAIAdvancedSchedulerWeightLoad { + changed = append(changed, "openai_advanced_scheduler_weight_load") + } + if before.OpenAIAdvancedSchedulerWeightQueue != after.OpenAIAdvancedSchedulerWeightQueue { + changed = append(changed, "openai_advanced_scheduler_weight_queue") + } + if before.OpenAIAdvancedSchedulerWeightErrorRate != after.OpenAIAdvancedSchedulerWeightErrorRate { + changed = append(changed, "openai_advanced_scheduler_weight_error_rate") + } + if before.OpenAIAdvancedSchedulerWeightTTFT != after.OpenAIAdvancedSchedulerWeightTTFT { + changed = append(changed, "openai_advanced_scheduler_weight_ttft") + } + if before.OpenAIAdvancedSchedulerWeightReset != after.OpenAIAdvancedSchedulerWeightReset { + changed = append(changed, "openai_advanced_scheduler_weight_reset") + } + if before.OpenAIAdvancedSchedulerWeightQuotaHeadroom != after.OpenAIAdvancedSchedulerWeightQuotaHeadroom { + changed = append(changed, "openai_advanced_scheduler_weight_quota_headroom") + } + if before.OpenAIAdvancedSchedulerWeightPreviousResponse != after.OpenAIAdvancedSchedulerWeightPreviousResponse { + changed = append(changed, "openai_advanced_scheduler_weight_previous_response") + } + if before.OpenAIAdvancedSchedulerWeightSessionSticky != after.OpenAIAdvancedSchedulerWeightSessionSticky { + changed = append(changed, "openai_advanced_scheduler_weight_session_sticky") + } + // 余额、订阅到期与账号限额通知 + if before.BalanceLowNotifyEnabled != after.BalanceLowNotifyEnabled { + changed = append(changed, "balance_low_notify_enabled") + } + if before.BalanceLowNotifyThreshold != after.BalanceLowNotifyThreshold { + changed = append(changed, "balance_low_notify_threshold") + } + if before.BalanceLowNotifyRechargeURL != after.BalanceLowNotifyRechargeURL { + changed = append(changed, "balance_low_notify_recharge_url") + } + if before.SubscriptionExpiryNotifyEnabled != after.SubscriptionExpiryNotifyEnabled { + changed = append(changed, "subscription_expiry_notify_enabled") + } + if before.AccountQuotaNotifyEnabled != after.AccountQuotaNotifyEnabled { + changed = append(changed, "account_quota_notify_enabled") + } + if !equalNotifyEmailEntries(before.AccountQuotaNotifyEmails, after.AccountQuotaNotifyEmails) { + changed = append(changed, "account_quota_notify_emails") + } + if before.ChannelMonitorEnabled != after.ChannelMonitorEnabled { + changed = append(changed, "channel_monitor_enabled") + } + if before.ChannelMonitorDefaultIntervalSeconds != after.ChannelMonitorDefaultIntervalSeconds { + changed = append(changed, "channel_monitor_default_interval_seconds") + } + if before.AvailableChannelsEnabled != after.AvailableChannelsEnabled { + changed = append(changed, "available_channels_enabled") + } + if before.AffiliateEnabled != after.AffiliateEnabled { + changed = append(changed, "affiliate_enabled") + } + if before.RiskControlEnabled != after.RiskControlEnabled { + changed = append(changed, "risk_control_enabled") + } + if before.CyberSessionBlockEnabled != after.CyberSessionBlockEnabled { + changed = append(changed, "cyber_session_block_enabled") + } + if before.CyberSessionBlockTTLSeconds != after.CyberSessionBlockTTLSeconds { + changed = append(changed, "cyber_session_block_ttl_seconds") + } + // Default platform quotas(JSON map,整体比较) + if !equalPlatformQuotaSettings(before.DefaultPlatformQuotas, after.DefaultPlatformQuotas) { + changed = append(changed, service.SettingKeyDefaultPlatformQuotas) + } + changed = appendAuthSourceDefaultChanges(changed, beforeAuthSourceDefaults, afterAuthSourceDefaults) + return changed +} + +func appendAuthSourceDefaultChanges(changed []string, before *service.AuthSourceDefaultSettings, after *service.AuthSourceDefaultSettings) []string { + if before == nil { + before = &service.AuthSourceDefaultSettings{} + } + if after == nil { + after = &service.AuthSourceDefaultSettings{} + } + + type providerDefaultGrantField struct { + name string + before service.ProviderDefaultGrantSettings + after service.ProviderDefaultGrantSettings + } + + fields := []providerDefaultGrantField{ + {name: "email", before: before.Email, after: after.Email}, + {name: "linuxdo", before: before.LinuxDo, after: after.LinuxDo}, + {name: "oidc", before: before.OIDC, after: after.OIDC}, + {name: "wechat", before: before.WeChat, after: after.WeChat}, + {name: "github", before: before.GitHub, after: after.GitHub}, + {name: "google", before: before.Google, after: after.Google}, + {name: "dingtalk", before: before.DingTalk, after: after.DingTalk}, + } + for _, field := range fields { + if field.before.Balance != field.after.Balance { + changed = append(changed, "auth_source_default_"+field.name+"_balance") + } + if field.before.Concurrency != field.after.Concurrency { + changed = append(changed, "auth_source_default_"+field.name+"_concurrency") + } + if !equalDefaultSubscriptions(field.before.Subscriptions, field.after.Subscriptions) { + changed = append(changed, "auth_source_default_"+field.name+"_subscriptions") + } + if field.before.GrantOnSignup != field.after.GrantOnSignup { + changed = append(changed, "auth_source_default_"+field.name+"_grant_on_signup") + } + if field.before.GrantOnFirstBind != field.after.GrantOnFirstBind { + changed = append(changed, "auth_source_default_"+field.name+"_grant_on_first_bind") + } + // Platform quotas diff:整体替换语义,发单个 JSON key。 + if !equalPlatformQuotaSettings(field.before.PlatformQuotas, field.after.PlatformQuotas) { + changed = append(changed, service.SettingKeyAuthSourcePlatformQuotas(field.name)) + } + } + if before.ForceEmailOnThirdPartySignup != after.ForceEmailOnThirdPartySignup { + changed = append(changed, "force_email_on_third_party_signup") + } + return changed +} + +func normalizeDefaultSubscriptions(input []dto.DefaultSubscriptionSetting) []dto.DefaultSubscriptionSetting { + if len(input) == 0 { + return nil + } + normalized := make([]dto.DefaultSubscriptionSetting, 0, len(input)) + for _, item := range input { + if item.GroupID <= 0 || item.ValidityDays <= 0 { + continue + } + if item.ValidityDays > service.MaxValidityDays { + item.ValidityDays = service.MaxValidityDays + } + normalized = append(normalized, item) + } + return normalized +} + +func normalizeOptionalDefaultSubscriptions(input *[]dto.DefaultSubscriptionSetting) *[]dto.DefaultSubscriptionSetting { + if input == nil { + return nil + } + normalized := normalizeDefaultSubscriptions(*input) + return &normalized +} + +func float64ValueOrDefault(value *float64, fallback float64) float64 { + if value == nil { + return fallback + } + return *value +} + +func intValueOrDefault(value *int, fallback int) int { + if value == nil { + return fallback + } + return *value +} + +func boolValueOrDefault(value *bool, fallback bool) bool { + if value == nil { + return fallback + } + return *value +} + +func defaultSubscriptionsValueOrDefault(input *[]dto.DefaultSubscriptionSetting, fallback []service.DefaultSubscriptionSetting) []service.DefaultSubscriptionSetting { + if input == nil { + return fallback + } + result := make([]service.DefaultSubscriptionSetting, 0, len(*input)) + for _, item := range *input { + result = append(result, service.DefaultSubscriptionSetting{ + GroupID: item.GroupID, + ValidityDays: item.ValidityDays, + }) + } + return result +} + +// platformQuotasValueOrDefault 处理 auth-source platform quota 的 nil 语义: +// nil = 请求未包含该字段(保留 fallback),non-nil(含 empty map)= 整体覆盖。 +// 注意:JSON null 与字段省略等价——两者均反序列化为 nil map,因此都保留旧值; +// 若要清空某 source 的所有 quota 配置,须显式发空对象 {}。 +func platformQuotasValueOrDefault(value, fallback map[string]*service.DefaultPlatformQuotaSetting) map[string]*service.DefaultPlatformQuotaSetting { + if value == nil { + return fallback + } + return value +} + +func equalStringSlice(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func equalDefaultSubscriptions(a, b []service.DefaultSubscriptionSetting) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].GroupID != b[i].GroupID || a[i].ValidityDays != b[i].ValidityDays { + return false + } + } + return true +} + +func equalLoginAgreementDocuments(a, b []service.LoginAgreementDocument) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].ID != b[i].ID || a[i].Title != b[i].Title || a[i].ContentMD != b[i].ContentMD { + return false + } + } + return true +} + +func equalIntSlice(a, b []int) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func equalNotifyEmailEntries(a, b []service.NotifyEmailEntry) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].Email != b[i].Email || a[i].Verified != b[i].Verified || a[i].Disabled != b[i].Disabled { + return false + } + } + return true +} + +// equalNullableFloat compares two *float64 values treating nil as a distinct case. +func equalNullableFloat(a, b *float64) bool { + if a == nil && b == nil { + return true + } + if a == nil || b == nil { + return false + } + return *a == *b +} + +// slotOf returns the *float64 for the given window from a DefaultPlatformQuotaSetting. +func slotOf(s *service.DefaultPlatformQuotaSetting, win string) *float64 { + if s == nil { + return nil + } + switch win { + case "daily": + return s.DailyLimitUSD + case "weekly": + return s.WeeklyLimitUSD + case "monthly": + return s.MonthlyLimitUSD + } + return nil +} + +// equalPlatformQuotaSettings reports whether two platform-quota maps are identical across all allowed slots. +func equalPlatformQuotaSettings(before, after map[string]*service.DefaultPlatformQuotaSetting) bool { + for _, platform := range service.AllowedQuotaPlatforms { + b := before[platform] + a := after[platform] + if !equalNullableFloat(slotOf(b, "daily"), slotOf(a, "daily")) { + return false + } + if !equalNullableFloat(slotOf(b, "weekly"), slotOf(a, "weekly")) { + return false + } + if !equalNullableFloat(slotOf(b, "monthly"), slotOf(a, "monthly")) { + return false + } + } + return true +} + +func stringSetting(value *string, fallback string) string { + if value == nil { + return fallback + } + return *value +} diff --git a/backend/internal/handler/admin/setting_handler_email.go b/backend/internal/handler/admin/setting_handler_email.go new file mode 100644 index 0000000000..68d76de6cc --- /dev/null +++ b/backend/internal/handler/admin/setting_handler_email.go @@ -0,0 +1,347 @@ +package admin + +import ( + "html" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/handler/dto" + "github.com/Wei-Shaw/sub2api/internal/pkg/response" + "github.com/Wei-Shaw/sub2api/internal/service" + + "github.com/gin-gonic/gin" +) + +// TestSMTPRequest 测试SMTP连接请求 +type TestSMTPRequest struct { + SMTPHost string `json:"smtp_host"` + SMTPPort int `json:"smtp_port"` + SMTPUsername string `json:"smtp_username"` + SMTPPassword string `json:"smtp_password"` + SMTPUseTLS bool `json:"smtp_use_tls"` +} + +// TestSMTPConnection 测试SMTP连接 +// POST /api/v1/admin/settings/test-smtp +func (h *SettingHandler) TestSMTPConnection(c *gin.Context) { + var req TestSMTPRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + req.SMTPHost = strings.TrimSpace(req.SMTPHost) + req.SMTPUsername = strings.TrimSpace(req.SMTPUsername) + + var savedConfig *service.SMTPConfig + if cfg, err := h.emailService.GetSMTPConfig(c.Request.Context()); err == nil && cfg != nil { + savedConfig = cfg + } + + if req.SMTPHost == "" && savedConfig != nil { + req.SMTPHost = savedConfig.Host + } + if req.SMTPPort <= 0 { + if savedConfig != nil && savedConfig.Port > 0 { + req.SMTPPort = savedConfig.Port + } else { + req.SMTPPort = 587 + } + } + if req.SMTPUsername == "" && savedConfig != nil { + req.SMTPUsername = savedConfig.Username + } + password := strings.TrimSpace(req.SMTPPassword) + if password == "" && savedConfig != nil { + password = savedConfig.Password + } + if req.SMTPHost == "" { + response.BadRequest(c, "SMTP host is required") + return + } + + config := &service.SMTPConfig{ + Host: req.SMTPHost, + Port: req.SMTPPort, + Username: req.SMTPUsername, + Password: password, + UseTLS: req.SMTPUseTLS, + } + + err := h.emailService.TestSMTPConnectionWithConfig(config) + if err != nil { + response.BadRequest(c, "SMTP connection test failed: "+err.Error()) + return + } + + response.Success(c, gin.H{"message": "SMTP connection successful"}) +} + +// SendTestEmailRequest 发送测试邮件请求 +type SendTestEmailRequest struct { + Email string `json:"email" binding:"required,email"` + SMTPHost string `json:"smtp_host"` + SMTPPort int `json:"smtp_port"` + SMTPUsername string `json:"smtp_username"` + SMTPPassword string `json:"smtp_password"` + SMTPFrom string `json:"smtp_from_email"` + SMTPFromName string `json:"smtp_from_name"` + SMTPUseTLS bool `json:"smtp_use_tls"` +} + +// SendTestEmail 发送测试邮件 +// POST /api/v1/admin/settings/send-test-email +func (h *SettingHandler) SendTestEmail(c *gin.Context) { + var req SendTestEmailRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + req.SMTPHost = strings.TrimSpace(req.SMTPHost) + req.SMTPUsername = strings.TrimSpace(req.SMTPUsername) + req.SMTPFrom = strings.TrimSpace(req.SMTPFrom) + req.SMTPFromName = strings.TrimSpace(req.SMTPFromName) + + var savedConfig *service.SMTPConfig + if cfg, err := h.emailService.GetSMTPConfig(c.Request.Context()); err == nil && cfg != nil { + savedConfig = cfg + } + + if req.SMTPHost == "" && savedConfig != nil { + req.SMTPHost = savedConfig.Host + } + if req.SMTPPort <= 0 { + if savedConfig != nil && savedConfig.Port > 0 { + req.SMTPPort = savedConfig.Port + } else { + req.SMTPPort = 587 + } + } + if req.SMTPUsername == "" && savedConfig != nil { + req.SMTPUsername = savedConfig.Username + } + password := strings.TrimSpace(req.SMTPPassword) + if password == "" && savedConfig != nil { + password = savedConfig.Password + } + if req.SMTPFrom == "" && savedConfig != nil { + req.SMTPFrom = savedConfig.From + } + if req.SMTPFromName == "" && savedConfig != nil { + req.SMTPFromName = savedConfig.FromName + } + if req.SMTPHost == "" { + response.BadRequest(c, "SMTP host is required") + return + } + + config := &service.SMTPConfig{ + Host: req.SMTPHost, + Port: req.SMTPPort, + Username: req.SMTPUsername, + Password: password, + From: req.SMTPFrom, + FromName: req.SMTPFromName, + UseTLS: req.SMTPUseTLS, + } + + siteName := h.settingService.GetSiteName(c.Request.Context()) + subject := "[" + siteName + "] Test Email" + body := ` + + + + + + + +
+
+

` + html.EscapeString(siteName) + `

+
+
+
✓
+

Email Configuration Successful!

+

This is a test email to verify your SMTP settings are working correctly.

+
+ +
+ + +` + + if err := h.emailService.SendEmailWithConfig(config, req.Email, subject, body); err != nil { + response.BadRequest(c, "Failed to send test email: "+err.Error()) + return + } + + response.Success(c, gin.H{"message": "Test email sent successfully"}) +} + +// ListEmailTemplates returns all editable notification email templates. +// GET /api/v1/admin/settings/email-templates +func (h *SettingHandler) ListEmailTemplates(c *gin.Context) { + if h.notificationEmailService == nil { + response.InternalError(c, "notification email service is not configured") + return + } + events := h.notificationEmailService.ListEventInfos() + templates, err := h.notificationEmailService.ListTemplates(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, dto.EmailTemplateListResponse{ + Events: emailTemplateEventOptionsToDTO(events), + Locales: h.notificationEmailService.SupportedLocales(), + Templates: emailTemplateSummariesToDTO(templates), + Placeholders: emailTemplatePlaceholderUnion(events), + }) +} + +// GetEmailTemplate returns one editable notification email template. +// GET /api/v1/admin/settings/email-templates/:event/:locale +func (h *SettingHandler) GetEmailTemplate(c *gin.Context) { + if h.notificationEmailService == nil { + response.InternalError(c, "notification email service is not configured") + return + } + tmpl, err := h.notificationEmailService.GetTemplate(c.Request.Context(), c.Param("event"), c.Param("locale")) + if err != nil { + response.BadRequest(c, err.Error()) + return + } + response.Success(c, emailTemplateDetailToDTO(tmpl)) +} + +// UpdateEmailTemplate saves an override for one event/locale template. +// PUT /api/v1/admin/settings/email-templates/:event/:locale +func (h *SettingHandler) UpdateEmailTemplate(c *gin.Context) { + if h.notificationEmailService == nil { + response.InternalError(c, "notification email service is not configured") + return + } + var req dto.UpdateEmailTemplateRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + tmpl, err := h.notificationEmailService.UpdateTemplate(c.Request.Context(), c.Param("event"), c.Param("locale"), req.Subject, req.HTML) + if err != nil { + response.BadRequest(c, err.Error()) + return + } + response.Success(c, emailTemplateDetailToDTO(tmpl)) +} + +// RestoreOfficialEmailTemplate removes an override and returns the built-in template. +// POST /api/v1/admin/settings/email-templates/:event/:locale/restore-official +func (h *SettingHandler) RestoreOfficialEmailTemplate(c *gin.Context) { + if h.notificationEmailService == nil { + response.InternalError(c, "notification email service is not configured") + return + } + tmpl, err := h.notificationEmailService.RestoreOfficialTemplate(c.Request.Context(), c.Param("event"), c.Param("locale")) + if err != nil { + response.BadRequest(c, err.Error()) + return + } + response.Success(c, emailTemplateDetailToDTO(tmpl)) +} + +// PreviewEmailTemplate renders a template with safe sample variables without saving it. +// POST /api/v1/admin/settings/email-templates/preview +func (h *SettingHandler) PreviewEmailTemplate(c *gin.Context) { + if h.notificationEmailService == nil { + response.InternalError(c, "notification email service is not configured") + return + } + var req dto.PreviewEmailTemplateRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + preview, err := h.notificationEmailService.PreviewTemplate(c.Request.Context(), service.NotificationEmailPreviewInput{ + Event: req.Event, + Locale: req.Locale, + Subject: req.Subject, + HTML: req.HTML, + Variables: req.Variables, + }) + if err != nil { + response.BadRequest(c, err.Error()) + return + } + response.Success(c, dto.EmailTemplatePreviewResponse{Subject: preview.Subject, HTML: preview.HTML}) +} + +func emailTemplateEventOptionsToDTO(events []service.NotificationEmailEventInfo) []dto.EmailTemplateEventOption { + items := make([]dto.EmailTemplateEventOption, 0, len(events)) + for _, event := range events { + items = append(items, dto.EmailTemplateEventOption{ + Value: event.Event, + Label: event.Label, + Description: event.Description, + Category: event.Category, + Optional: event.Optional, + }) + } + return items +} + +func emailTemplateSummariesToDTO(templates []service.NotificationEmailTemplate) []dto.EmailTemplateSummary { + items := make([]dto.EmailTemplateSummary, 0, len(templates)) + for _, tmpl := range templates { + items = append(items, dto.EmailTemplateSummary{ + Event: tmpl.Event, + Locale: tmpl.Locale, + Subject: tmpl.Subject, + IsCustom: tmpl.IsCustom, + UpdatedAt: emailTemplateUpdatedAt(tmpl), + }) + } + return items +} + +func emailTemplateDetailToDTO(tmpl service.NotificationEmailTemplate) dto.EmailTemplateDetail { + return dto.EmailTemplateDetail{ + Event: tmpl.Event, + Locale: tmpl.Locale, + Subject: tmpl.Subject, + HTML: tmpl.HTML, + IsCustom: tmpl.IsCustom, + UpdatedAt: emailTemplateUpdatedAt(tmpl), + Placeholders: tmpl.Placeholders, + } +} + +func emailTemplateUpdatedAt(tmpl service.NotificationEmailTemplate) string { + if tmpl.UpdatedAt == nil { + return "" + } + return tmpl.UpdatedAt.Format("2006-01-02T15:04:05Z07:00") +} + +func emailTemplatePlaceholderUnion(events []service.NotificationEmailEventInfo) []string { + seen := make(map[string]struct{}) + placeholders := make([]string, 0) + for _, event := range events { + for _, placeholder := range event.Placeholders { + if _, ok := seen[placeholder]; ok { + continue + } + seen[placeholder] = struct{}{} + placeholders = append(placeholders, placeholder) + } + } + return placeholders +} diff --git a/backend/internal/handler/admin/setting_handler_runtime.go b/backend/internal/handler/admin/setting_handler_runtime.go new file mode 100644 index 0000000000..4891c9b7ad --- /dev/null +++ b/backend/internal/handler/admin/setting_handler_runtime.go @@ -0,0 +1,445 @@ +package admin + +import ( + "strings" + + "github.com/Wei-Shaw/sub2api/internal/handler/dto" + "github.com/Wei-Shaw/sub2api/internal/pkg/response" + "github.com/Wei-Shaw/sub2api/internal/service" + + "github.com/gin-gonic/gin" +) + +// GetAdminAPIKey 获取管理员 API Key 状态 +// GET /api/v1/admin/settings/admin-api-key +func (h *SettingHandler) GetAdminAPIKey(c *gin.Context) { + maskedKey, exists, err := h.settingService.GetAdminAPIKeyStatus(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, gin.H{ + "exists": exists, + "masked_key": maskedKey, + }) +} + +// RegenerateAdminAPIKey 生成/重新生成管理员 API Key +// POST /api/v1/admin/settings/admin-api-key/regenerate +func (h *SettingHandler) RegenerateAdminAPIKey(c *gin.Context) { + key, err := h.settingService.GenerateAdminAPIKey(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, gin.H{ + "key": key, // 完整 key 只在生成时返回一次 + }) +} + +// DeleteAdminAPIKey 删除管理员 API Key +// DELETE /api/v1/admin/settings/admin-api-key +func (h *SettingHandler) DeleteAdminAPIKey(c *gin.Context) { + if err := h.settingService.DeleteAdminAPIKey(c.Request.Context()); err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, gin.H{"message": "Admin API key deleted"}) +} + +// GetOverloadCooldownSettings 获取529过载冷却配置 +// GET /api/v1/admin/settings/overload-cooldown +func (h *SettingHandler) GetOverloadCooldownSettings(c *gin.Context) { + settings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, dto.OverloadCooldownSettings{ + Enabled: settings.Enabled, + CooldownMinutes: settings.CooldownMinutes, + }) +} + +// UpdateOverloadCooldownSettingsRequest 更新529过载冷却配置请求 +type UpdateOverloadCooldownSettingsRequest struct { + Enabled bool `json:"enabled"` + CooldownMinutes int `json:"cooldown_minutes"` +} + +// UpdateOverloadCooldownSettings 更新529过载冷却配置 +// PUT /api/v1/admin/settings/overload-cooldown +func (h *SettingHandler) UpdateOverloadCooldownSettings(c *gin.Context) { + var req UpdateOverloadCooldownSettingsRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + settings := &service.OverloadCooldownSettings{ + Enabled: req.Enabled, + CooldownMinutes: req.CooldownMinutes, + } + + if err := h.settingService.SetOverloadCooldownSettings(c.Request.Context(), settings); err != nil { + response.BadRequest(c, err.Error()) + return + } + + updatedSettings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, dto.OverloadCooldownSettings{ + Enabled: updatedSettings.Enabled, + CooldownMinutes: updatedSettings.CooldownMinutes, + }) +} + +// GetRateLimit429CooldownSettings 获取429默认回避配置 +// GET /api/v1/admin/settings/rate-limit-429-cooldown +func (h *SettingHandler) GetRateLimit429CooldownSettings(c *gin.Context) { + settings, err := h.settingService.GetRateLimit429CooldownSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, dto.RateLimit429CooldownSettings{ + Enabled: settings.Enabled, + CooldownSeconds: settings.CooldownSeconds, + }) +} + +// UpdateRateLimit429CooldownSettingsRequest 更新429默认回避配置请求 +type UpdateRateLimit429CooldownSettingsRequest struct { + Enabled bool `json:"enabled"` + CooldownSeconds int `json:"cooldown_seconds"` +} + +// UpdateRateLimit429CooldownSettings 更新429默认回避配置 +// PUT /api/v1/admin/settings/rate-limit-429-cooldown +func (h *SettingHandler) UpdateRateLimit429CooldownSettings(c *gin.Context) { + var req UpdateRateLimit429CooldownSettingsRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + settings := &service.RateLimit429CooldownSettings{ + Enabled: req.Enabled, + CooldownSeconds: req.CooldownSeconds, + } + + if err := h.settingService.SetRateLimit429CooldownSettings(c.Request.Context(), settings); err != nil { + response.BadRequest(c, err.Error()) + return + } + + updatedSettings, err := h.settingService.GetRateLimit429CooldownSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, dto.RateLimit429CooldownSettings{ + Enabled: updatedSettings.Enabled, + CooldownSeconds: updatedSettings.CooldownSeconds, + }) +} + +// GetStreamTimeoutSettings 获取流超时处理配置 +// GET /api/v1/admin/settings/stream-timeout +func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) { + settings, err := h.settingService.GetStreamTimeoutSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, dto.StreamTimeoutSettings{ + Enabled: settings.Enabled, + Action: settings.Action, + TempUnschedMinutes: settings.TempUnschedMinutes, + ThresholdCount: settings.ThresholdCount, + ThresholdWindowMinutes: settings.ThresholdWindowMinutes, + }) +} + +// GetRectifierSettings 获取请求整流器配置 +// GET /api/v1/admin/settings/rectifier +func (h *SettingHandler) GetRectifierSettings(c *gin.Context) { + settings, err := h.settingService.GetRectifierSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + patterns := settings.APIKeySignaturePatterns + if patterns == nil { + patterns = []string{} + } + response.Success(c, dto.RectifierSettings{ + Enabled: settings.Enabled, + ThinkingSignatureEnabled: settings.ThinkingSignatureEnabled, + ThinkingBudgetEnabled: settings.ThinkingBudgetEnabled, + APIKeySignatureEnabled: settings.APIKeySignatureEnabled, + APIKeySignaturePatterns: patterns, + }) +} + +// UpdateRectifierSettingsRequest 更新整流器配置请求 +type UpdateRectifierSettingsRequest struct { + Enabled bool `json:"enabled"` + ThinkingSignatureEnabled bool `json:"thinking_signature_enabled"` + ThinkingBudgetEnabled bool `json:"thinking_budget_enabled"` + APIKeySignatureEnabled bool `json:"apikey_signature_enabled"` + APIKeySignaturePatterns []string `json:"apikey_signature_patterns"` +} + +// UpdateRectifierSettings 更新请求整流器配置 +// PUT /api/v1/admin/settings/rectifier +func (h *SettingHandler) UpdateRectifierSettings(c *gin.Context) { + var req UpdateRectifierSettingsRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + // 校验并清理自定义匹配关键词 + const maxPatterns = 50 + const maxPatternLen = 500 + if len(req.APIKeySignaturePatterns) > maxPatterns { + response.BadRequest(c, "Too many signature patterns (max 50)") + return + } + var cleanedPatterns []string + for _, p := range req.APIKeySignaturePatterns { + p = strings.TrimSpace(p) + if p == "" { + continue + } + if len(p) > maxPatternLen { + response.BadRequest(c, "Signature pattern too long (max 500 characters)") + return + } + cleanedPatterns = append(cleanedPatterns, p) + } + + settings := &service.RectifierSettings{ + Enabled: req.Enabled, + ThinkingSignatureEnabled: req.ThinkingSignatureEnabled, + ThinkingBudgetEnabled: req.ThinkingBudgetEnabled, + APIKeySignatureEnabled: req.APIKeySignatureEnabled, + APIKeySignaturePatterns: cleanedPatterns, + } + + if err := h.settingService.SetRectifierSettings(c.Request.Context(), settings); err != nil { + response.BadRequest(c, err.Error()) + return + } + + // 重新获取设置返回 + updatedSettings, err := h.settingService.GetRectifierSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + updatedPatterns := updatedSettings.APIKeySignaturePatterns + if updatedPatterns == nil { + updatedPatterns = []string{} + } + response.Success(c, dto.RectifierSettings{ + Enabled: updatedSettings.Enabled, + ThinkingSignatureEnabled: updatedSettings.ThinkingSignatureEnabled, + ThinkingBudgetEnabled: updatedSettings.ThinkingBudgetEnabled, + APIKeySignatureEnabled: updatedSettings.APIKeySignatureEnabled, + APIKeySignaturePatterns: updatedPatterns, + }) +} + +// GetBetaPolicySettings 获取 Beta 策略配置 +// GET /api/v1/admin/settings/beta-policy +func (h *SettingHandler) GetBetaPolicySettings(c *gin.Context) { + settings, err := h.settingService.GetBetaPolicySettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + rules := make([]dto.BetaPolicyRule, len(settings.Rules)) + for i, r := range settings.Rules { + rules[i] = dto.BetaPolicyRule(r) + } + response.Success(c, dto.BetaPolicySettings{Rules: rules}) +} + +// UpdateBetaPolicySettingsRequest 更新 Beta 策略配置请求 +type UpdateBetaPolicySettingsRequest struct { + Rules []dto.BetaPolicyRule `json:"rules"` +} + +// UpdateBetaPolicySettings 更新 Beta 策略配置 +// PUT /api/v1/admin/settings/beta-policy +func (h *SettingHandler) UpdateBetaPolicySettings(c *gin.Context) { + var req UpdateBetaPolicySettingsRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + rules := make([]service.BetaPolicyRule, len(req.Rules)) + for i, r := range req.Rules { + rules[i] = service.BetaPolicyRule(r) + } + + settings := &service.BetaPolicySettings{Rules: rules} + if err := h.settingService.SetBetaPolicySettings(c.Request.Context(), settings); err != nil { + response.BadRequest(c, err.Error()) + return + } + + // Re-fetch to return updated settings + updated, err := h.settingService.GetBetaPolicySettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + outRules := make([]dto.BetaPolicyRule, len(updated.Rules)) + for i, r := range updated.Rules { + outRules[i] = dto.BetaPolicyRule(r) + } + response.Success(c, dto.BetaPolicySettings{Rules: outRules}) +} + +// UpdateStreamTimeoutSettingsRequest 更新流超时配置请求 +type UpdateStreamTimeoutSettingsRequest struct { + Enabled bool `json:"enabled"` + Action string `json:"action"` + TempUnschedMinutes int `json:"temp_unsched_minutes"` + ThresholdCount int `json:"threshold_count"` + ThresholdWindowMinutes int `json:"threshold_window_minutes"` +} + +// UpdateStreamTimeoutSettings 更新流超时处理配置 +// PUT /api/v1/admin/settings/stream-timeout +func (h *SettingHandler) UpdateStreamTimeoutSettings(c *gin.Context) { + var req UpdateStreamTimeoutSettingsRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + settings := &service.StreamTimeoutSettings{ + Enabled: req.Enabled, + Action: req.Action, + TempUnschedMinutes: req.TempUnschedMinutes, + ThresholdCount: req.ThresholdCount, + ThresholdWindowMinutes: req.ThresholdWindowMinutes, + } + + if err := h.settingService.SetStreamTimeoutSettings(c.Request.Context(), settings); err != nil { + response.BadRequest(c, err.Error()) + return + } + + // 重新获取设置返回 + updatedSettings, err := h.settingService.GetStreamTimeoutSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + response.Success(c, dto.StreamTimeoutSettings{ + Enabled: updatedSettings.Enabled, + Action: updatedSettings.Action, + TempUnschedMinutes: updatedSettings.TempUnschedMinutes, + ThresholdCount: updatedSettings.ThresholdCount, + ThresholdWindowMinutes: updatedSettings.ThresholdWindowMinutes, + }) +} + +// GetWebSearchEmulationConfig 获取 Web Search 模拟配置 +// GET /api/v1/admin/settings/web-search-emulation +func (h *SettingHandler) GetWebSearchEmulationConfig(c *gin.Context) { + cfg, err := h.settingService.GetWebSearchEmulationConfig(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, service.PopulateWebSearchUsage(c.Request.Context(), cfg)) +} + +// UpdateWebSearchEmulationConfig 更新 Web Search 模拟配置 +// PUT /api/v1/admin/settings/web-search-emulation +func (h *SettingHandler) UpdateWebSearchEmulationConfig(c *gin.Context) { + var cfg service.WebSearchEmulationConfig + if err := c.ShouldBindJSON(&cfg); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + if err := h.settingService.SaveWebSearchEmulationConfig(c.Request.Context(), &cfg); err != nil { + response.ErrorFrom(c, err) + return + } + + // Re-read (with sanitized api keys) to return current state + updated, err := h.settingService.GetWebSearchEmulationConfig(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, service.PopulateWebSearchUsage(c.Request.Context(), updated)) +} + +// ResetWebSearchUsage 重置指定 provider 的配额用量 +// POST /api/v1/admin/settings/web-search-emulation/reset-usage +func (h *SettingHandler) ResetWebSearchUsage(c *gin.Context) { + var req struct { + ProviderType string `json:"provider_type"` + } + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + if req.ProviderType == "" { + response.BadRequest(c, "provider_type is required") + return + } + if err := service.ResetWebSearchUsage(c.Request.Context(), req.ProviderType); err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, nil) +} + +// TestWebSearchEmulation 测试 Web Search 搜索 +// POST /api/v1/admin/settings/web-search-emulation/test +func (h *SettingHandler) TestWebSearchEmulation(c *gin.Context) { + var req struct { + Query string `json:"query"` + } + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + if strings.TrimSpace(req.Query) == "" { + req.Query = "搜索今年世界大事件" + } + + result, err := service.TestWebSearch(c.Request.Context(), req.Query) + if err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, result) +} diff --git a/backend/internal/handler/admin/setting_handler_update.go b/backend/internal/handler/admin/setting_handler_update.go new file mode 100644 index 0000000000..b4c1f3afc1 --- /dev/null +++ b/backend/internal/handler/admin/setting_handler_update.go @@ -0,0 +1,1985 @@ +package admin + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "net/http" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/handler/dto" + "github.com/Wei-Shaw/sub2api/internal/pkg/response" + "github.com/Wei-Shaw/sub2api/internal/service" + + "github.com/gin-gonic/gin" +) + +// UpdateSettingsRequest 更新设置请求 +type UpdateSettingsRequest struct { + // 注册设置 + RegistrationEnabled bool `json:"registration_enabled"` + EmailVerifyEnabled bool `json:"email_verify_enabled"` + RegistrationEmailSuffixWhitelist []string `json:"registration_email_suffix_whitelist"` + PromoCodeEnabled bool `json:"promo_code_enabled"` + PasswordResetEnabled bool `json:"password_reset_enabled"` + FrontendURL string `json:"frontend_url"` + InvitationCodeEnabled bool `json:"invitation_code_enabled"` + TotpEnabled bool `json:"totp_enabled"` // TOTP 双因素认证 + LoginAgreementEnabled bool `json:"login_agreement_enabled"` + LoginAgreementMode string `json:"login_agreement_mode"` + LoginAgreementUpdatedAt string `json:"login_agreement_updated_at"` + LoginAgreementDocuments []dto.LoginAgreementDocument `json:"login_agreement_documents"` + + // 邮件服务设置 + SMTPHost string `json:"smtp_host"` + SMTPPort int `json:"smtp_port"` + SMTPUsername string `json:"smtp_username"` + SMTPPassword string `json:"smtp_password"` + SMTPFrom string `json:"smtp_from_email"` + SMTPFromName string `json:"smtp_from_name"` + SMTPUseTLS bool `json:"smtp_use_tls"` + + // Cloudflare Turnstile 设置 + TurnstileEnabled bool `json:"turnstile_enabled"` + TurnstileSiteKey string `json:"turnstile_site_key"` + TurnstileSecretKey string `json:"turnstile_secret_key"` + + // API Key IP 访问控制设置 + APIKeyACLTrustForwardedIP *bool `json:"api_key_acl_trust_forwarded_ip"` + + // LinuxDo Connect OAuth 登录 + LinuxDoConnectEnabled bool `json:"linuxdo_connect_enabled"` + LinuxDoConnectClientID string `json:"linuxdo_connect_client_id"` + LinuxDoConnectClientSecret string `json:"linuxdo_connect_client_secret"` + LinuxDoConnectRedirectURL string `json:"linuxdo_connect_redirect_url"` + + // DingTalk Connect OAuth 登录 + DingTalkConnectEnabled bool `json:"dingtalk_connect_enabled"` + DingTalkConnectClientID string `json:"dingtalk_connect_client_id"` + DingTalkConnectClientSecret string `json:"dingtalk_connect_client_secret"` + DingTalkConnectRedirectURL string `json:"dingtalk_connect_redirect_url"` + DingTalkConnectCorpRestrictionPolicy string `json:"dingtalk_connect_corp_restriction_policy"` + DingTalkConnectInternalCorpID string `json:"dingtalk_connect_internal_corp_id"` + DingTalkConnectBypassRegistration bool `json:"dingtalk_connect_bypass_registration"` + DingTalkConnectSyncCorpEmail bool `json:"dingtalk_connect_sync_corp_email"` + DingTalkConnectSyncDisplayName bool `json:"dingtalk_connect_sync_display_name"` + DingTalkConnectSyncDept bool `json:"dingtalk_connect_sync_dept"` + DingTalkConnectSyncCorpEmailAttrKey string `json:"dingtalk_connect_sync_corp_email_attr_key"` + DingTalkConnectSyncDisplayNameAttrKey string `json:"dingtalk_connect_sync_display_name_attr_key"` + DingTalkConnectSyncDeptAttrKey string `json:"dingtalk_connect_sync_dept_attr_key"` + DingTalkConnectSyncCorpEmailAttrName string `json:"dingtalk_connect_sync_corp_email_attr_name"` + DingTalkConnectSyncDisplayNameAttrName string `json:"dingtalk_connect_sync_display_name_attr_name"` + DingTalkConnectSyncDeptAttrName string `json:"dingtalk_connect_sync_dept_attr_name"` + + // WeChat Connect OAuth 登录 + WeChatConnectEnabled bool `json:"wechat_connect_enabled"` + WeChatConnectAppID string `json:"wechat_connect_app_id"` + WeChatConnectAppSecret string `json:"wechat_connect_app_secret"` + WeChatConnectOpenAppID string `json:"wechat_connect_open_app_id"` + WeChatConnectOpenAppSecret string `json:"wechat_connect_open_app_secret"` + WeChatConnectMPAppID string `json:"wechat_connect_mp_app_id"` + WeChatConnectMPAppSecret string `json:"wechat_connect_mp_app_secret"` + WeChatConnectMobileAppID string `json:"wechat_connect_mobile_app_id"` + WeChatConnectMobileAppSecret string `json:"wechat_connect_mobile_app_secret"` + WeChatConnectOpenEnabled bool `json:"wechat_connect_open_enabled"` + WeChatConnectMPEnabled bool `json:"wechat_connect_mp_enabled"` + WeChatConnectMobileEnabled bool `json:"wechat_connect_mobile_enabled"` + WeChatConnectMode string `json:"wechat_connect_mode"` + WeChatConnectScopes string `json:"wechat_connect_scopes"` + WeChatConnectRedirectURL string `json:"wechat_connect_redirect_url"` + WeChatConnectFrontendRedirectURL string `json:"wechat_connect_frontend_redirect_url"` + + // Generic OIDC OAuth 登录 + OIDCConnectEnabled bool `json:"oidc_connect_enabled"` + OIDCConnectProviderName string `json:"oidc_connect_provider_name"` + OIDCConnectClientID string `json:"oidc_connect_client_id"` + OIDCConnectClientSecret string `json:"oidc_connect_client_secret"` + OIDCConnectIssuerURL string `json:"oidc_connect_issuer_url"` + OIDCConnectDiscoveryURL string `json:"oidc_connect_discovery_url"` + OIDCConnectAuthorizeURL string `json:"oidc_connect_authorize_url"` + OIDCConnectTokenURL string `json:"oidc_connect_token_url"` + OIDCConnectUserInfoURL string `json:"oidc_connect_userinfo_url"` + OIDCConnectJWKSURL string `json:"oidc_connect_jwks_url"` + OIDCConnectScopes string `json:"oidc_connect_scopes"` + OIDCConnectRedirectURL string `json:"oidc_connect_redirect_url"` + OIDCConnectFrontendRedirectURL string `json:"oidc_connect_frontend_redirect_url"` + OIDCConnectTokenAuthMethod string `json:"oidc_connect_token_auth_method"` + OIDCConnectUsePKCE *bool `json:"oidc_connect_use_pkce"` + OIDCConnectValidateIDToken *bool `json:"oidc_connect_validate_id_token"` + OIDCConnectAllowedSigningAlgs string `json:"oidc_connect_allowed_signing_algs"` + OIDCConnectClockSkewSeconds int `json:"oidc_connect_clock_skew_seconds"` + OIDCConnectRequireEmailVerified bool `json:"oidc_connect_require_email_verified"` + OIDCConnectUserInfoEmailPath string `json:"oidc_connect_userinfo_email_path"` + OIDCConnectUserInfoIDPath string `json:"oidc_connect_userinfo_id_path"` + OIDCConnectUserInfoUsernamePath string `json:"oidc_connect_userinfo_username_path"` + + GitHubOAuthEnabled bool `json:"github_oauth_enabled"` + GitHubOAuthClientID string `json:"github_oauth_client_id"` + GitHubOAuthClientSecret string `json:"github_oauth_client_secret"` + GitHubOAuthRedirectURL string `json:"github_oauth_redirect_url"` + GitHubOAuthFrontendRedirectURL string `json:"github_oauth_frontend_redirect_url"` + GoogleOAuthEnabled bool `json:"google_oauth_enabled"` + GoogleOAuthClientID string `json:"google_oauth_client_id"` + GoogleOAuthClientSecret string `json:"google_oauth_client_secret"` + GoogleOAuthRedirectURL string `json:"google_oauth_redirect_url"` + GoogleOAuthFrontendRedirectURL string `json:"google_oauth_frontend_redirect_url"` + + // OEM设置 + SiteName string `json:"site_name"` + SiteLogo string `json:"site_logo"` + SiteSubtitle string `json:"site_subtitle"` + APIBaseURL string `json:"api_base_url"` + ContactInfo string `json:"contact_info"` + DocURL string `json:"doc_url"` + HomeContent string `json:"home_content"` + HideCcsImportButton bool `json:"hide_ccs_import_button"` + PurchaseSubscriptionEnabled *bool `json:"purchase_subscription_enabled"` + PurchaseSubscriptionURL *string `json:"purchase_subscription_url"` + TableDefaultPageSize int `json:"table_default_page_size"` + TablePageSizeOptions []int `json:"table_page_size_options"` + CustomMenuItems *[]dto.CustomMenuItem `json:"custom_menu_items"` + CustomEndpoints *[]dto.CustomEndpoint `json:"custom_endpoints"` + + // 默认配置 + DefaultConcurrency int `json:"default_concurrency"` + DefaultBalance float64 `json:"default_balance"` + AffiliateRebateRate *float64 `json:"affiliate_rebate_rate"` + AffiliateRebateFreezeHours *int `json:"affiliate_rebate_freeze_hours"` + AffiliateRebateDurationDays *int `json:"affiliate_rebate_duration_days"` + AffiliateRebatePerInviteeCap *float64 `json:"affiliate_rebate_per_invitee_cap"` + DefaultUserRPMLimit int `json:"default_user_rpm_limit"` + DefaultSubscriptions []dto.DefaultSubscriptionSetting `json:"default_subscriptions"` + AuthSourceDefaultEmailBalance *float64 `json:"auth_source_default_email_balance"` + AuthSourceDefaultEmailConcurrency *int `json:"auth_source_default_email_concurrency"` + AuthSourceDefaultEmailSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_email_subscriptions"` + AuthSourceDefaultEmailGrantOnSignup *bool `json:"auth_source_default_email_grant_on_signup"` + AuthSourceDefaultEmailGrantOnFirstBind *bool `json:"auth_source_default_email_grant_on_first_bind"` + AuthSourceDefaultLinuxDoBalance *float64 `json:"auth_source_default_linuxdo_balance"` + AuthSourceDefaultLinuxDoConcurrency *int `json:"auth_source_default_linuxdo_concurrency"` + AuthSourceDefaultLinuxDoSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_linuxdo_subscriptions"` + AuthSourceDefaultLinuxDoGrantOnSignup *bool `json:"auth_source_default_linuxdo_grant_on_signup"` + AuthSourceDefaultLinuxDoGrantOnFirstBind *bool `json:"auth_source_default_linuxdo_grant_on_first_bind"` + AuthSourceDefaultOIDCBalance *float64 `json:"auth_source_default_oidc_balance"` + AuthSourceDefaultOIDCConcurrency *int `json:"auth_source_default_oidc_concurrency"` + AuthSourceDefaultOIDCSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_oidc_subscriptions"` + AuthSourceDefaultOIDCGrantOnSignup *bool `json:"auth_source_default_oidc_grant_on_signup"` + AuthSourceDefaultOIDCGrantOnFirstBind *bool `json:"auth_source_default_oidc_grant_on_first_bind"` + AuthSourceDefaultWeChatBalance *float64 `json:"auth_source_default_wechat_balance"` + AuthSourceDefaultWeChatConcurrency *int `json:"auth_source_default_wechat_concurrency"` + AuthSourceDefaultWeChatSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_wechat_subscriptions"` + AuthSourceDefaultWeChatGrantOnSignup *bool `json:"auth_source_default_wechat_grant_on_signup"` + AuthSourceDefaultWeChatGrantOnFirstBind *bool `json:"auth_source_default_wechat_grant_on_first_bind"` + AuthSourceDefaultGitHubBalance *float64 `json:"auth_source_default_github_balance"` + AuthSourceDefaultGitHubConcurrency *int `json:"auth_source_default_github_concurrency"` + AuthSourceDefaultGitHubSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_github_subscriptions"` + AuthSourceDefaultGitHubGrantOnSignup *bool `json:"auth_source_default_github_grant_on_signup"` + AuthSourceDefaultGitHubGrantOnFirstBind *bool `json:"auth_source_default_github_grant_on_first_bind"` + AuthSourceDefaultGoogleBalance *float64 `json:"auth_source_default_google_balance"` + AuthSourceDefaultGoogleConcurrency *int `json:"auth_source_default_google_concurrency"` + AuthSourceDefaultGoogleSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_google_subscriptions"` + AuthSourceDefaultGoogleGrantOnSignup *bool `json:"auth_source_default_google_grant_on_signup"` + AuthSourceDefaultGoogleGrantOnFirstBind *bool `json:"auth_source_default_google_grant_on_first_bind"` + AuthSourceDefaultDingTalkBalance *float64 `json:"auth_source_default_dingtalk_balance"` + AuthSourceDefaultDingTalkConcurrency *int `json:"auth_source_default_dingtalk_concurrency"` + AuthSourceDefaultDingTalkSubscriptions *[]dto.DefaultSubscriptionSetting `json:"auth_source_default_dingtalk_subscriptions"` + AuthSourceDefaultDingTalkGrantOnSignup *bool `json:"auth_source_default_dingtalk_grant_on_signup"` + AuthSourceDefaultDingTalkGrantOnFirstBind *bool `json:"auth_source_default_dingtalk_grant_on_first_bind"` + ForceEmailOnThirdPartySignup *bool `json:"force_email_on_third_party_signup"` + + // Model fallback configuration + EnableModelFallback bool `json:"enable_model_fallback"` + FallbackModelAnthropic string `json:"fallback_model_anthropic"` + FallbackModelOpenAI string `json:"fallback_model_openai"` + FallbackModelGemini string `json:"fallback_model_gemini"` + FallbackModelAntigravity string `json:"fallback_model_antigravity"` + + // Identity patch configuration (Claude -> Gemini) + EnableIdentityPatch bool `json:"enable_identity_patch"` + IdentityPatchPrompt string `json:"identity_patch_prompt"` + + // Ops monitoring (vNext) + OpsMonitoringEnabled *bool `json:"ops_monitoring_enabled"` + OpsRealtimeMonitoringEnabled *bool `json:"ops_realtime_monitoring_enabled"` + OpsQueryModeDefault *string `json:"ops_query_mode_default"` + OpsMetricsIntervalSeconds *int `json:"ops_metrics_interval_seconds"` + + MinClaudeCodeVersion string `json:"min_claude_code_version"` + MaxClaudeCodeVersion string `json:"max_claude_code_version"` + + // 分组隔离 + AllowUngroupedKeyScheduling bool `json:"allow_ungrouped_key_scheduling"` + + // Backend Mode + BackendModeEnabled bool `json:"backend_mode_enabled"` + + // Gateway forwarding behavior + EnableFingerprintUnification *bool `json:"enable_fingerprint_unification"` + EnableMetadataPassthrough *bool `json:"enable_metadata_passthrough"` + EnableCCHSigning *bool `json:"enable_cch_signing"` + EnableClaudeOAuthSystemPromptInjection *bool `json:"enable_claude_oauth_system_prompt_injection"` + ClaudeOAuthSystemPrompt *string `json:"claude_oauth_system_prompt"` + ClaudeOAuthSystemPromptBlocks *string `json:"claude_oauth_system_prompt_blocks"` + EnableAnthropicCacheTTL1hInjection *bool `json:"enable_anthropic_cache_ttl_1h_injection"` + RewriteMessageCacheControl *bool `json:"rewrite_message_cache_control"` + EnableClientDatelineNormalization *bool `json:"enable_client_dateline_normalization"` + AntigravityUserAgentVersion *string `json:"antigravity_user_agent_version"` + OpenAICodexUserAgent *string `json:"openai_codex_user_agent"` + + // codex_cli_only 加固(global-only) + MinCodexVersion string `json:"min_codex_version"` + MaxCodexVersion string `json:"max_codex_version"` + CodexCLIOnlyBlacklist string `json:"codex_cli_only_blacklist"` + CodexCLIOnlyWhitelist string `json:"codex_cli_only_whitelist"` + CodexCLIOnlyAllowAppServerClients *bool `json:"codex_cli_only_allow_app_server_clients"` + CodexCLIOnlyEngineFingerprintSignals string `json:"codex_cli_only_engine_fingerprint_signals"` + + // Payment visible method routing + PaymentVisibleMethodAlipaySource *string `json:"payment_visible_method_alipay_source"` + PaymentVisibleMethodWxpaySource *string `json:"payment_visible_method_wxpay_source"` + PaymentVisibleMethodAlipayEnabled *bool `json:"payment_visible_method_alipay_enabled"` + PaymentVisibleMethodWxpayEnabled *bool `json:"payment_visible_method_wxpay_enabled"` + + // OpenAI account scheduling + OpenAIAdvancedSchedulerEnabled *bool `json:"openai_advanced_scheduler_enabled"` + OpenAIAdvancedSchedulerStickyWeightedEnabled *bool `json:"openai_advanced_scheduler_sticky_weighted_enabled"` + OpenAIAdvancedSchedulerSubscriptionPriorityEnabled *bool `json:"openai_advanced_scheduler_subscription_priority_enabled"` + OpenAIAdvancedSchedulerLBTopK *string `json:"openai_advanced_scheduler_lb_top_k"` + OpenAIAdvancedSchedulerWeightPriority *string `json:"openai_advanced_scheduler_weight_priority"` + OpenAIAdvancedSchedulerWeightLoad *string `json:"openai_advanced_scheduler_weight_load"` + OpenAIAdvancedSchedulerWeightQueue *string `json:"openai_advanced_scheduler_weight_queue"` + OpenAIAdvancedSchedulerWeightErrorRate *string `json:"openai_advanced_scheduler_weight_error_rate"` + OpenAIAdvancedSchedulerWeightTTFT *string `json:"openai_advanced_scheduler_weight_ttft"` + OpenAIAdvancedSchedulerWeightReset *string `json:"openai_advanced_scheduler_weight_reset"` + OpenAIAdvancedSchedulerWeightQuotaHeadroom *string `json:"openai_advanced_scheduler_weight_quota_headroom"` + OpenAIAdvancedSchedulerWeightPreviousResponse *string `json:"openai_advanced_scheduler_weight_previous_response"` + OpenAIAdvancedSchedulerWeightSessionSticky *string `json:"openai_advanced_scheduler_weight_session_sticky"` + + // 余额不足提醒 + BalanceLowNotifyEnabled *bool `json:"balance_low_notify_enabled"` + BalanceLowNotifyThreshold *float64 `json:"balance_low_notify_threshold"` + BalanceLowNotifyRechargeURL *string `json:"balance_low_notify_recharge_url"` + SubscriptionExpiryNotifyEnabled *bool `json:"subscription_expiry_notify_enabled"` + AccountQuotaNotifyEnabled *bool `json:"account_quota_notify_enabled"` + AccountQuotaNotifyEmails *[]dto.NotifyEmailEntry `json:"account_quota_notify_emails"` + + // Payment configuration (integrated into settings, full replace) + PaymentEnabled *bool `json:"payment_enabled"` + PaymentMinAmount *float64 `json:"payment_min_amount"` + PaymentMaxAmount *float64 `json:"payment_max_amount"` + PaymentDailyLimit *float64 `json:"payment_daily_limit"` + PaymentOrderTimeoutMin *int `json:"payment_order_timeout_minutes"` + PaymentMaxPendingOrders *int `json:"payment_max_pending_orders"` + PaymentEnabledTypes []string `json:"payment_enabled_types"` + PaymentBalanceDisabled *bool `json:"payment_balance_disabled"` + PaymentBalanceRechargeMultiplier *float64 `json:"payment_balance_recharge_multiplier"` + PaymentSubscriptionUSDToCNYRate *float64 `json:"payment_subscription_usd_to_cny_rate"` + PaymentRechargeFeeRate *float64 `json:"payment_recharge_fee_rate"` + PaymentLoadBalanceStrat *string `json:"payment_load_balance_strategy"` + PaymentProductNamePrefix *string `json:"payment_product_name_prefix"` + PaymentProductNameSuffix *string `json:"payment_product_name_suffix"` + PaymentHelpImageURL *string `json:"payment_help_image_url"` + PaymentHelpText *string `json:"payment_help_text"` + + // Cancel rate limit + PaymentCancelRateLimitEnabled *bool `json:"payment_cancel_rate_limit_enabled"` + PaymentCancelRateLimitMax *int `json:"payment_cancel_rate_limit_max"` + PaymentCancelRateLimitWindow *int `json:"payment_cancel_rate_limit_window"` + PaymentCancelRateLimitUnit *string `json:"payment_cancel_rate_limit_unit"` + PaymentCancelRateLimitMode *string `json:"payment_cancel_rate_limit_window_mode"` + + // Force Alipay mobile clients to use QR code payment instead of mobile redirect + PaymentAlipayForceQRCode *bool `json:"payment_alipay_force_qrcode"` + + // Channel Monitor feature switch + ChannelMonitorEnabled *bool `json:"channel_monitor_enabled"` + ChannelMonitorDefaultIntervalSeconds *int `json:"channel_monitor_default_interval_seconds"` + + // Available Channels feature switch (user-facing) + AvailableChannelsEnabled *bool `json:"available_channels_enabled"` + + // Affiliate (邀请返利) feature switch + AffiliateEnabled *bool `json:"affiliate_enabled"` + + // 风控中心功能开关 + RiskControlEnabled *bool `json:"risk_control_enabled"` + + // cyber 会话屏蔽开关 + TTL + CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"` + CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"` + + // OpenAI fast/flex policy (optional, only updated when provided) + OpenAIFastPolicySettings *dto.OpenAIFastPolicySettings `json:"openai_fast_policy_settings,omitempty"` + + // 系统全局 platform quota 默认值(整体替换语义:nil = 不修改,non-nil = 整体覆盖)。 + DefaultPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"default_platform_quotas"` + + // auth-source 层 platform quota 覆盖(override 语义:nil = 不修改,non-nil = 整体覆盖该 source 的 quota 配置)。 + AuthSourceEmailPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_email_platform_quotas"` + AuthSourceLinuxDoPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_linuxdo_platform_quotas"` + AuthSourceOIDCPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_oidc_platform_quotas"` + AuthSourceWeChatPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_wechat_platform_quotas"` + AuthSourceGitHubPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_github_platform_quotas"` + AuthSourceGooglePlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_google_platform_quotas"` + AuthSourceDingTalkPlatformQuotas map[string]*service.DefaultPlatformQuotaSetting `json:"auth_source_default_dingtalk_platform_quotas"` + + AllowUserViewErrorRequests *bool `json:"allow_user_view_error_requests"` +} + +// UpdateSettings 更新系统设置 +// PUT /api/v1/admin/settings +func (h *SettingHandler) UpdateSettings(c *gin.Context) { + var req UpdateSettingsRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + previousSettings, err := h.settingService.GetAllSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + previousAuthSourceDefaults, err := h.settingService.GetAuthSourceDefaultSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + + // 验证参数 + if req.DefaultConcurrency < 1 { + req.DefaultConcurrency = 1 + } + if req.DefaultBalance < 0 { + req.DefaultBalance = 0 + } + affiliateRebateRate := previousSettings.AffiliateRebateRate + if req.AffiliateRebateRate != nil { + affiliateRebateRate = *req.AffiliateRebateRate + } + if affiliateRebateRate < service.AffiliateRebateRateMin { + affiliateRebateRate = service.AffiliateRebateRateMin + } + if affiliateRebateRate > service.AffiliateRebateRateMax { + affiliateRebateRate = service.AffiliateRebateRateMax + } + affiliateRebateFreezeHours := previousSettings.AffiliateRebateFreezeHours + if req.AffiliateRebateFreezeHours != nil { + affiliateRebateFreezeHours = *req.AffiliateRebateFreezeHours + } + if affiliateRebateFreezeHours < 0 { + affiliateRebateFreezeHours = service.AffiliateRebateFreezeHoursDefault + } + if affiliateRebateFreezeHours > service.AffiliateRebateFreezeHoursMax { + affiliateRebateFreezeHours = service.AffiliateRebateFreezeHoursMax + } + affiliateRebateDurationDays := previousSettings.AffiliateRebateDurationDays + if req.AffiliateRebateDurationDays != nil { + affiliateRebateDurationDays = *req.AffiliateRebateDurationDays + } + if affiliateRebateDurationDays < 0 { + affiliateRebateDurationDays = service.AffiliateRebateDurationDaysDefault + } + if affiliateRebateDurationDays > service.AffiliateRebateDurationDaysMax { + affiliateRebateDurationDays = service.AffiliateRebateDurationDaysMax + } + affiliateRebatePerInviteeCap := previousSettings.AffiliateRebatePerInviteeCap + if req.AffiliateRebatePerInviteeCap != nil { + affiliateRebatePerInviteeCap = *req.AffiliateRebatePerInviteeCap + } + if affiliateRebatePerInviteeCap < 0 { + affiliateRebatePerInviteeCap = service.AffiliateRebatePerInviteeCapDefault + } + // 通用表格配置:兼容旧客户端未传字段时保留当前值。 + if req.TableDefaultPageSize <= 0 { + req.TableDefaultPageSize = previousSettings.TableDefaultPageSize + } + if req.TablePageSizeOptions == nil { + req.TablePageSizeOptions = previousSettings.TablePageSizeOptions + } + req.SMTPHost = strings.TrimSpace(req.SMTPHost) + req.SMTPUsername = strings.TrimSpace(req.SMTPUsername) + req.SMTPPassword = strings.TrimSpace(req.SMTPPassword) + req.SMTPFrom = strings.TrimSpace(req.SMTPFrom) + req.SMTPFromName = strings.TrimSpace(req.SMTPFromName) + if req.SMTPPort <= 0 { + req.SMTPPort = 587 + } + req.DefaultSubscriptions = normalizeDefaultSubscriptions(req.DefaultSubscriptions) + req.AuthSourceDefaultEmailSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultEmailSubscriptions) + req.AuthSourceDefaultLinuxDoSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultLinuxDoSubscriptions) + req.AuthSourceDefaultOIDCSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultOIDCSubscriptions) + req.AuthSourceDefaultWeChatSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultWeChatSubscriptions) + req.AuthSourceDefaultDingTalkSubscriptions = normalizeOptionalDefaultSubscriptions(req.AuthSourceDefaultDingTalkSubscriptions) + + // SMTP 配置保护:如果请求中 smtp_host 为空但数据库中已有配置,则保留已有 SMTP 配置 + // 防止前端加载设置失败时空表单覆盖已保存的 SMTP 配置 + if req.SMTPHost == "" && previousSettings.SMTPHost != "" { + req.SMTPHost = previousSettings.SMTPHost + req.SMTPPort = previousSettings.SMTPPort + req.SMTPUsername = previousSettings.SMTPUsername + req.SMTPFrom = previousSettings.SMTPFrom + req.SMTPFromName = previousSettings.SMTPFromName + req.SMTPUseTLS = previousSettings.SMTPUseTLS + } + + // Turnstile 参数验证 + if req.TurnstileEnabled { + // 检查必填字段 + if req.TurnstileSiteKey == "" { + response.BadRequest(c, "Turnstile Site Key is required when enabled") + return + } + // 如果未提供 secret key,使用已保存的值(留空保留当前值) + if req.TurnstileSecretKey == "" { + if previousSettings.TurnstileSecretKey == "" { + response.BadRequest(c, "Turnstile Secret Key is required when enabled") + return + } + req.TurnstileSecretKey = previousSettings.TurnstileSecretKey + } + + // 当 site_key 或 secret_key 任一变化时验证(避免配置错误导致无法登录) + siteKeyChanged := previousSettings.TurnstileSiteKey != req.TurnstileSiteKey + secretKeyChanged := previousSettings.TurnstileSecretKey != req.TurnstileSecretKey + if siteKeyChanged || secretKeyChanged { + if err := h.turnstileService.ValidateSecretKey(c.Request.Context(), req.TurnstileSecretKey); err != nil { + response.ErrorFrom(c, err) + return + } + } + } + + // TOTP 双因素认证参数验证 + // 只有手动配置了加密密钥才允许启用 TOTP 功能 + if req.TotpEnabled && !previousSettings.TotpEnabled { + // 尝试启用 TOTP,检查加密密钥是否已手动配置 + if !h.settingService.IsTotpEncryptionKeyConfigured() { + response.BadRequest(c, "Cannot enable TOTP: TOTP_ENCRYPTION_KEY environment variable must be configured first. Generate a key with 'openssl rand -hex 32' and set it in your environment.") + return + } + } + loginAgreementMode := strings.ToLower(strings.TrimSpace(req.LoginAgreementMode)) + if loginAgreementMode == "" { + loginAgreementMode = strings.ToLower(strings.TrimSpace(previousSettings.LoginAgreementMode)) + } + switch loginAgreementMode { + case "", "modal": + loginAgreementMode = "modal" + case "checkbox": + default: + response.BadRequest(c, "Login agreement mode must be modal or checkbox") + return + } + loginAgreementUpdatedAt := strings.TrimSpace(req.LoginAgreementUpdatedAt) + if loginAgreementUpdatedAt == "" { + loginAgreementUpdatedAt = strings.TrimSpace(previousSettings.LoginAgreementUpdatedAt) + } + loginAgreementDocuments := loginAgreementDocumentsToService(req.LoginAgreementDocuments) + if len(loginAgreementDocuments) == 0 { + loginAgreementDocuments = previousSettings.LoginAgreementDocuments + } + for _, doc := range loginAgreementDocuments { + if strings.TrimSpace(doc.Title) == "" { + response.BadRequest(c, "Login agreement document title is required") + return + } + if len(doc.Title) > 80 { + response.BadRequest(c, "Login agreement document title is too long (max 80 characters)") + return + } + if len(doc.ContentMD) > 200*1024 { + response.BadRequest(c, "Login agreement document content is too large (max 200KB)") + return + } + } + if req.LoginAgreementEnabled && len(loginAgreementDocuments) == 0 { + response.BadRequest(c, "Login agreement documents are required when enabled") + return + } + + // LinuxDo Connect 参数验证 + if req.LinuxDoConnectEnabled { + req.LinuxDoConnectClientID = strings.TrimSpace(req.LinuxDoConnectClientID) + req.LinuxDoConnectClientSecret = strings.TrimSpace(req.LinuxDoConnectClientSecret) + req.LinuxDoConnectRedirectURL = strings.TrimSpace(req.LinuxDoConnectRedirectURL) + + if req.LinuxDoConnectClientID == "" { + response.BadRequest(c, "LinuxDo Client ID is required when enabled") + return + } + if req.LinuxDoConnectRedirectURL == "" { + response.BadRequest(c, "LinuxDo Redirect URL is required when enabled") + return + } + if err := config.ValidateAbsoluteHTTPURL(req.LinuxDoConnectRedirectURL); err != nil { + response.BadRequest(c, "LinuxDo Redirect URL must be an absolute http(s) URL") + return + } + + // 如果未提供 client_secret,则保留现有值(如有)。 + if req.LinuxDoConnectClientSecret == "" { + if previousSettings.LinuxDoConnectClientSecret == "" { + response.BadRequest(c, "LinuxDo Client Secret is required when enabled") + return + } + req.LinuxDoConnectClientSecret = previousSettings.LinuxDoConnectClientSecret + } + } + + // DingTalk Connect 参数验证 + // 防御性:任何写入路径上把已废弃的 corp_restriction_policy=whitelist 入参 coerce 为 none, + // 避免任何直连 admin API 的客户端把死值写回 DB(前端 UI 已无此选项)。 + req.DingTalkConnectCorpRestrictionPolicy = service.CoerceDingTalkCorpPolicyForWrite(req.DingTalkConnectCorpRestrictionPolicy) + + if req.DingTalkConnectEnabled { + req.DingTalkConnectClientID = strings.TrimSpace(req.DingTalkConnectClientID) + req.DingTalkConnectClientSecret = strings.TrimSpace(req.DingTalkConnectClientSecret) + req.DingTalkConnectRedirectURL = strings.TrimSpace(req.DingTalkConnectRedirectURL) + req.DingTalkConnectCorpRestrictionPolicy = strings.TrimSpace(req.DingTalkConnectCorpRestrictionPolicy) + req.DingTalkConnectInternalCorpID = strings.TrimSpace(req.DingTalkConnectInternalCorpID) + + if req.DingTalkConnectClientID == "" { + response.BadRequest(c, "DingTalk Client ID is required when enabled") + return + } + if req.DingTalkConnectRedirectURL == "" { + response.BadRequest(c, "DingTalk Redirect URL is required when enabled") + return + } + if err := config.ValidateAbsoluteHTTPURL(req.DingTalkConnectRedirectURL); err != nil { + response.BadRequest(c, "DingTalk Redirect URL must be an absolute http(s) URL") + return + } + + // 如果未提供 client_secret,则保留现有值(如有)。 + if req.DingTalkConnectClientSecret == "" { + if previousSettings.DingTalkConnectClientSecret == "" { + response.BadRequest(c, "DingTalk Client Secret is required when enabled") + return + } + req.DingTalkConnectClientSecret = previousSettings.DingTalkConnectClientSecret + } + + // Corp 策略校验(V1/V4 fail-closed) + dingTalkCfg := config.DingTalkConnectConfig{ + Enabled: true, + DingTalkAppKind: "internal_app", // 硬编码:settings 层仅支持 internal_app + AppType: "internal", // 对于 internal_only 策略的默认值 + CorpRestrictionPolicy: req.DingTalkConnectCorpRestrictionPolicy, + InternalCorpID: req.DingTalkConnectInternalCorpID, + } + // 若未填 corp_restriction_policy,保留已有配置 + if dingTalkCfg.CorpRestrictionPolicy == "" { + dingTalkCfg.CorpRestrictionPolicy = previousSettings.DingTalkConnectCorpRestrictionPolicy + } + // 对于 internal_only 策略,app_type 必须为 internal(V1 校验) + if dingTalkCfg.CorpRestrictionPolicy == "internal_only" { + dingTalkCfg.AppType = "internal" + } else { + dingTalkCfg.AppType = "public" + } + if err := config.ValidateDingTalkConfig(dingTalkCfg); err != nil { + response.ErrorWithDetails(c, http.StatusBadRequest, err.Error(), mapDingTalkValidateError(err), nil) + return + } + + // bypass_registration 仅在 internal_only 模式下有意义;其它策略下强制为 false, + // 防止 admin 在切换 policy 时把 bypass 残留在 DB 中(前端 UI 也已隐藏该开关)。 + if dingTalkCfg.CorpRestrictionPolicy != "internal_only" { + req.DingTalkConnectBypassRegistration = false + // 身份同步三开关同理:仅 internal_only 模式下有意义,其它策略强制 false。 + req.DingTalkConnectSyncCorpEmail = false + req.DingTalkConnectSyncDisplayName = false + req.DingTalkConnectSyncDept = false + } + // 身份同步目标 attr key:trimSpace + 空值 fallback 到默认值 + req.DingTalkConnectSyncCorpEmailAttrKey = strings.TrimSpace(req.DingTalkConnectSyncCorpEmailAttrKey) + if req.DingTalkConnectSyncCorpEmailAttrKey == "" { + req.DingTalkConnectSyncCorpEmailAttrKey = "dingtalk_email" + } + req.DingTalkConnectSyncDisplayNameAttrKey = strings.TrimSpace(req.DingTalkConnectSyncDisplayNameAttrKey) + if req.DingTalkConnectSyncDisplayNameAttrKey == "" { + req.DingTalkConnectSyncDisplayNameAttrKey = "dingtalk_name" + } + req.DingTalkConnectSyncDeptAttrKey = strings.TrimSpace(req.DingTalkConnectSyncDeptAttrKey) + if req.DingTalkConnectSyncDeptAttrKey == "" { + req.DingTalkConnectSyncDeptAttrKey = "dingtalk_department" + } + // 身份同步目标 attr 显示名称:trim + 空值 fallback 到默认中文名 + req.DingTalkConnectSyncCorpEmailAttrName = strings.TrimSpace(req.DingTalkConnectSyncCorpEmailAttrName) + if req.DingTalkConnectSyncCorpEmailAttrName == "" { + req.DingTalkConnectSyncCorpEmailAttrName = "钉钉企业邮箱" + } + req.DingTalkConnectSyncDisplayNameAttrName = strings.TrimSpace(req.DingTalkConnectSyncDisplayNameAttrName) + if req.DingTalkConnectSyncDisplayNameAttrName == "" { + req.DingTalkConnectSyncDisplayNameAttrName = "钉钉姓名" + } + req.DingTalkConnectSyncDeptAttrName = strings.TrimSpace(req.DingTalkConnectSyncDeptAttrName) + if req.DingTalkConnectSyncDeptAttrName == "" { + req.DingTalkConnectSyncDeptAttrName = "钉钉部门" + } + } + + if req.WeChatConnectEnabled { + req.WeChatConnectAppID = strings.TrimSpace(req.WeChatConnectAppID) + req.WeChatConnectAppSecret = strings.TrimSpace(req.WeChatConnectAppSecret) + req.WeChatConnectOpenAppID = strings.TrimSpace(req.WeChatConnectOpenAppID) + req.WeChatConnectOpenAppSecret = strings.TrimSpace(req.WeChatConnectOpenAppSecret) + req.WeChatConnectMPAppID = strings.TrimSpace(req.WeChatConnectMPAppID) + req.WeChatConnectMPAppSecret = strings.TrimSpace(req.WeChatConnectMPAppSecret) + req.WeChatConnectMobileAppID = strings.TrimSpace(req.WeChatConnectMobileAppID) + req.WeChatConnectMobileAppSecret = strings.TrimSpace(req.WeChatConnectMobileAppSecret) + req.WeChatConnectMode = strings.ToLower(strings.TrimSpace(req.WeChatConnectMode)) + req.WeChatConnectScopes = strings.TrimSpace(req.WeChatConnectScopes) + req.WeChatConnectRedirectURL = strings.TrimSpace(req.WeChatConnectRedirectURL) + req.WeChatConnectFrontendRedirectURL = strings.TrimSpace(req.WeChatConnectFrontendRedirectURL) + req.WeChatConnectAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectAppID, previousSettings.WeChatConnectAppID)) + req.WeChatConnectRedirectURL = strings.TrimSpace(firstNonEmpty(req.WeChatConnectRedirectURL, previousSettings.WeChatConnectRedirectURL)) + req.WeChatConnectFrontendRedirectURL = strings.TrimSpace(firstNonEmpty(req.WeChatConnectFrontendRedirectURL, previousSettings.WeChatConnectFrontendRedirectURL)) + if req.WeChatConnectMode == "" { + req.WeChatConnectMode = strings.ToLower(strings.TrimSpace(previousSettings.WeChatConnectMode)) + } + if req.WeChatConnectScopes == "" { + req.WeChatConnectScopes = strings.TrimSpace(previousSettings.WeChatConnectScopes) + } + + if req.WeChatConnectMPEnabled && req.WeChatConnectMobileEnabled { + response.BadRequest(c, "WeChat Official Account and Mobile App cannot be enabled at the same time") + return + } + if req.WeChatConnectMode != "" { + switch req.WeChatConnectMode { + case "open", "mp", "mobile": + default: + response.BadRequest(c, "WeChat mode must be open, mp, or mobile") + return + } + } + if !req.WeChatConnectOpenEnabled && !req.WeChatConnectMPEnabled && !req.WeChatConnectMobileEnabled { + switch req.WeChatConnectMode { + case "mp": + req.WeChatConnectMPEnabled = true + case "mobile": + req.WeChatConnectMobileEnabled = true + default: + req.WeChatConnectOpenEnabled = true + } + } + if req.WeChatConnectMode == "" { + if req.WeChatConnectMPEnabled { + req.WeChatConnectMode = "mp" + } else if req.WeChatConnectMobileEnabled { + req.WeChatConnectMode = "mobile" + } else { + req.WeChatConnectMode = "open" + } + } + + req.WeChatConnectOpenAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectOpenAppID, req.WeChatConnectAppID, previousSettings.WeChatConnectOpenAppID, previousSettings.WeChatConnectAppID)) + req.WeChatConnectMPAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectMPAppID, req.WeChatConnectAppID, previousSettings.WeChatConnectMPAppID, previousSettings.WeChatConnectAppID)) + req.WeChatConnectMobileAppID = strings.TrimSpace(firstNonEmpty(req.WeChatConnectMobileAppID, req.WeChatConnectAppID, previousSettings.WeChatConnectMobileAppID, previousSettings.WeChatConnectAppID)) + + if req.WeChatConnectOpenAppSecret == "" { + req.WeChatConnectOpenAppSecret = strings.TrimSpace(firstNonEmpty(previousSettings.WeChatConnectOpenAppSecret, previousSettings.WeChatConnectAppSecret, req.WeChatConnectAppSecret)) + } + if req.WeChatConnectMPAppSecret == "" { + req.WeChatConnectMPAppSecret = strings.TrimSpace(firstNonEmpty(previousSettings.WeChatConnectMPAppSecret, previousSettings.WeChatConnectAppSecret, req.WeChatConnectAppSecret)) + } + if req.WeChatConnectMobileAppSecret == "" { + req.WeChatConnectMobileAppSecret = strings.TrimSpace(firstNonEmpty(previousSettings.WeChatConnectMobileAppSecret, previousSettings.WeChatConnectAppSecret, req.WeChatConnectAppSecret)) + } + if req.WeChatConnectAppSecret == "" { + req.WeChatConnectAppSecret = strings.TrimSpace(firstNonEmpty(req.WeChatConnectOpenAppSecret, req.WeChatConnectMPAppSecret, req.WeChatConnectMobileAppSecret, previousSettings.WeChatConnectAppSecret)) + } + + if req.WeChatConnectOpenEnabled { + if req.WeChatConnectOpenAppID == "" { + response.BadRequest(c, "WeChat PC App ID is required when enabled") + return + } + if req.WeChatConnectOpenAppSecret == "" { + response.BadRequest(c, "WeChat PC App Secret is required when enabled") + return + } + } + if req.WeChatConnectMPEnabled { + if req.WeChatConnectMPAppID == "" { + response.BadRequest(c, "WeChat Official Account App ID is required when enabled") + return + } + if req.WeChatConnectMPAppSecret == "" { + response.BadRequest(c, "WeChat Official Account App Secret is required when enabled") + return + } + } + if req.WeChatConnectMobileEnabled { + if req.WeChatConnectMobileAppID == "" { + response.BadRequest(c, "WeChat Mobile App ID is required when enabled") + return + } + if req.WeChatConnectMobileAppSecret == "" { + response.BadRequest(c, "WeChat Mobile App Secret is required when enabled") + return + } + } + + if req.WeChatConnectScopes == "" { + if req.WeChatConnectMPEnabled { + req.WeChatConnectScopes = service.DefaultWeChatConnectScopesForMode("mp") + } else { + req.WeChatConnectScopes = service.DefaultWeChatConnectScopesForMode(req.WeChatConnectMode) + } + } + if req.WeChatConnectOpenEnabled || req.WeChatConnectMPEnabled { + if req.WeChatConnectRedirectURL == "" { + response.BadRequest(c, "WeChat Redirect URL is required when web oauth is enabled") + return + } + if err := config.ValidateAbsoluteHTTPURL(req.WeChatConnectRedirectURL); err != nil { + response.BadRequest(c, "WeChat Redirect URL must be an absolute http(s) URL") + return + } + if req.WeChatConnectFrontendRedirectURL == "" { + req.WeChatConnectFrontendRedirectURL = "/auth/wechat/callback" + } + if err := config.ValidateFrontendRedirectURL(req.WeChatConnectFrontendRedirectURL); err != nil { + response.BadRequest(c, "WeChat Frontend Redirect URL is invalid") + return + } + } + } + + // Generic OIDC 参数验证 + oidcUsePKCE, oidcValidateIDToken, err := h.settingService.OIDCSecurityWriteDefaults(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + if req.OIDCConnectEnabled { + req.OIDCConnectProviderName = strings.TrimSpace(req.OIDCConnectProviderName) + req.OIDCConnectClientID = strings.TrimSpace(req.OIDCConnectClientID) + req.OIDCConnectClientSecret = strings.TrimSpace(req.OIDCConnectClientSecret) + req.OIDCConnectIssuerURL = strings.TrimSpace(req.OIDCConnectIssuerURL) + req.OIDCConnectDiscoveryURL = strings.TrimSpace(req.OIDCConnectDiscoveryURL) + req.OIDCConnectAuthorizeURL = strings.TrimSpace(req.OIDCConnectAuthorizeURL) + req.OIDCConnectTokenURL = strings.TrimSpace(req.OIDCConnectTokenURL) + req.OIDCConnectUserInfoURL = strings.TrimSpace(req.OIDCConnectUserInfoURL) + req.OIDCConnectJWKSURL = strings.TrimSpace(req.OIDCConnectJWKSURL) + req.OIDCConnectScopes = strings.TrimSpace(req.OIDCConnectScopes) + req.OIDCConnectRedirectURL = strings.TrimSpace(req.OIDCConnectRedirectURL) + req.OIDCConnectFrontendRedirectURL = strings.TrimSpace(req.OIDCConnectFrontendRedirectURL) + req.OIDCConnectTokenAuthMethod = strings.ToLower(strings.TrimSpace(req.OIDCConnectTokenAuthMethod)) + req.OIDCConnectAllowedSigningAlgs = strings.TrimSpace(req.OIDCConnectAllowedSigningAlgs) + req.OIDCConnectUserInfoEmailPath = strings.TrimSpace(req.OIDCConnectUserInfoEmailPath) + req.OIDCConnectUserInfoIDPath = strings.TrimSpace(req.OIDCConnectUserInfoIDPath) + req.OIDCConnectUserInfoUsernamePath = strings.TrimSpace(req.OIDCConnectUserInfoUsernamePath) + req.OIDCConnectProviderName = strings.TrimSpace(firstNonEmpty(req.OIDCConnectProviderName, previousSettings.OIDCConnectProviderName, "OIDC")) + req.OIDCConnectClientID = strings.TrimSpace(firstNonEmpty(req.OIDCConnectClientID, previousSettings.OIDCConnectClientID)) + req.OIDCConnectIssuerURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectIssuerURL, previousSettings.OIDCConnectIssuerURL)) + req.OIDCConnectDiscoveryURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectDiscoveryURL, previousSettings.OIDCConnectDiscoveryURL)) + req.OIDCConnectAuthorizeURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectAuthorizeURL, previousSettings.OIDCConnectAuthorizeURL)) + req.OIDCConnectTokenURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectTokenURL, previousSettings.OIDCConnectTokenURL)) + req.OIDCConnectUserInfoURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoURL, previousSettings.OIDCConnectUserInfoURL)) + req.OIDCConnectJWKSURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectJWKSURL, previousSettings.OIDCConnectJWKSURL)) + req.OIDCConnectScopes = strings.TrimSpace(firstNonEmpty(req.OIDCConnectScopes, previousSettings.OIDCConnectScopes, "openid email profile")) + req.OIDCConnectRedirectURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectRedirectURL, previousSettings.OIDCConnectRedirectURL)) + req.OIDCConnectFrontendRedirectURL = strings.TrimSpace(firstNonEmpty(req.OIDCConnectFrontendRedirectURL, previousSettings.OIDCConnectFrontendRedirectURL, "/auth/oidc/callback")) + req.OIDCConnectTokenAuthMethod = strings.ToLower(strings.TrimSpace(firstNonEmpty(req.OIDCConnectTokenAuthMethod, previousSettings.OIDCConnectTokenAuthMethod, "client_secret_post"))) + req.OIDCConnectAllowedSigningAlgs = strings.TrimSpace(firstNonEmpty(req.OIDCConnectAllowedSigningAlgs, previousSettings.OIDCConnectAllowedSigningAlgs, "RS256,ES256,PS256")) + req.OIDCConnectUserInfoEmailPath = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoEmailPath, previousSettings.OIDCConnectUserInfoEmailPath)) + req.OIDCConnectUserInfoIDPath = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoIDPath, previousSettings.OIDCConnectUserInfoIDPath)) + req.OIDCConnectUserInfoUsernamePath = strings.TrimSpace(firstNonEmpty(req.OIDCConnectUserInfoUsernamePath, previousSettings.OIDCConnectUserInfoUsernamePath)) + if req.OIDCConnectUsePKCE != nil { + oidcUsePKCE = *req.OIDCConnectUsePKCE + } + if req.OIDCConnectValidateIDToken != nil { + oidcValidateIDToken = *req.OIDCConnectValidateIDToken + } + if req.OIDCConnectClockSkewSeconds == 0 { + req.OIDCConnectClockSkewSeconds = previousSettings.OIDCConnectClockSkewSeconds + if req.OIDCConnectClockSkewSeconds == 0 { + req.OIDCConnectClockSkewSeconds = 120 + } + } + + if req.OIDCConnectClientID == "" { + response.BadRequest(c, "OIDC Client ID is required when enabled") + return + } + if req.OIDCConnectIssuerURL == "" { + response.BadRequest(c, "OIDC Issuer URL is required when enabled") + return + } + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectIssuerURL); err != nil { + response.BadRequest(c, "OIDC Issuer URL must be an absolute http(s) URL") + return + } + if req.OIDCConnectDiscoveryURL != "" { + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectDiscoveryURL); err != nil { + response.BadRequest(c, "OIDC Discovery URL must be an absolute http(s) URL") + return + } + } + if req.OIDCConnectAuthorizeURL != "" { + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectAuthorizeURL); err != nil { + response.BadRequest(c, "OIDC Authorize URL must be an absolute http(s) URL") + return + } + } + if req.OIDCConnectTokenURL != "" { + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectTokenURL); err != nil { + response.BadRequest(c, "OIDC Token URL must be an absolute http(s) URL") + return + } + } + if req.OIDCConnectUserInfoURL != "" { + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectUserInfoURL); err != nil { + response.BadRequest(c, "OIDC UserInfo URL must be an absolute http(s) URL") + return + } + } + if req.OIDCConnectRedirectURL == "" { + response.BadRequest(c, "OIDC Redirect URL is required when enabled") + return + } + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectRedirectURL); err != nil { + response.BadRequest(c, "OIDC Redirect URL must be an absolute http(s) URL") + return + } + if req.OIDCConnectFrontendRedirectURL == "" { + response.BadRequest(c, "OIDC Frontend Redirect URL is required when enabled") + return + } + if err := config.ValidateFrontendRedirectURL(req.OIDCConnectFrontendRedirectURL); err != nil { + response.BadRequest(c, "OIDC Frontend Redirect URL is invalid") + return + } + if !scopesContainOpenID(req.OIDCConnectScopes) { + response.BadRequest(c, "OIDC scopes must contain openid") + return + } + switch req.OIDCConnectTokenAuthMethod { + case "", "client_secret_post", "client_secret_basic", "none": + default: + response.BadRequest(c, "OIDC Token Auth Method must be one of client_secret_post/client_secret_basic/none") + return + } + if req.OIDCConnectClockSkewSeconds < 0 || req.OIDCConnectClockSkewSeconds > 600 { + response.BadRequest(c, "OIDC clock skew seconds must be between 0 and 600") + return + } + if oidcValidateIDToken && req.OIDCConnectAllowedSigningAlgs == "" { + response.BadRequest(c, "OIDC Allowed Signing Algs is required when validate_id_token=true") + return + } + if req.OIDCConnectJWKSURL != "" { + if err := config.ValidateAbsoluteHTTPURL(req.OIDCConnectJWKSURL); err != nil { + response.BadRequest(c, "OIDC JWKS URL must be an absolute http(s) URL") + return + } + } + if req.OIDCConnectTokenAuthMethod == "" || req.OIDCConnectTokenAuthMethod == "client_secret_post" || req.OIDCConnectTokenAuthMethod == "client_secret_basic" { + if req.OIDCConnectClientSecret == "" { + if previousSettings.OIDCConnectClientSecret == "" { + response.BadRequest(c, "OIDC Client Secret is required when enabled") + return + } + req.OIDCConnectClientSecret = previousSettings.OIDCConnectClientSecret + } + } + } + + // “购买订阅”页面配置验证 + purchaseEnabled := previousSettings.PurchaseSubscriptionEnabled + if req.PurchaseSubscriptionEnabled != nil { + purchaseEnabled = *req.PurchaseSubscriptionEnabled + } + purchaseURL := previousSettings.PurchaseSubscriptionURL + if req.PurchaseSubscriptionURL != nil { + purchaseURL = strings.TrimSpace(*req.PurchaseSubscriptionURL) + } + + // - 启用时要求 URL 合法且非空 + // - 禁用时允许为空;若提供了 URL 也做基本校验,避免误配置 + if purchaseEnabled { + if purchaseURL == "" { + response.BadRequest(c, "Purchase Subscription URL is required when enabled") + return + } + if err := config.ValidateAbsoluteHTTPURL(purchaseURL); err != nil { + response.BadRequest(c, "Purchase Subscription URL must be an absolute http(s) URL") + return + } + } else if purchaseURL != "" { + if err := config.ValidateAbsoluteHTTPURL(purchaseURL); err != nil { + response.BadRequest(c, "Purchase Subscription URL must be an absolute http(s) URL") + return + } + } + + // Frontend URL 验证 + req.FrontendURL = strings.TrimSpace(req.FrontendURL) + if req.FrontendURL != "" { + if err := config.ValidateAbsoluteHTTPURL(req.FrontendURL); err != nil { + response.BadRequest(c, "Frontend URL must be an absolute http(s) URL") + return + } + } + + // 自定义菜单项验证 + const ( + maxCustomMenuItems = 20 + maxMenuItemLabelLen = 50 + maxMenuItemURLLen = 2048 + maxMenuItemIconSVGLen = 10 * 1024 // 10KB + maxMenuItemIDLen = 32 + ) + + customMenuJSON := previousSettings.CustomMenuItems + if req.CustomMenuItems != nil { + items := *req.CustomMenuItems + if len(items) > maxCustomMenuItems { + response.BadRequest(c, "Too many custom menu items (max 20)") + return + } + for i, item := range items { + if strings.TrimSpace(item.Label) == "" { + response.BadRequest(c, "Custom menu item label is required") + return + } + if len(item.Label) > maxMenuItemLabelLen { + response.BadRequest(c, "Custom menu item label is too long (max 50 characters)") + return + } + urlTrimmed := strings.TrimSpace(item.URL) + if strings.HasPrefix(urlTrimmed, "md:") { + // Markdown page mode: URL = "md:" + slug := strings.TrimPrefix(urlTrimmed, "md:") + if slug == "" { + response.BadRequest(c, "Custom menu item markdown slug cannot be empty (use md:slug format)") + return + } + } else { + if urlTrimmed == "" { + response.BadRequest(c, "Custom menu item URL is required (use md:slug for markdown pages)") + return + } + if len(item.URL) > maxMenuItemURLLen { + response.BadRequest(c, "Custom menu item URL is too long (max 2048 characters)") + return + } + if err := config.ValidateAbsoluteHTTPURL(urlTrimmed); err != nil { + response.BadRequest(c, "Custom menu item URL must be an absolute http(s) URL or md:") + return + } + } + if item.Visibility != "user" && item.Visibility != "admin" { + response.BadRequest(c, "Custom menu item visibility must be 'user' or 'admin'") + return + } + if len(item.IconSVG) > maxMenuItemIconSVGLen { + response.BadRequest(c, "Custom menu item icon SVG is too large (max 10KB)") + return + } + // Auto-generate ID if missing + if strings.TrimSpace(item.ID) == "" { + id, err := generateMenuItemID() + if err != nil { + response.Error(c, http.StatusInternalServerError, "Failed to generate menu item ID") + return + } + items[i].ID = id + } else if len(item.ID) > maxMenuItemIDLen { + response.BadRequest(c, "Custom menu item ID is too long (max 32 characters)") + return + } else if !menuItemIDPattern.MatchString(item.ID) { + response.BadRequest(c, "Custom menu item ID contains invalid characters (only a-z, A-Z, 0-9, - and _ are allowed)") + return + } + } + // ID uniqueness check + seen := make(map[string]struct{}, len(items)) + for _, item := range items { + if _, exists := seen[item.ID]; exists { + response.BadRequest(c, "Duplicate custom menu item ID: "+item.ID) + return + } + seen[item.ID] = struct{}{} + } + menuBytes, err := json.Marshal(items) + if err != nil { + response.BadRequest(c, "Failed to serialize custom menu items") + return + } + customMenuJSON = string(menuBytes) + } + + // 自定义端点验证 + const ( + maxCustomEndpoints = 10 + maxEndpointNameLen = 50 + maxEndpointURLLen = 2048 + maxEndpointDescriptionLen = 200 + ) + + customEndpointsJSON := previousSettings.CustomEndpoints + if req.CustomEndpoints != nil { + endpoints := *req.CustomEndpoints + if len(endpoints) > maxCustomEndpoints { + response.BadRequest(c, "Too many custom endpoints (max 10)") + return + } + for _, ep := range endpoints { + if strings.TrimSpace(ep.Name) == "" { + response.BadRequest(c, "Custom endpoint name is required") + return + } + if len(ep.Name) > maxEndpointNameLen { + response.BadRequest(c, "Custom endpoint name is too long (max 50 characters)") + return + } + if strings.TrimSpace(ep.Endpoint) == "" { + response.BadRequest(c, "Custom endpoint URL is required") + return + } + if len(ep.Endpoint) > maxEndpointURLLen { + response.BadRequest(c, "Custom endpoint URL is too long (max 2048 characters)") + return + } + if err := config.ValidateAbsoluteHTTPURL(strings.TrimSpace(ep.Endpoint)); err != nil { + response.BadRequest(c, "Custom endpoint URL must be an absolute http(s) URL") + return + } + if len(ep.Description) > maxEndpointDescriptionLen { + response.BadRequest(c, "Custom endpoint description is too long (max 200 characters)") + return + } + } + endpointBytes, err := json.Marshal(endpoints) + if err != nil { + response.BadRequest(c, "Failed to serialize custom endpoints") + return + } + customEndpointsJSON = string(endpointBytes) + } + + // Ops metrics collector interval validation (seconds). + if req.OpsMetricsIntervalSeconds != nil { + v := *req.OpsMetricsIntervalSeconds + if v < 60 { + v = 60 + } + if v > 3600 { + v = 3600 + } + req.OpsMetricsIntervalSeconds = &v + } + defaultSubscriptions := make([]service.DefaultSubscriptionSetting, 0, len(req.DefaultSubscriptions)) + for _, sub := range req.DefaultSubscriptions { + defaultSubscriptions = append(defaultSubscriptions, service.DefaultSubscriptionSetting{ + GroupID: sub.GroupID, + ValidityDays: sub.ValidityDays, + }) + } + + // 验证最低版本号格式(空字符串=禁用,或合法 semver) + if req.MinClaudeCodeVersion != "" { + if !semverPattern.MatchString(req.MinClaudeCodeVersion) { + response.Error(c, http.StatusBadRequest, "min_claude_code_version must be empty or a valid semver (e.g. 2.1.63)") + return + } + } + + // 验证最高版本号格式(空字符串=禁用,或合法 semver) + if req.MaxClaudeCodeVersion != "" { + if !semverPattern.MatchString(req.MaxClaudeCodeVersion) { + response.Error(c, http.StatusBadRequest, "max_claude_code_version must be empty or a valid semver (e.g. 3.0.0)") + return + } + } + if req.AntigravityUserAgentVersion != nil { + normalized := strings.TrimSpace(*req.AntigravityUserAgentVersion) + req.AntigravityUserAgentVersion = &normalized + if normalized != "" && !semverPattern.MatchString(normalized) { + response.Error(c, http.StatusBadRequest, "antigravity_user_agent_version must be empty or a valid semver (e.g. 1.23.2)") + return + } + } + if req.OpenAICodexUserAgent != nil { + normalized := strings.TrimSpace(*req.OpenAICodexUserAgent) + req.OpenAICodexUserAgent = &normalized + // 仅做长度上限保护,不限制具体格式(运维需要可自由调整 codex 版本号) + if len(normalized) > 512 { + response.Error(c, http.StatusBadRequest, "openai_codex_user_agent must be at most 512 characters") + return + } + } + + // codex_cli_only 加固:最低/最高 Codex 版本(空=禁用,或合法 semver;max>=min) + if req.MinCodexVersion != "" && !semverPattern.MatchString(req.MinCodexVersion) { + response.Error(c, http.StatusBadRequest, "min_codex_version must be empty or a valid semver (e.g. 0.141.0)") + return + } + if req.MaxCodexVersion != "" && !semverPattern.MatchString(req.MaxCodexVersion) { + response.Error(c, http.StatusBadRequest, "max_codex_version must be empty or a valid semver (e.g. 0.200.0)") + return + } + if req.MinCodexVersion != "" && req.MaxCodexVersion != "" && service.CompareVersions(req.MaxCodexVersion, req.MinCodexVersion) < 0 { + response.Error(c, http.StatusBadRequest, "max_codex_version must be greater than or equal to min_codex_version") + return + } + // codex_cli_only 黑/白名单:非空须为合法 []AllowedClientEntry JSON。 + // 黑名单 OR 宽 deny(允许 originator-only);白名单双因子 AND,额外要求每条可命中(非空 originator + ua_contains)。 + if err := service.ValidateCodexClientEntriesJSON(req.CodexCLIOnlyBlacklist); err != nil { + response.Error(c, http.StatusBadRequest, "codex_cli_only_blacklist "+err.Error()) + return + } + if err := service.ValidateCodexWhitelistEntriesJSON(req.CodexCLIOnlyWhitelist); err != nil { + response.Error(c, http.StatusBadRequest, "codex_cli_only_whitelist "+err.Error()) + return + } + if err := service.ValidateEngineFingerprintSignalsJSON(req.CodexCLIOnlyEngineFingerprintSignals); err != nil { + response.Error(c, http.StatusBadRequest, "codex_cli_only_engine_fingerprint_signals "+err.Error()) + return + } + + // 交叉验证:如果同时设置了最低和最高版本号,最高版本号必须 >= 最低版本号 + if req.MinClaudeCodeVersion != "" && req.MaxClaudeCodeVersion != "" { + if service.CompareVersions(req.MaxClaudeCodeVersion, req.MinClaudeCodeVersion) < 0 { + response.Error(c, http.StatusBadRequest, "max_claude_code_version must be greater than or equal to min_claude_code_version") + return + } + } + + // cyber 会话屏蔽 TTL 校验:提供时必须 > 0 + if req.CyberSessionBlockTTLSeconds != nil && *req.CyberSessionBlockTTLSeconds <= 0 { + response.BadRequest(c, "cyber_session_block_ttl_seconds must be > 0") + return + } + + settings := &service.SystemSettings{ + // 系统全局 platform quota 默认值(整体替换语义) + DefaultPlatformQuotas: req.DefaultPlatformQuotas, + + RegistrationEnabled: req.RegistrationEnabled, + EmailVerifyEnabled: req.EmailVerifyEnabled, + RegistrationEmailSuffixWhitelist: req.RegistrationEmailSuffixWhitelist, + PromoCodeEnabled: req.PromoCodeEnabled, + PasswordResetEnabled: req.PasswordResetEnabled, + FrontendURL: req.FrontendURL, + InvitationCodeEnabled: req.InvitationCodeEnabled, + TotpEnabled: req.TotpEnabled, + LoginAgreementEnabled: req.LoginAgreementEnabled, + LoginAgreementMode: loginAgreementMode, + LoginAgreementUpdatedAt: loginAgreementUpdatedAt, + LoginAgreementDocuments: loginAgreementDocuments, + SMTPHost: req.SMTPHost, + SMTPPort: req.SMTPPort, + SMTPUsername: req.SMTPUsername, + SMTPPassword: req.SMTPPassword, + SMTPFrom: req.SMTPFrom, + SMTPFromName: req.SMTPFromName, + SMTPUseTLS: req.SMTPUseTLS, + TurnstileEnabled: req.TurnstileEnabled, + TurnstileSiteKey: req.TurnstileSiteKey, + TurnstileSecretKey: req.TurnstileSecretKey, + APIKeyACLTrustForwardedIP: func() bool { + if req.APIKeyACLTrustForwardedIP != nil { + return *req.APIKeyACLTrustForwardedIP + } + return previousSettings.APIKeyACLTrustForwardedIP + }(), + LinuxDoConnectEnabled: req.LinuxDoConnectEnabled, + LinuxDoConnectClientID: req.LinuxDoConnectClientID, + LinuxDoConnectClientSecret: req.LinuxDoConnectClientSecret, + LinuxDoConnectRedirectURL: req.LinuxDoConnectRedirectURL, + DingTalkConnectEnabled: req.DingTalkConnectEnabled, + DingTalkConnectClientID: req.DingTalkConnectClientID, + DingTalkConnectClientSecret: req.DingTalkConnectClientSecret, + DingTalkConnectRedirectURL: req.DingTalkConnectRedirectURL, + DingTalkConnectCorpRestrictionPolicy: req.DingTalkConnectCorpRestrictionPolicy, + DingTalkConnectInternalCorpID: req.DingTalkConnectInternalCorpID, + DingTalkConnectBypassRegistration: req.DingTalkConnectBypassRegistration, + DingTalkConnectSyncCorpEmail: req.DingTalkConnectSyncCorpEmail, + DingTalkConnectSyncDisplayName: req.DingTalkConnectSyncDisplayName, + DingTalkConnectSyncDept: req.DingTalkConnectSyncDept, + DingTalkConnectSyncCorpEmailAttrKey: req.DingTalkConnectSyncCorpEmailAttrKey, + DingTalkConnectSyncDisplayNameAttrKey: req.DingTalkConnectSyncDisplayNameAttrKey, + DingTalkConnectSyncDeptAttrKey: req.DingTalkConnectSyncDeptAttrKey, + DingTalkConnectSyncCorpEmailAttrName: req.DingTalkConnectSyncCorpEmailAttrName, + DingTalkConnectSyncDisplayNameAttrName: req.DingTalkConnectSyncDisplayNameAttrName, + DingTalkConnectSyncDeptAttrName: req.DingTalkConnectSyncDeptAttrName, + WeChatConnectEnabled: req.WeChatConnectEnabled, + WeChatConnectAppID: req.WeChatConnectAppID, + WeChatConnectAppSecret: req.WeChatConnectAppSecret, + WeChatConnectOpenAppID: req.WeChatConnectOpenAppID, + WeChatConnectOpenAppSecret: req.WeChatConnectOpenAppSecret, + WeChatConnectMPAppID: req.WeChatConnectMPAppID, + WeChatConnectMPAppSecret: req.WeChatConnectMPAppSecret, + WeChatConnectMobileAppID: req.WeChatConnectMobileAppID, + WeChatConnectMobileAppSecret: req.WeChatConnectMobileAppSecret, + WeChatConnectOpenEnabled: req.WeChatConnectOpenEnabled, + WeChatConnectMPEnabled: req.WeChatConnectMPEnabled, + WeChatConnectMobileEnabled: req.WeChatConnectMobileEnabled, + WeChatConnectMode: req.WeChatConnectMode, + WeChatConnectScopes: req.WeChatConnectScopes, + WeChatConnectRedirectURL: req.WeChatConnectRedirectURL, + WeChatConnectFrontendRedirectURL: req.WeChatConnectFrontendRedirectURL, + OIDCConnectEnabled: req.OIDCConnectEnabled, + OIDCConnectProviderName: req.OIDCConnectProviderName, + OIDCConnectClientID: req.OIDCConnectClientID, + OIDCConnectClientSecret: req.OIDCConnectClientSecret, + OIDCConnectIssuerURL: req.OIDCConnectIssuerURL, + OIDCConnectDiscoveryURL: req.OIDCConnectDiscoveryURL, + OIDCConnectAuthorizeURL: req.OIDCConnectAuthorizeURL, + OIDCConnectTokenURL: req.OIDCConnectTokenURL, + OIDCConnectUserInfoURL: req.OIDCConnectUserInfoURL, + OIDCConnectJWKSURL: req.OIDCConnectJWKSURL, + OIDCConnectScopes: req.OIDCConnectScopes, + OIDCConnectRedirectURL: req.OIDCConnectRedirectURL, + OIDCConnectFrontendRedirectURL: req.OIDCConnectFrontendRedirectURL, + OIDCConnectTokenAuthMethod: req.OIDCConnectTokenAuthMethod, + OIDCConnectUsePKCE: oidcUsePKCE, + OIDCConnectValidateIDToken: oidcValidateIDToken, + OIDCConnectAllowedSigningAlgs: req.OIDCConnectAllowedSigningAlgs, + OIDCConnectClockSkewSeconds: req.OIDCConnectClockSkewSeconds, + OIDCConnectRequireEmailVerified: req.OIDCConnectRequireEmailVerified, + OIDCConnectUserInfoEmailPath: req.OIDCConnectUserInfoEmailPath, + OIDCConnectUserInfoIDPath: req.OIDCConnectUserInfoIDPath, + OIDCConnectUserInfoUsernamePath: req.OIDCConnectUserInfoUsernamePath, + GitHubOAuthEnabled: req.GitHubOAuthEnabled, + GitHubOAuthClientID: req.GitHubOAuthClientID, + GitHubOAuthClientSecret: req.GitHubOAuthClientSecret, + GitHubOAuthRedirectURL: req.GitHubOAuthRedirectURL, + GitHubOAuthFrontendRedirectURL: req.GitHubOAuthFrontendRedirectURL, + GoogleOAuthEnabled: req.GoogleOAuthEnabled, + GoogleOAuthClientID: req.GoogleOAuthClientID, + GoogleOAuthClientSecret: req.GoogleOAuthClientSecret, + GoogleOAuthRedirectURL: req.GoogleOAuthRedirectURL, + GoogleOAuthFrontendRedirectURL: req.GoogleOAuthFrontendRedirectURL, + SiteName: req.SiteName, + SiteLogo: req.SiteLogo, + SiteSubtitle: req.SiteSubtitle, + APIBaseURL: req.APIBaseURL, + ContactInfo: req.ContactInfo, + DocURL: req.DocURL, + HomeContent: req.HomeContent, + HideCcsImportButton: req.HideCcsImportButton, + PurchaseSubscriptionEnabled: purchaseEnabled, + PurchaseSubscriptionURL: purchaseURL, + TableDefaultPageSize: req.TableDefaultPageSize, + TablePageSizeOptions: req.TablePageSizeOptions, + CustomMenuItems: customMenuJSON, + CustomEndpoints: customEndpointsJSON, + DefaultConcurrency: req.DefaultConcurrency, + DefaultBalance: req.DefaultBalance, + AffiliateRebateRate: affiliateRebateRate, + AffiliateRebateFreezeHours: affiliateRebateFreezeHours, + AffiliateRebateDurationDays: affiliateRebateDurationDays, + AffiliateRebatePerInviteeCap: affiliateRebatePerInviteeCap, + DefaultUserRPMLimit: req.DefaultUserRPMLimit, + DefaultSubscriptions: defaultSubscriptions, + EnableModelFallback: req.EnableModelFallback, + FallbackModelAnthropic: req.FallbackModelAnthropic, + FallbackModelOpenAI: req.FallbackModelOpenAI, + FallbackModelGemini: req.FallbackModelGemini, + FallbackModelAntigravity: req.FallbackModelAntigravity, + EnableIdentityPatch: req.EnableIdentityPatch, + IdentityPatchPrompt: req.IdentityPatchPrompt, + MinClaudeCodeVersion: req.MinClaudeCodeVersion, + MaxClaudeCodeVersion: req.MaxClaudeCodeVersion, + AllowUngroupedKeyScheduling: req.AllowUngroupedKeyScheduling, + BackendModeEnabled: req.BackendModeEnabled, + AllowUserViewErrorRequests: func() bool { + if req.AllowUserViewErrorRequests != nil { + return *req.AllowUserViewErrorRequests + } + return previousSettings.AllowUserViewErrorRequests + }(), + OpsMonitoringEnabled: func() bool { + if req.OpsMonitoringEnabled != nil { + return *req.OpsMonitoringEnabled + } + return previousSettings.OpsMonitoringEnabled + }(), + OpsRealtimeMonitoringEnabled: func() bool { + if req.OpsRealtimeMonitoringEnabled != nil { + return *req.OpsRealtimeMonitoringEnabled + } + return previousSettings.OpsRealtimeMonitoringEnabled + }(), + OpsQueryModeDefault: func() string { + if req.OpsQueryModeDefault != nil { + return *req.OpsQueryModeDefault + } + return previousSettings.OpsQueryModeDefault + }(), + OpsMetricsIntervalSeconds: func() int { + if req.OpsMetricsIntervalSeconds != nil { + return *req.OpsMetricsIntervalSeconds + } + return previousSettings.OpsMetricsIntervalSeconds + }(), + EnableFingerprintUnification: func() bool { + if req.EnableFingerprintUnification != nil { + return *req.EnableFingerprintUnification + } + return previousSettings.EnableFingerprintUnification + }(), + EnableMetadataPassthrough: func() bool { + if req.EnableMetadataPassthrough != nil { + return *req.EnableMetadataPassthrough + } + return previousSettings.EnableMetadataPassthrough + }(), + EnableCCHSigning: func() bool { + if req.EnableCCHSigning != nil { + return *req.EnableCCHSigning + } + return previousSettings.EnableCCHSigning + }(), + EnableClaudeOAuthSystemPromptInjection: func() bool { + if req.EnableClaudeOAuthSystemPromptInjection != nil { + return *req.EnableClaudeOAuthSystemPromptInjection + } + return previousSettings.EnableClaudeOAuthSystemPromptInjection + }(), + ClaudeOAuthSystemPrompt: func() string { + if req.ClaudeOAuthSystemPrompt != nil { + return *req.ClaudeOAuthSystemPrompt + } + return previousSettings.ClaudeOAuthSystemPrompt + }(), + ClaudeOAuthSystemPromptBlocks: func() string { + if req.ClaudeOAuthSystemPromptBlocks != nil { + return *req.ClaudeOAuthSystemPromptBlocks + } + return previousSettings.ClaudeOAuthSystemPromptBlocks + }(), + EnableAnthropicCacheTTL1hInjection: func() bool { + if req.EnableAnthropicCacheTTL1hInjection != nil { + return *req.EnableAnthropicCacheTTL1hInjection + } + return previousSettings.EnableAnthropicCacheTTL1hInjection + }(), + RewriteMessageCacheControl: func() bool { + if req.RewriteMessageCacheControl != nil { + return *req.RewriteMessageCacheControl + } + return previousSettings.RewriteMessageCacheControl + }(), + EnableClientDatelineNormalization: func() bool { + if req.EnableClientDatelineNormalization != nil { + return *req.EnableClientDatelineNormalization + } + return previousSettings.EnableClientDatelineNormalization + }(), + AntigravityUserAgentVersion: func() string { + if req.AntigravityUserAgentVersion != nil { + return *req.AntigravityUserAgentVersion + } + return previousSettings.AntigravityUserAgentVersion + }(), + OpenAICodexUserAgent: func() string { + if req.OpenAICodexUserAgent != nil { + return *req.OpenAICodexUserAgent + } + return previousSettings.OpenAICodexUserAgent + }(), + MinCodexVersion: strings.TrimSpace(req.MinCodexVersion), + MaxCodexVersion: strings.TrimSpace(req.MaxCodexVersion), + CodexCLIOnlyBlacklist: strings.TrimSpace(req.CodexCLIOnlyBlacklist), + CodexCLIOnlyWhitelist: strings.TrimSpace(req.CodexCLIOnlyWhitelist), + CodexCLIOnlyAllowAppServerClients: func() bool { + if req.CodexCLIOnlyAllowAppServerClients != nil { + return *req.CodexCLIOnlyAllowAppServerClients + } + return previousSettings.CodexCLIOnlyAllowAppServerClients + }(), + CodexCLIOnlyEngineFingerprintSignals: strings.TrimSpace(req.CodexCLIOnlyEngineFingerprintSignals), + PaymentVisibleMethodAlipaySource: func() string { + if req.PaymentVisibleMethodAlipaySource != nil { + return strings.TrimSpace(*req.PaymentVisibleMethodAlipaySource) + } + return previousSettings.PaymentVisibleMethodAlipaySource + }(), + PaymentVisibleMethodWxpaySource: func() string { + if req.PaymentVisibleMethodWxpaySource != nil { + return strings.TrimSpace(*req.PaymentVisibleMethodWxpaySource) + } + return previousSettings.PaymentVisibleMethodWxpaySource + }(), + PaymentVisibleMethodAlipayEnabled: func() bool { + if req.PaymentVisibleMethodAlipayEnabled != nil { + return *req.PaymentVisibleMethodAlipayEnabled + } + return previousSettings.PaymentVisibleMethodAlipayEnabled + }(), + PaymentVisibleMethodWxpayEnabled: func() bool { + if req.PaymentVisibleMethodWxpayEnabled != nil { + return *req.PaymentVisibleMethodWxpayEnabled + } + return previousSettings.PaymentVisibleMethodWxpayEnabled + }(), + OpenAIAdvancedSchedulerEnabled: func() bool { + if req.OpenAIAdvancedSchedulerEnabled != nil { + return *req.OpenAIAdvancedSchedulerEnabled + } + return previousSettings.OpenAIAdvancedSchedulerEnabled + }(), + OpenAIAdvancedSchedulerStickyWeightedEnabled: func() bool { + if req.OpenAIAdvancedSchedulerStickyWeightedEnabled != nil { + return *req.OpenAIAdvancedSchedulerStickyWeightedEnabled + } + return previousSettings.OpenAIAdvancedSchedulerStickyWeightedEnabled + }(), + OpenAIAdvancedSchedulerSubscriptionPriorityEnabled: func() bool { + if req.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled != nil { + return *req.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled + } + return previousSettings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled + }(), + OpenAIAdvancedSchedulerLBTopK: stringSetting(req.OpenAIAdvancedSchedulerLBTopK, previousSettings.OpenAIAdvancedSchedulerLBTopK), + OpenAIAdvancedSchedulerWeightPriority: stringSetting(req.OpenAIAdvancedSchedulerWeightPriority, previousSettings.OpenAIAdvancedSchedulerWeightPriority), + OpenAIAdvancedSchedulerWeightLoad: stringSetting(req.OpenAIAdvancedSchedulerWeightLoad, previousSettings.OpenAIAdvancedSchedulerWeightLoad), + OpenAIAdvancedSchedulerWeightQueue: stringSetting(req.OpenAIAdvancedSchedulerWeightQueue, previousSettings.OpenAIAdvancedSchedulerWeightQueue), + OpenAIAdvancedSchedulerWeightErrorRate: stringSetting(req.OpenAIAdvancedSchedulerWeightErrorRate, previousSettings.OpenAIAdvancedSchedulerWeightErrorRate), + OpenAIAdvancedSchedulerWeightTTFT: stringSetting(req.OpenAIAdvancedSchedulerWeightTTFT, previousSettings.OpenAIAdvancedSchedulerWeightTTFT), + OpenAIAdvancedSchedulerWeightReset: stringSetting(req.OpenAIAdvancedSchedulerWeightReset, previousSettings.OpenAIAdvancedSchedulerWeightReset), + OpenAIAdvancedSchedulerWeightQuotaHeadroom: stringSetting(req.OpenAIAdvancedSchedulerWeightQuotaHeadroom, previousSettings.OpenAIAdvancedSchedulerWeightQuotaHeadroom), + OpenAIAdvancedSchedulerWeightPreviousResponse: stringSetting(req.OpenAIAdvancedSchedulerWeightPreviousResponse, previousSettings.OpenAIAdvancedSchedulerWeightPreviousResponse), + OpenAIAdvancedSchedulerWeightSessionSticky: stringSetting(req.OpenAIAdvancedSchedulerWeightSessionSticky, previousSettings.OpenAIAdvancedSchedulerWeightSessionSticky), + BalanceLowNotifyEnabled: func() bool { + if req.BalanceLowNotifyEnabled != nil { + return *req.BalanceLowNotifyEnabled + } + return previousSettings.BalanceLowNotifyEnabled + }(), + BalanceLowNotifyThreshold: func() float64 { + if req.BalanceLowNotifyThreshold != nil { + return *req.BalanceLowNotifyThreshold + } + return previousSettings.BalanceLowNotifyThreshold + }(), + BalanceLowNotifyRechargeURL: func() string { + if req.BalanceLowNotifyRechargeURL != nil { + return *req.BalanceLowNotifyRechargeURL + } + return previousSettings.BalanceLowNotifyRechargeURL + }(), + SubscriptionExpiryNotifyEnabled: func() bool { + if req.SubscriptionExpiryNotifyEnabled != nil { + return *req.SubscriptionExpiryNotifyEnabled + } + return previousSettings.SubscriptionExpiryNotifyEnabled + }(), + AccountQuotaNotifyEnabled: func() bool { + if req.AccountQuotaNotifyEnabled != nil { + return *req.AccountQuotaNotifyEnabled + } + return previousSettings.AccountQuotaNotifyEnabled + }(), + AccountQuotaNotifyEmails: func() []service.NotifyEmailEntry { + if req.AccountQuotaNotifyEmails != nil { + return dto.NotifyEmailEntriesToService(*req.AccountQuotaNotifyEmails) + } + return previousSettings.AccountQuotaNotifyEmails + }(), + ChannelMonitorEnabled: func() bool { + if req.ChannelMonitorEnabled != nil { + return *req.ChannelMonitorEnabled + } + return previousSettings.ChannelMonitorEnabled + }(), + ChannelMonitorDefaultIntervalSeconds: func() int { + if req.ChannelMonitorDefaultIntervalSeconds != nil { + return *req.ChannelMonitorDefaultIntervalSeconds + } + return previousSettings.ChannelMonitorDefaultIntervalSeconds + }(), + AvailableChannelsEnabled: func() bool { + if req.AvailableChannelsEnabled != nil { + return *req.AvailableChannelsEnabled + } + return previousSettings.AvailableChannelsEnabled + }(), + AffiliateEnabled: func() bool { + if req.AffiliateEnabled != nil { + return *req.AffiliateEnabled + } + return previousSettings.AffiliateEnabled + }(), + RiskControlEnabled: func() bool { + if req.RiskControlEnabled != nil { + return *req.RiskControlEnabled + } + return previousSettings.RiskControlEnabled + }(), + CyberSessionBlockEnabled: func() bool { + if req.CyberSessionBlockEnabled != nil { + return *req.CyberSessionBlockEnabled + } + return previousSettings.CyberSessionBlockEnabled + }(), + CyberSessionBlockTTLSeconds: func() int { + if req.CyberSessionBlockTTLSeconds != nil { + return *req.CyberSessionBlockTTLSeconds + } + return previousSettings.CyberSessionBlockTTLSeconds + }(), + } + + // req.AuthSourceXxxPlatformQuotas 为 nil 表示本次请求未包含该 source 的 quota 配置(保留 previousAuthSourceDefaults 中的值); + // non-nil(含 empty map)表示整体覆盖:empty map = 清空该 source 的所有 quota 配置。 + authSourceDefaults := &service.AuthSourceDefaultSettings{ + Email: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultEmailBalance, previousAuthSourceDefaults.Email.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultEmailConcurrency, previousAuthSourceDefaults.Email.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultEmailSubscriptions, previousAuthSourceDefaults.Email.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultEmailGrantOnSignup, previousAuthSourceDefaults.Email.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultEmailGrantOnFirstBind, previousAuthSourceDefaults.Email.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceEmailPlatformQuotas, previousAuthSourceDefaults.Email.PlatformQuotas), + }, + LinuxDo: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultLinuxDoBalance, previousAuthSourceDefaults.LinuxDo.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultLinuxDoConcurrency, previousAuthSourceDefaults.LinuxDo.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultLinuxDoSubscriptions, previousAuthSourceDefaults.LinuxDo.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultLinuxDoGrantOnSignup, previousAuthSourceDefaults.LinuxDo.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultLinuxDoGrantOnFirstBind, previousAuthSourceDefaults.LinuxDo.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceLinuxDoPlatformQuotas, previousAuthSourceDefaults.LinuxDo.PlatformQuotas), + }, + OIDC: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultOIDCBalance, previousAuthSourceDefaults.OIDC.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultOIDCConcurrency, previousAuthSourceDefaults.OIDC.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultOIDCSubscriptions, previousAuthSourceDefaults.OIDC.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultOIDCGrantOnSignup, previousAuthSourceDefaults.OIDC.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultOIDCGrantOnFirstBind, previousAuthSourceDefaults.OIDC.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceOIDCPlatformQuotas, previousAuthSourceDefaults.OIDC.PlatformQuotas), + }, + WeChat: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultWeChatBalance, previousAuthSourceDefaults.WeChat.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultWeChatConcurrency, previousAuthSourceDefaults.WeChat.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultWeChatSubscriptions, previousAuthSourceDefaults.WeChat.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultWeChatGrantOnSignup, previousAuthSourceDefaults.WeChat.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultWeChatGrantOnFirstBind, previousAuthSourceDefaults.WeChat.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceWeChatPlatformQuotas, previousAuthSourceDefaults.WeChat.PlatformQuotas), + }, + GitHub: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultGitHubBalance, previousAuthSourceDefaults.GitHub.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultGitHubConcurrency, previousAuthSourceDefaults.GitHub.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultGitHubSubscriptions, previousAuthSourceDefaults.GitHub.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultGitHubGrantOnSignup, previousAuthSourceDefaults.GitHub.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultGitHubGrantOnFirstBind, previousAuthSourceDefaults.GitHub.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceGitHubPlatformQuotas, previousAuthSourceDefaults.GitHub.PlatformQuotas), + }, + Google: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultGoogleBalance, previousAuthSourceDefaults.Google.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultGoogleConcurrency, previousAuthSourceDefaults.Google.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultGoogleSubscriptions, previousAuthSourceDefaults.Google.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultGoogleGrantOnSignup, previousAuthSourceDefaults.Google.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultGoogleGrantOnFirstBind, previousAuthSourceDefaults.Google.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceGooglePlatformQuotas, previousAuthSourceDefaults.Google.PlatformQuotas), + }, + DingTalk: service.ProviderDefaultGrantSettings{ + Balance: float64ValueOrDefault(req.AuthSourceDefaultDingTalkBalance, previousAuthSourceDefaults.DingTalk.Balance), + Concurrency: intValueOrDefault(req.AuthSourceDefaultDingTalkConcurrency, previousAuthSourceDefaults.DingTalk.Concurrency), + Subscriptions: defaultSubscriptionsValueOrDefault(req.AuthSourceDefaultDingTalkSubscriptions, previousAuthSourceDefaults.DingTalk.Subscriptions), + GrantOnSignup: boolValueOrDefault(req.AuthSourceDefaultDingTalkGrantOnSignup, previousAuthSourceDefaults.DingTalk.GrantOnSignup), + GrantOnFirstBind: boolValueOrDefault(req.AuthSourceDefaultDingTalkGrantOnFirstBind, previousAuthSourceDefaults.DingTalk.GrantOnFirstBind), + PlatformQuotas: platformQuotasValueOrDefault(req.AuthSourceDingTalkPlatformQuotas, previousAuthSourceDefaults.DingTalk.PlatformQuotas), + }, + ForceEmailOnThirdPartySignup: boolValueOrDefault(req.ForceEmailOnThirdPartySignup, previousAuthSourceDefaults.ForceEmailOnThirdPartySignup), + } + if err := h.settingService.UpdateSettingsWithAuthSourceDefaults(c.Request.Context(), settings, authSourceDefaults); err != nil { + response.ErrorFrom(c, err) + return + } + + // Update OpenAI fast policy (stored under dedicated key, only when provided). + if req.OpenAIFastPolicySettings != nil { + if err := h.settingService.SetOpenAIFastPolicySettings(c.Request.Context(), openaiFastPolicySettingsFromDTO(req.OpenAIFastPolicySettings)); err != nil { + response.BadRequest(c, err.Error()) + return + } + } + + // Update payment configuration (integrated into system settings). + // Skip if no payment fields were provided (prevents accidental wipe). + if h.paymentConfigService != nil && hasPaymentFields(req) { + paymentReq := service.UpdatePaymentConfigRequest{ + Enabled: req.PaymentEnabled, + MinAmount: req.PaymentMinAmount, + MaxAmount: req.PaymentMaxAmount, + DailyLimit: req.PaymentDailyLimit, + OrderTimeoutMin: req.PaymentOrderTimeoutMin, + MaxPendingOrders: req.PaymentMaxPendingOrders, + EnabledTypes: req.PaymentEnabledTypes, + BalanceDisabled: req.PaymentBalanceDisabled, + BalanceRechargeMultiplier: req.PaymentBalanceRechargeMultiplier, + SubscriptionUSDToCNYRate: req.PaymentSubscriptionUSDToCNYRate, + RechargeFeeRate: req.PaymentRechargeFeeRate, + LoadBalanceStrategy: req.PaymentLoadBalanceStrat, + ProductNamePrefix: req.PaymentProductNamePrefix, + ProductNameSuffix: req.PaymentProductNameSuffix, + HelpImageURL: req.PaymentHelpImageURL, + HelpText: req.PaymentHelpText, + CancelRateLimitEnabled: req.PaymentCancelRateLimitEnabled, + CancelRateLimitMax: req.PaymentCancelRateLimitMax, + CancelRateLimitWindow: req.PaymentCancelRateLimitWindow, + CancelRateLimitUnit: req.PaymentCancelRateLimitUnit, + CancelRateLimitMode: req.PaymentCancelRateLimitMode, + AlipayForceQRCode: req.PaymentAlipayForceQRCode, + } + if err := h.paymentConfigService.UpdatePaymentConfig(c.Request.Context(), paymentReq); err != nil { + response.ErrorFrom(c, err) + return + } + // Refresh in-memory provider registry so config changes take effect immediately + if h.paymentService != nil { + h.paymentService.RefreshProviders(c.Request.Context()) + } + } + + h.auditSettingsUpdate(c, previousSettings, settings, previousAuthSourceDefaults, authSourceDefaults, req) + + // 重新获取设置返回 + updatedSettings, err := h.settingService.GetAllSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + h.ensureDingTalkSyncAttributes(c.Request.Context(), updatedSettings) + updatedAuthSourceDefaults, err := h.settingService.GetAuthSourceDefaultSettings(c.Request.Context()) + if err != nil { + response.ErrorFrom(c, err) + return + } + updatedDefaultSubscriptions := make([]dto.DefaultSubscriptionSetting, 0, len(updatedSettings.DefaultSubscriptions)) + for _, sub := range updatedSettings.DefaultSubscriptions { + updatedDefaultSubscriptions = append(updatedDefaultSubscriptions, dto.DefaultSubscriptionSetting{ + GroupID: sub.GroupID, + ValidityDays: sub.ValidityDays, + }) + } + + // Reload payment config for response + var updatedPaymentCfg *service.PaymentConfig + if h.paymentConfigService != nil { + updatedPaymentCfg, _ = h.paymentConfigService.GetPaymentConfig(c.Request.Context()) + } + if updatedPaymentCfg == nil { + updatedPaymentCfg = &service.PaymentConfig{} + } + + payload := dto.SystemSettings{ + RegistrationEnabled: updatedSettings.RegistrationEnabled, + EmailVerifyEnabled: updatedSettings.EmailVerifyEnabled, + RegistrationEmailSuffixWhitelist: updatedSettings.RegistrationEmailSuffixWhitelist, + PromoCodeEnabled: updatedSettings.PromoCodeEnabled, + PasswordResetEnabled: updatedSettings.PasswordResetEnabled, + FrontendURL: updatedSettings.FrontendURL, + InvitationCodeEnabled: updatedSettings.InvitationCodeEnabled, + TotpEnabled: updatedSettings.TotpEnabled, + TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(), + LoginAgreementEnabled: updatedSettings.LoginAgreementEnabled, + LoginAgreementMode: updatedSettings.LoginAgreementMode, + LoginAgreementUpdatedAt: updatedSettings.LoginAgreementUpdatedAt, + LoginAgreementDocuments: loginAgreementDocumentsToDTO(updatedSettings.LoginAgreementDocuments), + SMTPHost: updatedSettings.SMTPHost, + SMTPPort: updatedSettings.SMTPPort, + SMTPUsername: updatedSettings.SMTPUsername, + SMTPPasswordConfigured: updatedSettings.SMTPPasswordConfigured, + SMTPFrom: updatedSettings.SMTPFrom, + SMTPFromName: updatedSettings.SMTPFromName, + SMTPUseTLS: updatedSettings.SMTPUseTLS, + TurnstileEnabled: updatedSettings.TurnstileEnabled, + TurnstileSiteKey: updatedSettings.TurnstileSiteKey, + TurnstileSecretKeyConfigured: updatedSettings.TurnstileSecretKeyConfigured, + APIKeyACLTrustForwardedIP: updatedSettings.APIKeyACLTrustForwardedIP, + LinuxDoConnectEnabled: updatedSettings.LinuxDoConnectEnabled, + LinuxDoConnectClientID: updatedSettings.LinuxDoConnectClientID, + LinuxDoConnectClientSecretConfigured: updatedSettings.LinuxDoConnectClientSecretConfigured, + LinuxDoConnectRedirectURL: updatedSettings.LinuxDoConnectRedirectURL, + DingTalkConnectEnabled: updatedSettings.DingTalkConnectEnabled, + DingTalkConnectClientID: updatedSettings.DingTalkConnectClientID, + DingTalkConnectClientSecretConfigured: updatedSettings.DingTalkConnectClientSecretConfigured, + DingTalkConnectRedirectURL: updatedSettings.DingTalkConnectRedirectURL, + DingTalkConnectCorpRestrictionPolicy: updatedSettings.DingTalkConnectCorpRestrictionPolicy, + DingTalkConnectInternalCorpID: updatedSettings.DingTalkConnectInternalCorpID, + DingTalkConnectBypassRegistration: updatedSettings.DingTalkConnectBypassRegistration, + DingTalkConnectSyncCorpEmail: updatedSettings.DingTalkConnectSyncCorpEmail, + DingTalkConnectSyncDisplayName: updatedSettings.DingTalkConnectSyncDisplayName, + DingTalkConnectSyncDept: updatedSettings.DingTalkConnectSyncDept, + DingTalkConnectSyncCorpEmailAttrKey: updatedSettings.DingTalkConnectSyncCorpEmailAttrKey, + DingTalkConnectSyncDisplayNameAttrKey: updatedSettings.DingTalkConnectSyncDisplayNameAttrKey, + DingTalkConnectSyncDeptAttrKey: updatedSettings.DingTalkConnectSyncDeptAttrKey, + DingTalkConnectSyncCorpEmailAttrName: updatedSettings.DingTalkConnectSyncCorpEmailAttrName, + DingTalkConnectSyncDisplayNameAttrName: updatedSettings.DingTalkConnectSyncDisplayNameAttrName, + DingTalkConnectSyncDeptAttrName: updatedSettings.DingTalkConnectSyncDeptAttrName, + WeChatConnectEnabled: updatedSettings.WeChatConnectEnabled, + WeChatConnectAppID: updatedSettings.WeChatConnectAppID, + WeChatConnectAppSecretConfigured: updatedSettings.WeChatConnectAppSecretConfigured, + WeChatConnectOpenAppID: updatedSettings.WeChatConnectOpenAppID, + WeChatConnectOpenAppSecretConfigured: updatedSettings.WeChatConnectOpenAppSecretConfigured, + WeChatConnectMPAppID: updatedSettings.WeChatConnectMPAppID, + WeChatConnectMPAppSecretConfigured: updatedSettings.WeChatConnectMPAppSecretConfigured, + WeChatConnectMobileAppID: updatedSettings.WeChatConnectMobileAppID, + WeChatConnectMobileAppSecretConfigured: updatedSettings.WeChatConnectMobileAppSecretConfigured, + WeChatConnectOpenEnabled: updatedSettings.WeChatConnectOpenEnabled, + WeChatConnectMPEnabled: updatedSettings.WeChatConnectMPEnabled, + WeChatConnectMobileEnabled: updatedSettings.WeChatConnectMobileEnabled, + WeChatConnectMode: updatedSettings.WeChatConnectMode, + WeChatConnectScopes: updatedSettings.WeChatConnectScopes, + WeChatConnectRedirectURL: updatedSettings.WeChatConnectRedirectURL, + WeChatConnectFrontendRedirectURL: updatedSettings.WeChatConnectFrontendRedirectURL, + OIDCConnectEnabled: updatedSettings.OIDCConnectEnabled, + OIDCConnectProviderName: updatedSettings.OIDCConnectProviderName, + OIDCConnectClientID: updatedSettings.OIDCConnectClientID, + OIDCConnectClientSecretConfigured: updatedSettings.OIDCConnectClientSecretConfigured, + OIDCConnectIssuerURL: updatedSettings.OIDCConnectIssuerURL, + OIDCConnectDiscoveryURL: updatedSettings.OIDCConnectDiscoveryURL, + OIDCConnectAuthorizeURL: updatedSettings.OIDCConnectAuthorizeURL, + OIDCConnectTokenURL: updatedSettings.OIDCConnectTokenURL, + OIDCConnectUserInfoURL: updatedSettings.OIDCConnectUserInfoURL, + OIDCConnectJWKSURL: updatedSettings.OIDCConnectJWKSURL, + OIDCConnectScopes: updatedSettings.OIDCConnectScopes, + OIDCConnectRedirectURL: updatedSettings.OIDCConnectRedirectURL, + OIDCConnectFrontendRedirectURL: updatedSettings.OIDCConnectFrontendRedirectURL, + OIDCConnectTokenAuthMethod: updatedSettings.OIDCConnectTokenAuthMethod, + OIDCConnectUsePKCE: updatedSettings.OIDCConnectUsePKCE, + OIDCConnectValidateIDToken: updatedSettings.OIDCConnectValidateIDToken, + OIDCConnectAllowedSigningAlgs: updatedSettings.OIDCConnectAllowedSigningAlgs, + OIDCConnectClockSkewSeconds: updatedSettings.OIDCConnectClockSkewSeconds, + OIDCConnectRequireEmailVerified: updatedSettings.OIDCConnectRequireEmailVerified, + OIDCConnectUserInfoEmailPath: updatedSettings.OIDCConnectUserInfoEmailPath, + OIDCConnectUserInfoIDPath: updatedSettings.OIDCConnectUserInfoIDPath, + OIDCConnectUserInfoUsernamePath: updatedSettings.OIDCConnectUserInfoUsernamePath, + GitHubOAuthEnabled: updatedSettings.GitHubOAuthEnabled, + GitHubOAuthClientID: updatedSettings.GitHubOAuthClientID, + GitHubOAuthClientSecretConfigured: updatedSettings.GitHubOAuthClientSecretConfigured, + GitHubOAuthRedirectURL: updatedSettings.GitHubOAuthRedirectURL, + GitHubOAuthFrontendRedirectURL: updatedSettings.GitHubOAuthFrontendRedirectURL, + GoogleOAuthEnabled: updatedSettings.GoogleOAuthEnabled, + GoogleOAuthClientID: updatedSettings.GoogleOAuthClientID, + GoogleOAuthClientSecretConfigured: updatedSettings.GoogleOAuthClientSecretConfigured, + GoogleOAuthRedirectURL: updatedSettings.GoogleOAuthRedirectURL, + GoogleOAuthFrontendRedirectURL: updatedSettings.GoogleOAuthFrontendRedirectURL, + SiteName: updatedSettings.SiteName, + SiteLogo: updatedSettings.SiteLogo, + SiteSubtitle: updatedSettings.SiteSubtitle, + APIBaseURL: updatedSettings.APIBaseURL, + ContactInfo: updatedSettings.ContactInfo, + DocURL: updatedSettings.DocURL, + HomeContent: updatedSettings.HomeContent, + HideCcsImportButton: updatedSettings.HideCcsImportButton, + PurchaseSubscriptionEnabled: updatedSettings.PurchaseSubscriptionEnabled, + PurchaseSubscriptionURL: updatedSettings.PurchaseSubscriptionURL, + TableDefaultPageSize: updatedSettings.TableDefaultPageSize, + TablePageSizeOptions: updatedSettings.TablePageSizeOptions, + CustomMenuItems: dto.ParseCustomMenuItems(updatedSettings.CustomMenuItems), + CustomEndpoints: dto.ParseCustomEndpoints(updatedSettings.CustomEndpoints), + DefaultConcurrency: updatedSettings.DefaultConcurrency, + DefaultBalance: updatedSettings.DefaultBalance, + AffiliateRebateRate: updatedSettings.AffiliateRebateRate, + AffiliateRebateFreezeHours: updatedSettings.AffiliateRebateFreezeHours, + AffiliateRebateDurationDays: updatedSettings.AffiliateRebateDurationDays, + AffiliateRebatePerInviteeCap: updatedSettings.AffiliateRebatePerInviteeCap, + DefaultUserRPMLimit: updatedSettings.DefaultUserRPMLimit, + DefaultSubscriptions: updatedDefaultSubscriptions, + EnableModelFallback: updatedSettings.EnableModelFallback, + FallbackModelAnthropic: updatedSettings.FallbackModelAnthropic, + FallbackModelOpenAI: updatedSettings.FallbackModelOpenAI, + FallbackModelGemini: updatedSettings.FallbackModelGemini, + FallbackModelAntigravity: updatedSettings.FallbackModelAntigravity, + EnableIdentityPatch: updatedSettings.EnableIdentityPatch, + IdentityPatchPrompt: updatedSettings.IdentityPatchPrompt, + OpsMonitoringEnabled: updatedSettings.OpsMonitoringEnabled, + OpsRealtimeMonitoringEnabled: updatedSettings.OpsRealtimeMonitoringEnabled, + OpsQueryModeDefault: updatedSettings.OpsQueryModeDefault, + OpsMetricsIntervalSeconds: updatedSettings.OpsMetricsIntervalSeconds, + MinClaudeCodeVersion: updatedSettings.MinClaudeCodeVersion, + MaxClaudeCodeVersion: updatedSettings.MaxClaudeCodeVersion, + AllowUngroupedKeyScheduling: updatedSettings.AllowUngroupedKeyScheduling, + BackendModeEnabled: updatedSettings.BackendModeEnabled, + EnableFingerprintUnification: updatedSettings.EnableFingerprintUnification, + EnableMetadataPassthrough: updatedSettings.EnableMetadataPassthrough, + EnableCCHSigning: updatedSettings.EnableCCHSigning, + EnableClaudeOAuthSystemPromptInjection: updatedSettings.EnableClaudeOAuthSystemPromptInjection, + ClaudeOAuthSystemPrompt: updatedSettings.ClaudeOAuthSystemPrompt, + ClaudeOAuthSystemPromptBlocks: updatedSettings.ClaudeOAuthSystemPromptBlocks, + EnableAnthropicCacheTTL1hInjection: updatedSettings.EnableAnthropicCacheTTL1hInjection, + RewriteMessageCacheControl: updatedSettings.RewriteMessageCacheControl, + EnableClientDatelineNormalization: updatedSettings.EnableClientDatelineNormalization, + AntigravityUserAgentVersion: updatedSettings.AntigravityUserAgentVersion, + OpenAICodexUserAgent: updatedSettings.OpenAICodexUserAgent, + MinCodexVersion: updatedSettings.MinCodexVersion, + MaxCodexVersion: updatedSettings.MaxCodexVersion, + CodexCLIOnlyBlacklist: updatedSettings.CodexCLIOnlyBlacklist, + CodexCLIOnlyWhitelist: updatedSettings.CodexCLIOnlyWhitelist, + CodexCLIOnlyAllowAppServerClients: updatedSettings.CodexCLIOnlyAllowAppServerClients, + CodexCLIOnlyEngineFingerprintSignals: updatedSettings.CodexCLIOnlyEngineFingerprintSignals, + PaymentVisibleMethodAlipaySource: updatedSettings.PaymentVisibleMethodAlipaySource, + PaymentVisibleMethodWxpaySource: updatedSettings.PaymentVisibleMethodWxpaySource, + PaymentVisibleMethodAlipayEnabled: updatedSettings.PaymentVisibleMethodAlipayEnabled, + PaymentVisibleMethodWxpayEnabled: updatedSettings.PaymentVisibleMethodWxpayEnabled, + OpenAIAdvancedSchedulerEnabled: updatedSettings.OpenAIAdvancedSchedulerEnabled, + OpenAIAdvancedSchedulerStickyWeightedEnabled: updatedSettings.OpenAIAdvancedSchedulerStickyWeightedEnabled, + OpenAIAdvancedSchedulerSubscriptionPriorityEnabled: updatedSettings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled, + OpenAIAdvancedSchedulerLBTopK: updatedSettings.OpenAIAdvancedSchedulerLBTopK, + OpenAIAdvancedSchedulerWeightPriority: updatedSettings.OpenAIAdvancedSchedulerWeightPriority, + OpenAIAdvancedSchedulerWeightLoad: updatedSettings.OpenAIAdvancedSchedulerWeightLoad, + OpenAIAdvancedSchedulerWeightQueue: updatedSettings.OpenAIAdvancedSchedulerWeightQueue, + OpenAIAdvancedSchedulerWeightErrorRate: updatedSettings.OpenAIAdvancedSchedulerWeightErrorRate, + OpenAIAdvancedSchedulerWeightTTFT: updatedSettings.OpenAIAdvancedSchedulerWeightTTFT, + OpenAIAdvancedSchedulerWeightReset: updatedSettings.OpenAIAdvancedSchedulerWeightReset, + OpenAIAdvancedSchedulerWeightQuotaHeadroom: updatedSettings.OpenAIAdvancedSchedulerWeightQuotaHeadroom, + OpenAIAdvancedSchedulerWeightPreviousResponse: updatedSettings.OpenAIAdvancedSchedulerWeightPreviousResponse, + OpenAIAdvancedSchedulerWeightSessionSticky: updatedSettings.OpenAIAdvancedSchedulerWeightSessionSticky, + OpenAIAdvancedSchedulerEffectiveLBTopK: updatedSettings.OpenAIAdvancedSchedulerEffectiveLBTopK, + OpenAIAdvancedSchedulerEffectiveWeightPriority: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightPriority, + OpenAIAdvancedSchedulerEffectiveWeightLoad: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightLoad, + OpenAIAdvancedSchedulerEffectiveWeightQueue: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightQueue, + OpenAIAdvancedSchedulerEffectiveWeightErrorRate: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightErrorRate, + OpenAIAdvancedSchedulerEffectiveWeightTTFT: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightTTFT, + OpenAIAdvancedSchedulerEffectiveWeightReset: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightReset, + OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom, + OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse, + OpenAIAdvancedSchedulerEffectiveWeightSessionSticky: updatedSettings.OpenAIAdvancedSchedulerEffectiveWeightSessionSticky, + BalanceLowNotifyEnabled: updatedSettings.BalanceLowNotifyEnabled, + BalanceLowNotifyThreshold: updatedSettings.BalanceLowNotifyThreshold, + BalanceLowNotifyRechargeURL: updatedSettings.BalanceLowNotifyRechargeURL, + SubscriptionExpiryNotifyEnabled: updatedSettings.SubscriptionExpiryNotifyEnabled, + AccountQuotaNotifyEnabled: updatedSettings.AccountQuotaNotifyEnabled, + AccountQuotaNotifyEmails: dto.NotifyEmailEntriesFromService(updatedSettings.AccountQuotaNotifyEmails), + PaymentEnabled: updatedPaymentCfg.Enabled, + PaymentMinAmount: updatedPaymentCfg.MinAmount, + PaymentMaxAmount: updatedPaymentCfg.MaxAmount, + PaymentDailyLimit: updatedPaymentCfg.DailyLimit, + PaymentOrderTimeoutMin: updatedPaymentCfg.OrderTimeoutMin, + PaymentMaxPendingOrders: updatedPaymentCfg.MaxPendingOrders, + PaymentEnabledTypes: updatedPaymentCfg.EnabledTypes, + PaymentBalanceDisabled: updatedPaymentCfg.BalanceDisabled, + PaymentBalanceRechargeMultiplier: updatedPaymentCfg.BalanceRechargeMultiplier, + PaymentSubscriptionUSDToCNYRate: updatedPaymentCfg.SubscriptionUSDToCNYRate, + PaymentRechargeFeeRate: updatedPaymentCfg.RechargeFeeRate, + PaymentLoadBalanceStrat: updatedPaymentCfg.LoadBalanceStrategy, + PaymentProductNamePrefix: updatedPaymentCfg.ProductNamePrefix, + PaymentProductNameSuffix: updatedPaymentCfg.ProductNameSuffix, + PaymentHelpImageURL: updatedPaymentCfg.HelpImageURL, + PaymentHelpText: updatedPaymentCfg.HelpText, + PaymentCancelRateLimitEnabled: updatedPaymentCfg.CancelRateLimitEnabled, + PaymentCancelRateLimitMax: updatedPaymentCfg.CancelRateLimitMax, + PaymentCancelRateLimitWindow: updatedPaymentCfg.CancelRateLimitWindow, + PaymentCancelRateLimitUnit: updatedPaymentCfg.CancelRateLimitUnit, + PaymentCancelRateLimitMode: updatedPaymentCfg.CancelRateLimitMode, + PaymentAlipayForceQRCode: updatedPaymentCfg.AlipayForceQRCode, + + ChannelMonitorEnabled: updatedSettings.ChannelMonitorEnabled, + ChannelMonitorDefaultIntervalSeconds: updatedSettings.ChannelMonitorDefaultIntervalSeconds, + + AvailableChannelsEnabled: updatedSettings.AvailableChannelsEnabled, + + AffiliateEnabled: updatedSettings.AffiliateEnabled, + + RiskControlEnabled: updatedSettings.RiskControlEnabled, + CyberSessionBlockEnabled: updatedSettings.CyberSessionBlockEnabled, + CyberSessionBlockTTLSeconds: updatedSettings.CyberSessionBlockTTLSeconds, + AllowUserViewErrorRequests: updatedSettings.AllowUserViewErrorRequests, + } + if fastPolicy, err := h.settingService.GetOpenAIFastPolicySettings(c.Request.Context()); err != nil { + slog.Error("openai_fast_policy_settings_get_failed", "error", err) + } else if fastPolicy != nil { + payload.OpenAIFastPolicySettings = openaiFastPolicySettingsToDTO(fastPolicy) + } + + // Default platform quotas(JSON map)—— 与 GetSettings 一致,避免保存后响应缺失该字段 + if platformQuotas, err := h.settingService.GetDefaultPlatformQuotas(c.Request.Context()); err != nil { + slog.Error("default_platform_quotas_get_failed", "error", err) + } else { + payload.DefaultPlatformQuotas = platformQuotas + } + response.Success(c, systemSettingsResponseData(payload, updatedAuthSourceDefaults)) +} + +// hasPaymentFields returns true if any payment-related field was explicitly provided. +// mapDingTalkValidateError maps ValidateDingTalkConfig errors to machine-readable reason codes. +func mapDingTalkValidateError(err error) string { + switch { + case errors.Is(err, config.ErrDingTalkV1AppTypeMismatch): + return "dingtalk_apptype_mismatch" + case errors.Is(err, config.ErrDingTalkV4InvalidAppKind): + return "dingtalk_app_kind_invalid" + default: + return "dingtalk_corp_config_invalid" + } +} + +func hasPaymentFields(req UpdateSettingsRequest) bool { + return req.PaymentEnabled != nil || req.PaymentMinAmount != nil || + req.PaymentMaxAmount != nil || req.PaymentDailyLimit != nil || + req.PaymentOrderTimeoutMin != nil || req.PaymentMaxPendingOrders != nil || + req.PaymentEnabledTypes != nil || req.PaymentBalanceDisabled != nil || + req.PaymentBalanceRechargeMultiplier != nil || req.PaymentSubscriptionUSDToCNYRate != nil || + req.PaymentRechargeFeeRate != nil || + req.PaymentLoadBalanceStrat != nil || req.PaymentProductNamePrefix != nil || + req.PaymentProductNameSuffix != nil || req.PaymentHelpImageURL != nil || + req.PaymentHelpText != nil || req.PaymentCancelRateLimitEnabled != nil || + req.PaymentCancelRateLimitMax != nil || req.PaymentCancelRateLimitWindow != nil || + req.PaymentCancelRateLimitUnit != nil || req.PaymentCancelRateLimitMode != nil || + req.PaymentAlipayForceQRCode != nil +} + +// ensureDingTalkSyncAttributes 在保存 settings 后,按 admin 配置的 (attr key, attr name) +// 兜底 upsert 对应 user attribute definition:不存在则创建;存在但 name 不同则更新 name +// (type/options/required 不变)。仅 internal_only + 对应 sync 开关开启时执行。 +// 失败仅记录日志,不阻塞 settings 保存。 +func (h *SettingHandler) ensureDingTalkSyncAttributes(ctx context.Context, settings *service.SystemSettings) { + if h.userAttributeService == nil || settings == nil { + return + } + if settings.DingTalkConnectCorpRestrictionPolicy != "internal_only" { + return + } + if settings.DingTalkConnectSyncDisplayName { + h.ensureUserAttributeDefinition(ctx, settings.DingTalkConnectSyncDisplayNameAttrKey, settings.DingTalkConnectSyncDisplayNameAttrName, "钉钉 internal_only 登录时同步的钉钉姓名", service.AttributeTypeText) + } + if settings.DingTalkConnectSyncCorpEmail { + h.ensureUserAttributeDefinition(ctx, settings.DingTalkConnectSyncCorpEmailAttrKey, settings.DingTalkConnectSyncCorpEmailAttrName, "钉钉 internal_only 登录时同步的企业邮箱", service.AttributeTypeEmail) + } + if settings.DingTalkConnectSyncDept { + h.ensureUserAttributeDefinition(ctx, settings.DingTalkConnectSyncDeptAttrKey, settings.DingTalkConnectSyncDeptAttrName, "钉钉 internal_only 登录时同步的完整部门路径(如:公司/研发部)", service.AttributeTypeText) + } +} + +func (h *SettingHandler) ensureUserAttributeDefinition(ctx context.Context, key, name, description string, attrType service.UserAttributeType) { + key = strings.TrimSpace(key) + if key == "" { + return + } + existing, err := h.userAttributeService.GetDefinitionByKey(ctx, key) + if err == nil && existing != nil { + if strings.TrimSpace(name) != "" && existing.Name != name { + if _, err := h.userAttributeService.UpdateDefinition(ctx, existing.ID, service.UpdateAttributeDefinitionInput{ + Name: &name, + }); err != nil { + slog.Warn("dingtalk: update user attribute definition name failed", "key", key, "err", err.Error()) + return + } + slog.Info("dingtalk: updated user attribute definition name", "key", key, "name", name) + } + return + } + if _, err := h.userAttributeService.CreateDefinition(ctx, service.CreateAttributeDefinitionInput{ + Key: key, + Name: name, + Description: description, + Type: attrType, + Enabled: true, + }); err != nil { + slog.Warn("dingtalk: ensure user attribute definition failed", "key", key, "err", err.Error()) + return + } + slog.Info("dingtalk: created user attribute definition", "key", key, "name", name, "type", attrType) +} diff --git a/backend/internal/handler/admin/system_handler.go b/backend/internal/handler/admin/system_handler.go index fb6c0ef726..49bd5bc20a 100644 --- a/backend/internal/handler/admin/system_handler.go +++ b/backend/internal/handler/admin/system_handler.go @@ -26,6 +26,8 @@ type systemUpdateService interface { CheckUpdate(ctx context.Context, force bool) (*service.UpdateInfo, error) PerformUpdate(ctx context.Context) error Rollback() error + ListRollbackVersions(ctx context.Context) ([]service.RollbackVersion, error) + RollbackToVersion(ctx context.Context, version string) error } // NewSystemHandler creates a new SystemHandler @@ -102,11 +104,42 @@ func (h *SystemHandler) PerformUpdate(c *gin.Context) { }) } -// Rollback restores the previous version +// GetRollbackVersions lists versions available for rollback +// GET /api/v1/admin/system/rollback-versions +func (h *SystemHandler) GetRollbackVersions(c *gin.Context) { + versions, err := h.updateSvc.ListRollbackVersions(c.Request.Context()) + if err != nil { + response.Error(c, http.StatusInternalServerError, err.Error()) + return + } + response.Success(c, gin.H{ + "versions": versions, + }) +} + +// Rollback restores a previous version. +// Without a body (or with an empty version) it restores the local .backup binary +// left by the last in-place update. With {"version": "x.y.z"} it downloads and +// installs that specific release (must be one of the recent rollback versions). // POST /api/v1/admin/system/rollback func (h *SystemHandler) Rollback(c *gin.Context) { - operationID := buildSystemOperationID(c, "rollback") - payload := gin.H{"operation_id": operationID} + var req struct { + Version string `json:"version"` + } + if c.Request.Body != nil && c.Request.ContentLength > 0 { + if err := c.ShouldBindJSON(&req); err != nil { + response.Error(c, http.StatusBadRequest, "invalid request body") + return + } + } + targetVersion := strings.TrimSpace(req.Version) + + operation := "rollback" + if targetVersion != "" { + operation = "rollback:" + targetVersion + } + operationID := buildSystemOperationID(c, operation) + payload := gin.H{"operation_id": operationID, "version": targetVersion} executeAdminIdempotentJSON(c, "admin.system.rollback", payload, service.DefaultSystemOperationIdempotencyTTL(), func(ctx context.Context) (any, error) { lock, release, err := h.acquireSystemLock(ctx, operationID) if err != nil { @@ -118,7 +151,12 @@ func (h *SystemHandler) Rollback(c *gin.Context) { release(releaseReason, succeeded) }() - if err := h.updateSvc.Rollback(); err != nil { + if targetVersion != "" { + err = h.updateSvc.RollbackToVersion(ctx, targetVersion) + } else { + err = h.updateSvc.Rollback() + } + if err != nil { releaseReason = "SYSTEM_ROLLBACK_FAILED" return nil, err } @@ -127,6 +165,7 @@ func (h *SystemHandler) Rollback(c *gin.Context) { return gin.H{ "message": "Rollback completed. Please restart the service.", "need_restart": true, + "version": targetVersion, "operation_id": lock.OperationID(), }, nil }) diff --git a/backend/internal/handler/admin/system_handler_test.go b/backend/internal/handler/admin/system_handler_test.go index 0f33a4528d..78e24d008e 100644 --- a/backend/internal/handler/admin/system_handler_test.go +++ b/backend/internal/handler/admin/system_handler_test.go @@ -8,6 +8,7 @@ import ( "errors" "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -17,11 +18,18 @@ import ( ) type systemHandlerUpdateServiceStub struct { - performErr error - updateInfo *service.UpdateInfo - checkErr error - checkForces []bool - performCall int + performErr error + updateInfo *service.UpdateInfo + checkErr error + checkForces []bool + performCall int + rollbackCall int + rollbackToCall int + rollbackToVersions []string + rollbackToErr error + rollbackVersions []service.RollbackVersion + rollbackVersionsErr error + rollbackVersionsCall int } func (s *systemHandlerUpdateServiceStub) CheckUpdate(_ context.Context, force bool) (*service.UpdateInfo, error) { @@ -35,9 +43,21 @@ func (s *systemHandlerUpdateServiceStub) PerformUpdate(context.Context) error { } func (s *systemHandlerUpdateServiceStub) Rollback() error { + s.rollbackCall++ return nil } +func (s *systemHandlerUpdateServiceStub) ListRollbackVersions(context.Context) ([]service.RollbackVersion, error) { + s.rollbackVersionsCall++ + return s.rollbackVersions, s.rollbackVersionsErr +} + +func (s *systemHandlerUpdateServiceStub) RollbackToVersion(_ context.Context, version string) error { + s.rollbackToCall++ + s.rollbackToVersions = append(s.rollbackToVersions, version) + return s.rollbackToErr +} + type systemUpdateResponseEnvelope struct { Code int `json:"code"` Message string `json:"message"` @@ -71,6 +91,8 @@ func newSystemHandlerTestRouter(t *testing.T, updateSvc *systemHandlerUpdateServ router := gin.New() router.POST("/api/v1/admin/system/update", handler.PerformUpdate) + router.POST("/api/v1/admin/system/rollback", handler.Rollback) + router.GET("/api/v1/admin/system/rollback-versions", handler.GetRollbackVersions) return router } @@ -142,3 +164,104 @@ func TestSystemHandlerPerformUpdateFailureStillReturnsInternalError(t *testing.T require.Equal(t, http.StatusInternalServerError, body.Code) require.Equal(t, "internal error", body.Message) } + +func TestSystemHandlerRollbackWithoutBodyUsesLegacyBackup(t *testing.T) { + updateSvc := &systemHandlerUpdateServiceStub{} + repo := newMemoryIdempotencyRepoStub() + router := newSystemHandlerTestRouter(t, updateSvc, repo) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/rollback", nil) + req.Header.Set("Idempotency-Key", "legacy-rollback") + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, 1, updateSvc.rollbackCall) + require.Equal(t, 0, updateSvc.rollbackToCall) + requireSystemLockStatus(t, repo, service.IdempotencyStatusSucceeded) +} + +func TestSystemHandlerRollbackWithVersionCallsRollbackToVersion(t *testing.T) { + updateSvc := &systemHandlerUpdateServiceStub{} + repo := newMemoryIdempotencyRepoStub() + router := newSystemHandlerTestRouter(t, updateSvc, repo) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/rollback", + strings.NewReader(`{"version":"0.1.146"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Idempotency-Key", "rollback-to-146") + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, 0, updateSvc.rollbackCall) + require.Equal(t, 1, updateSvc.rollbackToCall) + require.Equal(t, []string{"0.1.146"}, updateSvc.rollbackToVersions) + requireSystemLockStatus(t, repo, service.IdempotencyStatusSucceeded) + + var body systemUpdateResponseEnvelope + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + require.Equal(t, 0, body.Code) + require.Equal(t, "Rollback completed. Please restart the service.", body.Data.Message) +} + +func TestSystemHandlerRollbackWithDisallowedVersionReturnsBadRequest(t *testing.T) { + updateSvc := &systemHandlerUpdateServiceStub{ + rollbackToErr: service.ErrRollbackVersionNotAllowed, + } + repo := newMemoryIdempotencyRepoStub() + router := newSystemHandlerTestRouter(t, updateSvc, repo) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/rollback", + strings.NewReader(`{"version":"9.9.9"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Idempotency-Key", "rollback-to-bad") + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Equal(t, 1, updateSvc.rollbackToCall) +} + +func TestSystemHandlerGetRollbackVersions(t *testing.T) { + updateSvc := &systemHandlerUpdateServiceStub{ + rollbackVersions: []service.RollbackVersion{ + {Version: "0.1.146", PublishedAt: "2026-07-07T00:00:00Z", HTMLURL: "https://example.com/v0.1.146"}, + {Version: "0.1.145", PublishedAt: "2026-07-06T00:00:00Z", HTMLURL: "https://example.com/v0.1.145"}, + }, + } + repo := newMemoryIdempotencyRepoStub() + router := newSystemHandlerTestRouter(t, updateSvc, repo) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/system/rollback-versions", nil) + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, 1, updateSvc.rollbackVersionsCall) + + var body struct { + Code int `json:"code"` + Data struct { + Versions []service.RollbackVersion `json:"versions"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + require.Equal(t, 0, body.Code) + require.Len(t, body.Data.Versions, 2) + require.Equal(t, "0.1.146", body.Data.Versions[0].Version) +} + +func TestSystemHandlerGetRollbackVersionsError(t *testing.T) { + updateSvc := &systemHandlerUpdateServiceStub{ + rollbackVersionsErr: errors.New("github unavailable"), + } + repo := newMemoryIdempotencyRepoStub() + router := newSystemHandlerTestRouter(t, updateSvc, repo) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/system/rollback-versions", nil) + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusInternalServerError, rec.Code) +} diff --git a/backend/internal/handler/admin/user_handler.go b/backend/internal/handler/admin/user_handler.go index b47310da1d..534f774e88 100644 --- a/backend/internal/handler/admin/user_handler.go +++ b/backend/internal/handler/admin/user_handler.go @@ -53,6 +53,7 @@ type CreateUserRequest struct { Password string `json:"password" binding:"required,min=6"` Username string `json:"username"` Notes string `json:"notes"` + Role string `json:"role" binding:"omitempty,oneof=admin user"` Balance *float64 `json:"balance"` Concurrency int `json:"concurrency"` RPMLimit int `json:"rpm_limit"` @@ -66,6 +67,7 @@ type UpdateUserRequest struct { Password string `json:"password" binding:"omitempty,min=6"` Username *string `json:"username"` Notes *string `json:"notes"` + Role string `json:"role" binding:"omitempty,oneof=admin user"` Balance *float64 `json:"balance"` Concurrency *int `json:"concurrency"` RPMLimit *int `json:"rpm_limit"` @@ -269,10 +271,12 @@ func (h *UserHandler) Create(c *gin.Context) { Password: req.Password, Username: req.Username, Notes: req.Notes, + Role: req.Role, Balance: req.Balance, Concurrency: req.Concurrency, RPMLimit: req.RPMLimit, AllowedGroups: req.AllowedGroups, + ActorAdminID: getAdminIDFromContext(c), }) if err != nil { response.ErrorFrom(c, err) @@ -297,18 +301,27 @@ func (h *UserHandler) Update(c *gin.Context) { return } + // 防锁死保护:管理员不能把自己降级为普通用户(单管理员场景下会失去后台访问权)。 + // 与既有"不能禁用/删除 admin"保护一致。降级其他管理员仍然允许。 + if req.Role == service.RoleUser && userID == getAdminIDFromContext(c) { + response.BadRequest(c, "cannot demote yourself from admin") + return + } + // 使用指针类型直接传递,nil 表示未提供该字段 user, err := h.adminService.UpdateUser(c.Request.Context(), userID, &service.UpdateUserInput{ Email: req.Email, Password: req.Password, Username: req.Username, Notes: req.Notes, + Role: req.Role, Balance: req.Balance, Concurrency: req.Concurrency, RPMLimit: req.RPMLimit, Status: req.Status, AllowedGroups: req.AllowedGroups, GroupRates: req.GroupRates, + ActorAdminID: getAdminIDFromContext(c), }) if err != nil { response.ErrorFrom(c, err) diff --git a/backend/internal/handler/batch_image_handler.go b/backend/internal/handler/batch_image_handler.go new file mode 100644 index 0000000000..d739cba998 --- /dev/null +++ b/backend/internal/handler/batch_image_handler.go @@ -0,0 +1,270 @@ +package handler + +import ( + "errors" + "io" + "net/http" + "strconv" + "strings" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type BatchImageHandler struct { + service *service.BatchImagePublicService + download *service.BatchImageDownloadService + cleanup *service.BatchImageCleanupService +} + +func NewBatchImageHandler(service *service.BatchImagePublicService, download *service.BatchImageDownloadService, cleanup *service.BatchImageCleanupService) *BatchImageHandler { + return &BatchImageHandler{service: service, download: download, cleanup: cleanup} +} + +func (h *BatchImageHandler) Submit(c *gin.Context) { + var req service.BatchImageSubmitRequest + if err := c.ShouldBindJSON(&req); err != nil { + batchImageError(c, service.ErrBatchImageInvalidItems) + return + } + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + got, err := h.service.Submit(c.Request.Context(), owner, req, c.GetHeader("Idempotency-Key")) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func (h *BatchImageHandler) Get(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + got, err := h.service.Get(c.Request.Context(), owner, c.Param("id")) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func (h *BatchImageHandler) List(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + limit, _ := strconv.Atoi(c.Query("limit")) + got, err := h.service.List(c.Request.Context(), owner, service.BatchImageJobsQuery{ + Status: c.Query("status"), + TaskName: c.Query("task_name"), + Downloaded: c.Query("downloaded"), + From: c.Query("from"), + To: c.Query("to"), + Limit: limit, + Cursor: c.Query("cursor"), + }) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func (h *BatchImageHandler) Models(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + got, err := h.service.ListModels(c.Request.Context(), owner) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func (h *BatchImageHandler) Items(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + limit, _ := strconv.Atoi(c.Query("limit")) + got, err := h.service.ListItems(c.Request.Context(), owner, c.Param("id"), service.BatchImageItemsQuery{ + Status: c.Query("status"), + Limit: limit, + Cursor: c.Query("cursor"), + }) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func (h *BatchImageHandler) Cancel(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + got, err := h.service.Cancel(c.Request.Context(), owner, c.Param("id")) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func (h *BatchImageHandler) ItemContent(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + imageIndex := 0 + if raw := c.Query("image_index"); raw != "" { + parsed, err := strconv.Atoi(raw) + if err != nil { + batchImageError(c, service.ErrBatchImageItemImageIndexOutOfRange) + return + } + imageIndex = parsed + } + stream, err := h.download.OpenItemContent(c.Request.Context(), owner, c.Param("id"), c.Param("custom_id"), imageIndex) + if err != nil { + batchImageError(c, err) + return + } + defer func() { _ = stream.Reader.Close() }() + + c.Header("Content-Type", stream.ContentType) + c.Header("Content-Disposition", service.BatchImageContentDispositionAttachment(stream.Filename)) + c.Header("Cache-Control", "private, max-age=300") + c.Header("X-Content-Type-Options", "nosniff") + if stream.ContentLength != nil && *stream.ContentLength >= 0 { + c.Header("Content-Length", strconv.FormatInt(*stream.ContentLength, 10)) + } + c.Status(http.StatusOK) + if _, err := io.Copy(c.Writer, stream.Reader); err != nil { + return + } + h.markDownloadedBestEffort(c, owner) +} + +// markDownloadedBestEffort 在响应体已写出后标记下载状态; +// 此时无法再向客户端返回错误,失败只能记日志(不能静默丢弃)。 +func (h *BatchImageHandler) markDownloadedBestEffort(c *gin.Context, owner service.BatchImageOwner) { + if err := h.service.MarkDownloaded(c.Request.Context(), owner, c.Param("id")); err != nil { + logger.L().Warn("batch_image.mark_downloaded_failed", + zap.String("batch_id", c.Param("id")), + zap.Error(err), + ) + } +} + +func (h *BatchImageHandler) Download(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + maxItems, _ := strconv.Atoi(c.Query("max_items")) + + c.Header("Content-Type", "application/zip") + c.Header("Content-Disposition", service.BatchImageContentDispositionAttachment(c.Param("id")+".zip")) + c.Header("Cache-Control", "private, no-store") + c.Header("X-Content-Type-Options", "nosniff") + result, err := h.download.StreamZip(c.Request.Context(), owner, c.Param("id"), service.BatchImageZipOptions{ + Status: c.Query("status"), + MaxItems: maxItems, + IncludeManifest: true, + }, c.Writer) + if err != nil { + if result == nil || !c.Writer.Written() { + batchImageError(c, err) + } + return + } + h.markDownloadedBestEffort(c, owner) +} + +func (h *BatchImageHandler) DeleteRecord(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + if err := h.service.DeleteRecord(c.Request.Context(), owner, c.Param("id")); err != nil { + batchImageError(c, err) + return + } + c.Status(http.StatusNoContent) +} + +func (h *BatchImageHandler) DeleteOutputs(c *gin.Context) { + owner, ok := batchImageOwnerFromContext(c) + if !ok { + batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required")) + return + } + got, err := h.cleanup.DeleteOutputsForOwner(c.Request.Context(), owner, c.Param("id")) + if err != nil { + batchImageError(c, err) + return + } + c.JSON(http.StatusOK, got) +} + +func batchImageOwnerFromContext(c *gin.Context) (service.BatchImageOwner, bool) { + apiKey, ok := middleware.GetAPIKeyFromContext(c) + if !ok || apiKey == nil || apiKey.ID <= 0 || apiKey.UserID <= 0 { + return service.BatchImageOwner{}, false + } + return service.BatchImageOwner{ + UserID: apiKey.UserID, + APIKeyID: apiKey.ID, + GroupID: apiKey.GroupID, + }, true +} + +func batchImageError(c *gin.Context, err error) { + status := infraerrors.Code(err) + code := infraerrors.Reason(err) + message := infraerrors.Message(err) + if err == nil { + status = http.StatusInternalServerError + code = "INTERNAL_ERROR" + message = "internal error" + } + if status == 0 || (status == http.StatusInternalServerError && strings.TrimSpace(code) == "") { + status = http.StatusInternalServerError + code = "INTERNAL_ERROR" + message = "internal error" + } + if errors.Is(err, service.ErrBatchImageJobNotFound) { + status = http.StatusNotFound + code = "BATCH_IMAGE_NOT_FOUND" + message = "batch image job not found" + } + c.JSON(status, gin.H{ + "error": gin.H{ + "type": "invalid_request_error", + "code": code, + "message": message, + }, + }) +} diff --git a/backend/internal/handler/dto/api_key_mapper_last_used_test.go b/backend/internal/handler/dto/api_key_mapper_last_used_test.go index d63baba91a..a9ccf94524 100644 --- a/backend/internal/handler/dto/api_key_mapper_last_used_test.go +++ b/backend/internal/handler/dto/api_key_mapper_last_used_test.go @@ -10,6 +10,7 @@ import ( func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) { lastUsed := time.Now().UTC().Truncate(time.Second) + lastUsedIP := "203.0.113.10" src := &service.APIKey{ ID: 1, UserID: 2, @@ -17,6 +18,7 @@ func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) { Name: "Mapper", Status: service.StatusActive, LastUsedAt: &lastUsed, + LastUsedIP: &lastUsedIP, CurrentConcurrency: 3, } @@ -24,6 +26,8 @@ func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) { require.NotNil(t, out) require.NotNil(t, out.LastUsedAt) require.WithinDuration(t, lastUsed, *out.LastUsedAt, time.Second) + require.NotNil(t, out.LastUsedIP) + require.Equal(t, lastUsedIP, *out.LastUsedIP) require.Equal(t, 3, out.CurrentConcurrency) } @@ -39,4 +43,5 @@ func TestAPIKeyFromService_MapsNilLastUsedAt(t *testing.T) { out := APIKeyFromService(src) require.NotNil(t, out) require.Nil(t, out.LastUsedAt) + require.Nil(t, out.LastUsedIP) } diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 2b5bdccad7..6270c2b982 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -18,6 +18,7 @@ func UserFromServiceShallow(u *service.User) *User { Username: u.Username, Role: u.Role, Balance: u.Balance, + FrozenBalance: u.FrozenBalance, Concurrency: u.Concurrency, Status: u.Status, AllowedGroups: u.AllowedGroups, @@ -88,6 +89,7 @@ func APIKeyFromService(k *service.APIKey) *APIKey { IPWhitelist: k.IPWhitelist, IPBlacklist: k.IPBlacklist, LastUsedAt: k.LastUsedAt, + LastUsedIP: k.LastUsedIP, Quota: k.Quota, QuotaUsed: k.QuotaUsed, ExpiresAt: k.ExpiresAt, @@ -180,8 +182,13 @@ func groupFromServiceBase(g *service.Group) Group { WeeklyLimitUSD: g.WeeklyLimitUSD, MonthlyLimitUSD: g.MonthlyLimitUSD, AllowImageGeneration: g.AllowImageGeneration, + AllowBatchImageGeneration: g.AllowBatchImageGeneration, ImageRateIndependent: g.ImageRateIndependent, ImageRateMultiplier: g.ImageRateMultiplier, + BatchImageDiscountMultiplier: g.BatchImageDiscountMultiplier, + BatchImageHoldMultiplier: g.BatchImageHoldMultiplier, + VideoRateIndependent: g.VideoRateIndependent, + VideoRateMultiplier: g.VideoRateMultiplier, PeakRateEnabled: g.PeakRateEnabled, PeakStart: g.PeakStart, PeakEnd: g.PeakEnd, @@ -189,6 +196,9 @@ func groupFromServiceBase(g *service.Group) Group { ImagePrice1K: g.ImagePrice1K, ImagePrice2K: g.ImagePrice2K, ImagePrice4K: g.ImagePrice4K, + VideoPrice480P: g.VideoPrice480P, + VideoPrice720P: g.VideoPrice720P, + VideoPrice1080P: g.VideoPrice1080P, ClaudeCodeOnly: g.ClaudeCodeOnly, FallbackGroupID: g.FallbackGroupID, FallbackGroupIDOnInvalidRequest: g.FallbackGroupIDOnInvalidRequest, diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 8f0d23be65..7cfd102880 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -14,6 +14,7 @@ type User struct { Username string `json:"username"` Role string `json:"role"` Balance float64 `json:"balance"` + FrozenBalance float64 `json:"frozen_balance"` Concurrency int `json:"concurrency"` Status string `json:"status"` AllowedGroups []int64 `json:"allowed_groups"` @@ -58,6 +59,7 @@ type APIKey struct { IPWhitelist []string `json:"ip_whitelist"` IPBlacklist []string `json:"ip_blacklist"` LastUsedAt *time.Time `json:"last_used_at"` + LastUsedIP *string `json:"last_used_ip"` Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited) QuotaUsed float64 `json:"quota_used"` // Used quota amount in USD ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = never expires) @@ -99,9 +101,14 @@ type Group struct { MonthlyLimitUSD *float64 `json:"monthly_limit_usd"` // 图片生成计费配置(仅 antigravity 平台使用) - AllowImageGeneration bool `json:"allow_image_generation"` - ImageRateIndependent bool `json:"image_rate_independent"` - ImageRateMultiplier float64 `json:"image_rate_multiplier"` + AllowImageGeneration bool `json:"allow_image_generation"` + AllowBatchImageGeneration bool `json:"allow_batch_image_generation"` + ImageRateIndependent bool `json:"image_rate_independent"` + ImageRateMultiplier float64 `json:"image_rate_multiplier"` + BatchImageDiscountMultiplier float64 `json:"batch_image_discount_multiplier"` + BatchImageHoldMultiplier float64 `json:"batch_image_hold_multiplier"` + VideoRateIndependent bool `json:"video_rate_independent"` + VideoRateMultiplier float64 `json:"video_rate_multiplier"` // 高峰时段倍率配置 PeakRateEnabled bool `json:"peak_rate_enabled"` PeakStart string `json:"peak_start"` @@ -110,6 +117,9 @@ type Group struct { ImagePrice1K *float64 `json:"image_price_1k"` ImagePrice2K *float64 `json:"image_price_2k"` ImagePrice4K *float64 `json:"image_price_4k"` + VideoPrice480P *float64 `json:"video_price_480p"` + VideoPrice720P *float64 `json:"video_price_720p"` + VideoPrice1080P *float64 `json:"video_price_1080p"` // Claude Code 客户端限制 ClaudeCodeOnly bool `json:"claude_code_only"` diff --git a/backend/internal/handler/endpoint.go b/backend/internal/handler/endpoint.go index f8689a6e3a..0b9930c5cc 100644 --- a/backend/internal/handler/endpoint.go +++ b/backend/internal/handler/endpoint.go @@ -19,6 +19,7 @@ const ( EndpointChatCompletions = "/v1/chat/completions" EndpointEmbeddings = "/v1/embeddings" EndpointResponses = "/v1/responses" + EndpointResponsesCompact = "/v1/responses/compact" EndpointImagesGenerations = "/v1/images/generations" EndpointImagesEdits = "/v1/images/edits" EndpointVideosGenerations = "/v1/videos/generations" @@ -42,6 +43,33 @@ const ( // "/v1/chat/completions" → "/v1/chat/completions" // "/openai/v1/responses/foo" → "/v1/responses" // "/v1beta/models/gemini:gen" → "/v1beta/models" +// +// The OpenAI Responses API is also exposed via a few bare/alias +// routes that do not carry a "/v1/" prefix (top-level bare route and +// the Codex direct route). "/responses/compact" (and "/backend-api/ +// codex/responses/compact") is a distinct client endpoint — the +// "compact" client — and is normalized to its OWN canonical inbound +// endpoint, EndpointResponsesCompact, rather than being folded into +// the root Responses endpoint. Any other subpath under the bare/alias +// roots (i.e. not "compact" itself or nested under it) remains a +// subresource suffix of the root Responses endpoint: +// +// "/v1/responses/compact" → EndpointResponsesCompact +// "/v1/responses/compact/detail" → EndpointResponsesCompact +// "/openai/v1/responses/compact" → EndpointResponsesCompact +// "/openai/v1/responses/compact/detail" → EndpointResponsesCompact +// "/responses/compact" → EndpointResponsesCompact +// "/responses/compact/detail" → EndpointResponsesCompact +// "/backend-api/codex/responses/compact" → EndpointResponsesCompact +// "/backend-api/codex/responses/compact/detail" → EndpointResponsesCompact +// "/v1/responses" → EndpointResponses +// "/openai/v1/responses" → EndpointResponses +// "/responses" → EndpointResponses +// "/backend-api/codex/responses" → EndpointResponses +// +// The compact check MUST be evaluated before the root Responses check, +// otherwise "/v1/responses" (a prefix of "/v1/responses/compact") +// would erroneously match first. func NormalizeInboundEndpoint(path string) string { path = strings.TrimSpace(path) switch { @@ -59,7 +87,9 @@ func NormalizeInboundEndpoint(path string) string { return EndpointVideosGenerations case strings.Contains(path, EndpointVideos) || strings.Contains(path, "/videos/"): return EndpointVideos - case strings.Contains(path, EndpointResponses): + case strings.Contains(path, EndpointResponsesCompact) || isResponsesCompactAliasPath(path): + return EndpointResponsesCompact + case strings.Contains(path, EndpointResponses) || isResponsesRootAliasPath(path): return EndpointResponses case strings.Contains(path, EndpointGeminiModels): return EndpointGeminiModels @@ -68,6 +98,59 @@ func NormalizeInboundEndpoint(path string) string { } } +// isResponsesCompactAliasPath reports whether path is the bare/alias +// "compact" client endpoint — i.e. it is rooted at "/responses/compact" +// or "/backend-api/codex/responses/compact" (bare routes that serve +// the OpenAI Responses API "compact" client without a "/v1/" prefix), +// or any subpath nested under either of those roots: +// +// - "/responses/compact" (bare route, compact client) +// - "/responses/compact/*subpath" (nested, e.g. "/responses/compact/detail") +// - "/backend-api/codex/responses/compact" (Codex direct route, compact client) +// - "/backend-api/codex/responses/compact/*subpath" (nested, e.g. +// "/backend-api/codex/responses/compact/detail") +// +// This MUST be checked before isResponsesRootAliasPath, since +// "/responses" is a prefix of "/responses/compact". +func isResponsesCompactAliasPath(path string) bool { + trimmed := strings.TrimRight(strings.TrimSpace(path), "/") + if trimmed == "" { + return false + } + return isBareOrSubpathOf(trimmed, "/responses/compact") || isBareOrSubpathOf(trimmed, "/backend-api/codex/responses/compact") +} + +// isResponsesRootAliasPath reports whether path is one of the bare/alias +// routes that serve the root OpenAI Responses API without a "/v1/" +// prefix, or any non-"compact" subpath registered under them: +// +// - "/responses" (top-level bare route) +// - "/responses/*subpath" (any subpath other than "compact", +// since "compact" is its own distinct inbound endpoint) +// - "/backend-api/codex/responses" (Codex direct route) +// - "/backend-api/codex/responses/*subpath" (any subpath other than +// "compact") +// +// Only the top-level bare route and the Codex direct route (and their +// subpaths) are recognized here — this deliberately does NOT generalize +// to any path merely ending in "/responses" (e.g. an unrelated +// "/foo/responses" must not match). +func isResponsesRootAliasPath(path string) bool { + trimmed := strings.TrimRight(strings.TrimSpace(path), "/") + if trimmed == "" { + return false + } + return isBareOrSubpathOf(trimmed, "/responses") || isBareOrSubpathOf(trimmed, "/backend-api/codex/responses") +} + +// isBareOrSubpathOf reports whether path is exactly root, or a subpath +// rooted at root (i.e. root followed by "/"). This anchors the match +// at the start of path so it cannot match paths where root appears +// nested under some other unrelated prefix. +func isBareOrSubpathOf(path, root string) bool { + return path == root || strings.HasPrefix(path, root+"/") +} + // DeriveUpstreamEndpoint determines the upstream endpoint from the // account platform and the normalized inbound endpoint. // @@ -88,10 +171,20 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string { return inbound } // OpenAI forwards everything to the Responses API. - // Preserve subresource suffix (e.g. /v1/responses/compact). + // Preserve subresource suffix (e.g. /v1/responses/compact, + // /v1/responses/compact/detail) as derived from the raw path. if suffix := responsesSubpathSuffix(rawRequestPath); suffix != "" { return EndpointResponses + suffix } + // The raw path carried no derivable suffix (e.g. it was already + // normalized upstream, or the caller only has the canonical + // inbound endpoint available) — fall back to the canonical + // compact endpoint when that's what the inbound request was + // recognized as, so it isn't silently treated as the root + // Responses endpoint. + if inbound == EndpointResponsesCompact { + return EndpointResponsesCompact + } return EndpointResponses case service.PlatformAnthropic: @@ -142,10 +235,13 @@ func responsesSubpathSuffix(rawPath string) string { // Apply this middleware to all gateway route groups. func InboundEndpointMiddleware() gin.HandlerFunc { return func(c *gin.Context) { - path := c.FullPath() - if path == "" && c.Request != nil && c.Request.URL != nil { + path := "" + if c.Request != nil && c.Request.URL != nil { path = c.Request.URL.Path } + if path == "" { + path = c.FullPath() + } c.Set(ctxKeyInboundEndpoint, NormalizeInboundEndpoint(path)) c.Next() } @@ -158,7 +254,11 @@ func InboundEndpointMiddleware() gin.HandlerFunc { // GetInboundEndpoint returns the canonical inbound endpoint stored by // InboundEndpointMiddleware. If the middleware did not run (e.g. in -// tests), it falls back to normalizing c.FullPath() on the fly. +// tests), it falls back to normalizing c.Request.URL.Path on the fly +// (preferring the raw request path over c.FullPath(), which collapses +// wildcard route patterns such as "/v1/responses/*subpath" and would +// otherwise mis-normalize concrete requests like "/v1/responses/compact" +// to the root Responses endpoint). func GetInboundEndpoint(c *gin.Context) string { if v, ok := c.Get(ctxKeyInboundEndpoint); ok { if s, ok := v.(string); ok && s != "" { @@ -168,10 +268,12 @@ func GetInboundEndpoint(c *gin.Context) string { // Fallback: normalize on the fly. path := "" if c != nil { - path = c.FullPath() - if path == "" && c.Request != nil && c.Request.URL != nil { + if c.Request != nil && c.Request.URL != nil { path = c.Request.URL.Path } + if path == "" { + path = c.FullPath() + } } return NormalizeInboundEndpoint(path) } diff --git a/backend/internal/handler/endpoint_test.go b/backend/internal/handler/endpoint_test.go index 55e3845ed4..96ed1292b3 100644 --- a/backend/internal/handler/endpoint_test.go +++ b/backend/internal/handler/endpoint_test.go @@ -26,23 +26,42 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/v1/chat/completions", EndpointChatCompletions}, {"/v1/embeddings", EndpointEmbeddings}, {"/v1/responses", EndpointResponses}, + {"/v1/responses/compact", EndpointResponsesCompact}, + {"/v1/responses/compact/detail", EndpointResponsesCompact}, {"/v1/images/generations", EndpointImagesGenerations}, {"/v1/images/edits", EndpointImagesEdits}, {"/v1/videos/generations", EndpointVideosGenerations}, {"/v1/videos/req_123", EndpointVideos}, {"/v1beta/models", EndpointGeminiModels}, - // Prefixed paths (antigravity, openai). + // Prefixed paths (antigravity, openai) — root Responses. {"/antigravity/v1/messages", EndpointMessages}, {"/openai/v1/responses", EndpointResponses}, - {"/openai/v1/responses/compact", EndpointResponses}, {"/openai/v1/images/generations", EndpointImagesGenerations}, {"/openai/v1/images/edits", EndpointImagesEdits}, {"/antigravity/v1beta/models/gemini:generateContent", EndpointGeminiModels}, - // Gin route patterns with wildcards. - {"/v1beta/models/*modelAction", EndpointGeminiModels}, - {"/v1/responses/*subpath", EndpointResponses}, + // Prefixed paths — "/responses/compact" is its OWN distinct + // inbound endpoint, not folded into the root Responses endpoint. + {"/openai/v1/responses/compact", EndpointResponsesCompact}, + {"/openai/v1/responses/compact/detail", EndpointResponsesCompact}, + + // Bare top-level alias route "/responses" — root vs. compact. + {"/responses", EndpointResponses}, + {"/responses/compact", EndpointResponsesCompact}, + {"/responses/compact/detail", EndpointResponsesCompact}, + + // Bare Codex direct alias route — root vs. compact. + {"/backend-api/codex/responses", EndpointResponses}, + {"/backend-api/codex/responses/compact", EndpointResponsesCompact}, + {"/backend-api/codex/responses/compact/detail", EndpointResponsesCompact}, + + // Must NOT generalize to arbitrary paths merely ending in + // "/responses" (or "/responses/compact") that are unrelated to + // the two known bare alias roots, unless they already carry a + // supported "/v1/responses..." prefix form. + {"/foo/responses", "/foo/responses"}, + {"/foo/responses/compact", "/foo/responses/compact"}, // Unknown path is returned as-is. {"/v1/embeddings", "/v1/embeddings"}, @@ -74,10 +93,29 @@ func TestDeriveUpstreamEndpoint(t *testing.T) { // Gemini. {"gemini models", EndpointGeminiModels, "/v1beta/models/gemini:gen", service.PlatformGemini, EndpointGeminiModels}, - // OpenAI — always /v1/responses. + // OpenAI — root Responses. {"openai responses root", EndpointResponses, "/v1/responses", service.PlatformOpenAI, EndpointResponses}, - {"openai responses compact", EndpointResponses, "/openai/v1/responses/compact", service.PlatformOpenAI, "/v1/responses/compact"}, - {"openai responses nested", EndpointResponses, "/openai/v1/responses/compact/detail", service.PlatformOpenAI, "/v1/responses/compact/detail"}, + + // OpenAI — compact, raw path carries the derivable "/compact" + // (or nested) suffix, which must be preserved on the upstream + // endpoint. + {"openai responses compact", EndpointResponsesCompact, "/openai/v1/responses/compact", service.PlatformOpenAI, "/v1/responses/compact"}, + {"openai responses nested", EndpointResponsesCompact, "/openai/v1/responses/compact/detail", service.PlatformOpenAI, "/v1/responses/compact/detail"}, + {"openai bare responses compact", EndpointResponsesCompact, "/responses/compact", service.PlatformOpenAI, "/v1/responses/compact"}, + {"openai bare responses compact detail", EndpointResponsesCompact, "/responses/compact/detail", service.PlatformOpenAI, "/v1/responses/compact/detail"}, + {"openai codex direct responses compact", EndpointResponsesCompact, "/backend-api/codex/responses/compact", service.PlatformOpenAI, "/v1/responses/compact"}, + {"openai codex direct responses compact detail", EndpointResponsesCompact, "/backend-api/codex/responses/compact/detail", service.PlatformOpenAI, "/v1/responses/compact/detail"}, + + // OpenAI — bare root alias routes normalize to root Responses. + {"openai bare responses", EndpointResponses, "/responses", service.PlatformOpenAI, EndpointResponses}, + {"openai codex direct responses", EndpointResponses, "/backend-api/codex/responses", service.PlatformOpenAI, EndpointResponses}, + + // OpenAI — inbound is already the canonical compact endpoint but + // the raw path carries no derivable "/responses..." suffix (e.g. + // it was already normalized upstream). Must not silently fall + // back to the root Responses endpoint. + {"openai responses compact inbound only, unrelated raw path", EndpointResponsesCompact, "/v1/messages", service.PlatformOpenAI, EndpointResponsesCompact}, + {"openai from messages", EndpointMessages, "/v1/messages", service.PlatformOpenAI, EndpointResponses}, {"openai from completions", EndpointChatCompletions, "/v1/chat/completions", service.PlatformOpenAI, EndpointResponses}, {"openai embeddings", EndpointEmbeddings, "/v1/embeddings", service.PlatformOpenAI, EndpointEmbeddings}, @@ -113,6 +151,12 @@ func TestResponsesSubpathSuffix(t *testing.T) { {"/v1/responses/", ""}, {"/v1/responses/compact", "/compact"}, {"/openai/v1/responses/compact/detail", "/compact/detail"}, + {"/responses", ""}, + {"/responses/compact", "/compact"}, + {"/responses/compact/detail", "/compact/detail"}, + {"/backend-api/codex/responses", ""}, + {"/backend-api/codex/responses/compact", "/compact"}, + {"/backend-api/codex/responses/compact/detail", "/compact/detail"}, {"/v1/messages", ""}, {"", ""}, } @@ -154,6 +198,132 @@ func TestGetInboundEndpoint_FallbackWithoutMiddleware(t *testing.T) { require.Equal(t, EndpointMessages, got) } +// TestInboundEndpointMiddleware_WildcardRoutes verifies that, when a +// gateway route is registered with a Gin wildcard pattern (e.g. +// "/v1/responses/*subpath"), InboundEndpointMiddleware normalizes based +// on the concrete request path (c.Request.URL.Path) rather than the +// route pattern (c.FullPath()). Using c.FullPath() here would collapse +// every request under the wildcard — including "/v1/responses/compact" +// — down to the literal pattern string, which never matches the +// "compact" alias detection and would incorrectly normalize to the root +// Responses endpoint. +func TestInboundEndpointMiddleware_WildcardRoutes(t *testing.T) { + tests := []struct { + name string + routePath string + requestPath string + want string + }{ + { + name: "v1 responses wildcard route, compact request", + routePath: "/v1/responses/*subpath", + requestPath: "/v1/responses/compact", + want: EndpointResponsesCompact, + }, + { + name: "bare responses wildcard route, compact request", + routePath: "/responses/*subpath", + requestPath: "/responses/compact", + want: EndpointResponsesCompact, + }, + { + name: "codex direct wildcard route, compact request", + routePath: "/backend-api/codex/responses/*subpath", + requestPath: "/backend-api/codex/responses/compact", + want: EndpointResponsesCompact, + }, + { + name: "v1 responses wildcard route, non-compact subpath request", + routePath: "/v1/responses/*subpath", + requestPath: "/v1/responses/foo", + want: EndpointResponses, + }, + { + name: "bare responses wildcard route, non-compact subpath request", + routePath: "/responses/*subpath", + requestPath: "/responses/foo", + want: EndpointResponses, + }, + { + name: "codex direct wildcard route, non-compact subpath request", + routePath: "/backend-api/codex/responses/*subpath", + requestPath: "/backend-api/codex/responses/foo", + want: EndpointResponses, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + router := gin.New() + router.Use(InboundEndpointMiddleware()) + + var captured string + router.POST(tt.routePath, func(c *gin.Context) { + captured = GetInboundEndpoint(c) + c.Status(http.StatusOK) + }) + + req := httptest.NewRequest(http.MethodPost, tt.requestPath, nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, tt.want, captured) + }) + } +} + +// TestInboundEndpointMiddleware_GeminiWildcardRoute verifies that a Gemini +// wildcard route (e.g. "/v1beta/models/*modelAction", used to capture the +// ":generateContent"-style action suffix embedded in the path) is normalized +// to EndpointGeminiModels via InboundEndpointMiddleware, using the same real +// Gin routing path as TestInboundEndpointMiddleware_WildcardRoutes above. +func TestInboundEndpointMiddleware_GeminiWildcardRoute(t *testing.T) { + router := gin.New() + router.Use(InboundEndpointMiddleware()) + + var captured string + router.POST("/v1beta/models/*modelAction", func(c *gin.Context) { + captured = GetInboundEndpoint(c) + c.Status(http.StatusOK) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, EndpointGeminiModels, captured) +} + +// TestGetInboundEndpoint_FallbackWildcardRouteWithoutMiddleware verifies +// that when InboundEndpointMiddleware did NOT run (so no value is stored +// in gin.Context), the GetInboundEndpoint fallback path still prefers +// c.Request.URL.Path over c.FullPath(). This guards against the fallback +// regressing to prefer c.FullPath() again, which would misnormalize +// concrete requests matched by a wildcard route pattern (e.g. +// "/v1/responses/*subpath" matching "/v1/responses/compact") down to +// the root Responses endpoint. +func TestGetInboundEndpoint_FallbackWildcardRouteWithoutMiddleware(t *testing.T) { + router := gin.New() + // Deliberately do NOT register InboundEndpointMiddleware. + + var captured string + router.POST("/v1/responses/*subpath", func(c *gin.Context) { + // Sanity check: FullPath returns the route pattern, not the + // concrete request path, when a wildcard route matches. + require.Equal(t, "/v1/responses/*subpath", c.FullPath()) + captured = GetInboundEndpoint(c) + c.Status(http.StatusOK) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, EndpointResponsesCompact, captured) +} + func TestGetUpstreamEndpoint_FullFlow(t *testing.T) { rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index b20d9ef652..116346b4a6 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -20,7 +20,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" - pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/pkg/openai" @@ -138,7 +137,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { defer h.maybeLogCompatibilityFallbackMetrics(reqLog) // 读取请求体 - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -158,6 +157,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { bodyRef := service.NewRequestBodyRef(body) parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { + logRequestBodyParseFailure(reqLog, body, err) h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } @@ -1628,6 +1628,11 @@ func (h *GatewayHandler) mapUpstreamError(statusCode int) (int, string, string) // handleStreamingAwareError handles errors that may occur after streaming has started func (h *GatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) { if streamStarted { + // 响应状态码已固化为 200(ping/部分数据已 flush),错误只能就地以 SSE 帧回传。 + // 标记本次流内错误,供 ops_error_logger 补记——否则该中间件按 status>=400 采集, + // 这类挂在 200 流上的失败(如并发限流回退)不会进错误看板。 + service.MarkOpsStreamError(c, errType, message, status) + // /v1/responses 的严格 SDK(Codex CLI)要求终止事件必须属于 // response.completed/failed/incomplete/cancelled 集合。 // Anthropic-backed Responses 路径同样会因为通用 error 帧被拒。 @@ -1776,7 +1781,7 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) { defer h.maybeLogCompatibilityFallbackMetrics(reqLog) // 读取请求体 - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -1796,6 +1801,7 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) { bodyRef := service.NewRequestBodyRef(body) parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { + logRequestBodyParseFailure(reqLog, body, err) h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index d0ecc01e6a..f3805f3a53 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -7,7 +7,6 @@ import ( "strconv" "time" - pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" @@ -45,7 +44,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { ) // Read request body - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.chatCompletionsErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -64,6 +63,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { // Validate JSON if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 4a8d752193..5b49ca69a2 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -7,7 +7,6 @@ import ( "strconv" "time" - pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" @@ -45,7 +44,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { ) // Read request body - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.responsesErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -64,6 +63,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { // Validate JSON if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.responsesErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/handler.go b/backend/internal/handler/handler.go index 014cf7d2ba..58c524889c 100644 --- a/backend/internal/handler/handler.go +++ b/backend/internal/handler/handler.go @@ -58,6 +58,7 @@ type Handlers struct { Payment *PaymentHandler PaymentWebhook *PaymentWebhookHandler AvailableChannel *AvailableChannelHandler + BatchImage *BatchImageHandler } // BuildInfo contains build-time information diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index baff1dcbd6..f5f2522e49 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -7,7 +7,6 @@ import ( "strconv" "time" - pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" @@ -49,7 +48,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { return } - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -64,6 +63,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { } if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/openai_codex_models_handler.go b/backend/internal/handler/openai_codex_models_handler.go new file mode 100644 index 0000000000..e64c555d14 --- /dev/null +++ b/backend/internal/handler/openai_codex_models_handler.go @@ -0,0 +1,53 @@ +package handler + +import ( + "net/http" + + "github.com/gin-gonic/gin" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" +) + +// CodexModels serves the Codex models manifest for Codex clients. +// +// Codex CLI and the Codex desktop app refresh their model picker from +// GET {base_url}/models?client_version=... (custom provider mode) or +// GET /backend-api/codex/models (chatgpt_base_url mode). Both routes land +// here. The manifest is proxied verbatim from the ChatGPT backend with a +// schedulable OAuth account's credentials, so clients pointed at the gateway +// see the account's real, always-current model entitlements instead of a +// frozen local cache. +func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) { + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok || apiKey.Group == nil { + h.errorResponse(c, http.StatusUnauthorized, "invalid_request_error", "API key group is required") + return + } + if apiKey.Group.Platform != service.PlatformOpenAI { + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Codex models manifest is only available for OpenAI groups") + return + } + + account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, "", "") + if err != nil { + h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts") + return + } + + manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match")) + if err != nil { + h.errorResponse(c, infraerrors.Code(err), "upstream_error", infraerrors.Message(err)) + return + } + + if manifest.ETag != "" { + c.Header("ETag", manifest.ETag) + } + if manifest.NotModified { + c.Status(http.StatusNotModified) + return + } + c.Data(http.StatusOK, "application/json", manifest.Body) +} diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index 8be533c723..56d775eb7c 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -60,6 +60,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { return } if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/openai_gateway_compact_body_signal_test.go b/backend/internal/handler/openai_gateway_compact_body_signal_test.go new file mode 100644 index 0000000000..a4bfb90466 --- /dev/null +++ b/backend/internal/handler/openai_gateway_compact_body_signal_test.go @@ -0,0 +1,162 @@ +package handler + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" + "go.uber.org/zap" +) + +func newCompactBodySignalTestContext(t *testing.T, path string, body []byte) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + return c +} + +// body-signal 提升后必须与 path-based compact 走同一条链路: +// path 改写、requireCompact 判定、stream/store/prompt_cache_key 归一化删除。 +// 回归防护:若 stream 字段存活,Forward 会用流式 handler 解析 compact 的 +// JSON 响应,导致 "stream ended before a terminal event" 的换号 failover 风暴。 +func TestNormalizeOpenAIResponsesCompactRequest_BodySignalPromoted(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{ + "model":"gpt-5.5", + "stream":true, + "store":true, + "prompt_cache_key":"pck-signal-1", + "input":[ + {"type":"message","role":"user","content":"hello"}, + {"type":"compaction_trigger"} + ] + }`) + c := newCompactBodySignalTestContext(t, "/v1/responses", body) + + normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + + require.Equal(t, "/v1/responses/compact", c.Request.URL.Path) + require.True(t, isOpenAIRemoteCompactPath(c)) + + require.False(t, gjson.GetBytes(normalized, "stream").Exists()) + require.False(t, gjson.GetBytes(normalized, "store").Exists()) + require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists()) + require.Equal(t, "gpt-5.5", gjson.GetBytes(normalized, "model").String()) + require.True(t, gjson.GetBytes(normalized, "input").IsArray()) + + reqStream, streamOK := parseOpenAICompatibleStream(normalized) + require.True(t, streamOK) + require.False(t, reqStream) + + seed, exists := c.Get(service.OpenAICompactSessionSeedKeyForTest()) + require.True(t, exists) + require.Equal(t, "pck-signal-1", seed) +} + +func TestNormalizeOpenAIResponsesCompactRequest_BodySignalTrailingSlash(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`) + c := newCompactBodySignalTestContext(t, "/v1/responses/", body) + + _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + require.Equal(t, "/v1/responses/compact", c.Request.URL.Path) +} + +func TestNormalizeOpenAIResponsesCompactRequest_CodexDirectAliasPromoted(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`) + c := newCompactBodySignalTestContext(t, "/backend-api/codex/responses", body) + + _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + require.Equal(t, "/backend-api/codex/responses/compact", c.Request.URL.Path) +} + +func TestNormalizeOpenAIResponsesCompactRequest_NoTriggerUntouched(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"message","role":"user","content":"hello"}]}`) + c := newCompactBodySignalTestContext(t, "/v1/responses", body) + + normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + require.Equal(t, "/v1/responses", c.Request.URL.Path) + require.False(t, isOpenAIRemoteCompactPath(c)) + require.Equal(t, body, normalized) + require.True(t, gjson.GetBytes(normalized, "stream").Bool()) +} + +func TestNormalizeOpenAIResponsesCompactRequest_PathBasedNoDoubleSuffix(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","stream":true,"store":true,"input":[{"type":"message","role":"user","content":"hello"}]}`) + c := newCompactBodySignalTestContext(t, "/v1/responses/compact", body) + + normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + require.Equal(t, "/v1/responses/compact", c.Request.URL.Path) + require.False(t, gjson.GetBytes(normalized, "stream").Exists()) + require.False(t, gjson.GetBytes(normalized, "store").Exists()) +} + +func TestNormalizeOpenAIResponsesCompactRequest_SubpathNotPromoted(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`) + c := newCompactBodySignalTestContext(t, "/v1/responses/resp_123/cancel", body) + + normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + require.Equal(t, "/v1/responses/resp_123/cancel", c.Request.URL.Path) + require.Equal(t, body, normalized) +} + +// 回归 #3875:body-signal 原始请求 stream:true 时必须标记 client-stream, +// 供响应写回阶段把上游 unary JSON 合成回 Codex remote compact v2 所需的 SSE。 +func TestNormalizeOpenAIResponsesCompactRequest_BodySignalStreamTrueMarksClientStream(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`) + c := newCompactBodySignalTestContext(t, "/v1/responses", body) + + _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + + marked, exists := c.Get(service.OpenAICompactClientStreamKeyForTest()) + require.True(t, exists) + require.Equal(t, true, marked) +} + +func TestNormalizeOpenAIResponsesCompactRequest_BodySignalStreamFalseNotMarked(t *testing.T) { + h := &OpenAIGatewayHandler{} + for name, body := range map[string][]byte{ + "stream_false": []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`), + "stream_absent": []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`), + } { + c := newCompactBodySignalTestContext(t, "/v1/responses", body) + _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok, name) + require.Equal(t, "/v1/responses/compact", c.Request.URL.Path, name) + _, exists := c.Get(service.OpenAICompactClientStreamKeyForTest()) + require.False(t, exists, "case %s 不应标记 client-stream", name) + } +} + +// path-based compact(Codex v1 unary 协议)即使 body 带 stream:true 也不标记, +// 保持 JSON 写回行为不变。 +func TestNormalizeOpenAIResponsesCompactRequest_PathBasedStreamTrueNotMarked(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"message","role":"user","content":"hello"}]}`) + c := newCompactBodySignalTestContext(t, "/v1/responses/compact", body) + + _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + _, exists := c.Get(service.OpenAICompactClientStreamKeyForTest()) + require.False(t, exists) +} diff --git a/backend/internal/handler/openai_gateway_count_tokens.go b/backend/internal/handler/openai_gateway_count_tokens.go index fc9c4d5df7..0461017067 100644 --- a/backend/internal/handler/openai_gateway_count_tokens.go +++ b/backend/internal/handler/openai_gateway_count_tokens.go @@ -6,7 +6,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/domain" - pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" @@ -47,7 +46,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) { return } - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.anthropicErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -64,6 +63,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) { bodyRef := service.NewRequestBodyRef(body) parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { + logRequestBodyParseFailure(reqLog, body, err) h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 7f097afa4b..feba261103 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -13,7 +13,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" - pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" @@ -185,7 +184,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { } // Read request body - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -202,22 +201,14 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { setOpsRequestContext(c, "", false) sessionHashBody := body - if service.IsOpenAIResponsesCompactPathForTest(c) { - if compactSeed := strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()); compactSeed != "" { - c.Set(service.OpenAICompactSessionSeedKeyForTest(), compactSeed) - } - normalizedCompactBody, normalizedCompact, compactErr := service.NormalizeOpenAICompactRequestBodyForTest(body) - if compactErr != nil { - h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to normalize compact request body") - return - } - if normalizedCompact { - body = normalizedCompactBody - } + body, ok = h.normalizeOpenAIResponsesCompactRequest(c, reqLog, body) + if !ok { + return } // 校验请求体 JSON 合法性 if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } @@ -572,6 +563,54 @@ func isOpenAIRemoteCompactPath(c *gin.Context) bool { return strings.HasSuffix(normalizedPath, "/responses/compact") } +// isBareOpenAIResponsesPath 仅匹配裸 /responses 端点(无 /compact 等子路径), +// body-signal 提升只允许发生在这里,避免误伤 /responses/{id}/... 形态的请求。 +func isBareOpenAIResponsesPath(c *gin.Context) bool { + if c == nil || c.Request == nil || c.Request.URL == nil { + return false + } + normalizedPath := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/") + return strings.HasSuffix(normalizedPath, "/responses") +} + +// normalizeOpenAIResponsesCompactRequest 统一处理两种入站 compact 形态: +// path-based(POST /v1/responses/compact)与 Codex remote compact v2 的 +// body-signal(普通 POST /v1/responses 的 input 中携带 type=compaction_trigger, +// 见 #3777)。body-signal 命中时在 stream 解析、compact body 归一化与 +// requireCompact 调度判定之前改写 URL path,使后续全部链路(含 passthrough +// 分支与上游 URL 构建)与 path-based 完全一致。 +// 返回归一化后的 body;ok=false 表示错误响应已写出,调用方应直接 return。 +func (h *OpenAIGatewayHandler) normalizeOpenAIResponsesCompactRequest(c *gin.Context, reqLog *zap.Logger, body []byte) ([]byte, bool) { + isCompactRequest := service.IsOpenAIResponsesCompactPathForTest(c) + if !isCompactRequest && isBareOpenAIResponsesPath(c) && service.HasCompactionTriggerInInput(body) { + c.Request.URL.Path = strings.TrimRight(c.Request.URL.Path, "/") + "/compact" + isCompactRequest = true + // Codex remote compact v2 的原始请求是流式 /responses:白名单归一化会删除 + // stream 并让上游走 unary JSON,但客户端仍按 SSE 消费响应。记录原始 + // stream 意图,响应写回阶段据此把 JSON 合成回 SSE(#3875)。 + clientStream := gjson.GetBytes(body, "stream").Bool() + if clientStream { + service.MarkOpenAICompactClientStream(c) + } + reqLog.Info("codex.remote_compact.detected_body_signal", zap.Bool("client_stream", clientStream)) + } + if !isCompactRequest { + return body, true + } + if compactSeed := strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()); compactSeed != "" { + c.Set(service.OpenAICompactSessionSeedKeyForTest(), compactSeed) + } + normalizedCompactBody, normalizedCompact, compactErr := service.NormalizeOpenAICompactRequestBodyForTest(body) + if compactErr != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to normalize compact request body") + return nil, false + } + if normalizedCompact { + body = normalizedCompactBody + } + return body, true +} + func (h *OpenAIGatewayHandler) logOpenAIRemoteCompactOutcome(c *gin.Context, startedAt time.Time) { if !isOpenAIRemoteCompactPath(c) { return @@ -682,7 +721,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } - body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) if err != nil { if maxErr, ok := extractMaxBytesError(err); ok { h.anthropicErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) @@ -697,6 +736,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { } if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index c4ccb5a024..b7f43079ef 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -440,7 +440,7 @@ func TestResolveOpenAIMessagesDispatchMappedModel(t *testing.T) { Platform: service.PlatformGrok, }, } - require.Equal(t, "grok-4.3", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5")) + require.Equal(t, "grok-4.5", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5")) require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(apiKey, "grok")) }) diff --git a/backend/internal/handler/ops_error_logger.go b/backend/internal/handler/ops_error_logger.go index 8cd9d02ba0..5a1e57ff7d 100644 --- a/backend/internal/handler/ops_error_logger.go +++ b/backend/internal/handler/ops_error_logger.go @@ -590,6 +590,10 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } if !hasUpstreamContext { + // 没有上游错误上下文,但网关可能在已固化的 200 流上就地补发了 SSE 错误帧 + // (如 ping 等待后并发超限、Wait 后二次计费校验失败)。这类失败若不在此补记, + // 会因 wire 状态码为 200 而在错误看板里彻底隐形。 + logOpsStreamError(c, ops, status) return } @@ -999,6 +1003,138 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } +// logOpsStreamError 记录一次挂在已固化 HTTP 200 SSE 流上的就地错误。 +// 由于 wire 状态码停留在 200,常规的 status>=400 捕获路径永远不会触发; +// handleStreamingAwareError 通过 service.MarkOpsStreamError 标记这类错误, +// 此函数据此补记一条错误日志,让并发限流/流内失败在错误看板里可见。 +// +// 仅在 status<400 且不存在上游错误上下文时调用:上游透传错误已由中间件的 +// upstream-context 分支落库,无需在此重复记录。 +func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int) { + streamErr, ok := service.GetOpsStreamError(c) + if !ok { + return + } + + // 命中 skip_monitoring=true 透传规则的请求跳过落库,与其它分支一致。 + if v, ok := c.Get(service.OpsSkipPassthroughKey); ok { + if skip, _ := v.(bool); skip { + return + } + } + + // 复用与 status>=400 分支相同的设置过滤(context canceled / 无可用账号等)。 + if shouldSkipOpsErrorLog(c.Request.Context(), ops, streamErr.Message, streamErr.Message, c.Request.URL.Path) { + return + } + + // 分级用「本应返回的状态码」(如并发限流 429),wire 状态码缺省时回退。 + classifyStatus := streamErr.IntendedStatus + if classifyStatus <= 0 { + classifyStatus = wireStatus + } + normalizedType := normalizeOpsErrorType(streamErr.ErrType, "") + phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, "", classifyStatus) + + apiKey := getOpsAPIKey(c) + clientRequestID, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) + + model, _ := c.Get(opsModelKey) + var modelName string + if s, ok := model.(string); ok { + modelName = s + } + accountIDV, _ := c.Get(opsAccountIDKey) + var accountID *int64 + if v, ok := accountIDV.(int64); ok && v > 0 { + accountID = &v + } + + fallbackPlatform := guessPlatformFromPath(c.Request.URL.Path) + platform := resolveOpsPlatform(apiKey, fallbackPlatform) + + requestID := c.Writer.Header().Get("X-Request-Id") + if requestID == "" { + requestID = c.Writer.Header().Get("x-request-id") + } + + entry := &service.OpsInsertErrorLogInput{ + RequestID: requestID, + ClientRequestID: clientRequestID, + + AccountID: accountID, + Platform: platform, + Model: modelName, + RequestPath: func() string { + if c.Request != nil && c.Request.URL != nil { + return c.Request.URL.Path + } + return "" + }(), + // 就地 SSE 错误只出现在流式请求上。 + Stream: true, + InboundEndpoint: GetInboundEndpoint(c), + UpstreamEndpoint: GetUpstreamEndpoint(c, platform), + RequestedModel: modelName, + UpstreamModel: func() string { + if v, ok := c.Get(opsUpstreamModelKey); ok { + if s, ok := v.(string); ok { + return strings.TrimSpace(s) + } + } + return "" + }(), + RequestType: func() *int16 { + if v, ok := c.Get(opsRequestTypeKey); ok { + switch t := v.(type) { + case int16: + return &t + case int: + v16 := int16(t) + return &v16 + } + } + return nil + }(), + UserAgent: c.GetHeader("User-Agent"), + + ErrorPhase: phase, + ErrorType: normalizedType, + Severity: classifyOpsSeverity(normalizedType, classifyStatus), + StatusCode: wireStatus, + IsBusinessLimited: isBusinessLimited, + IsCountTokens: isCountTokensRequest(c), + + ErrorMessage: streamErr.Message, + ErrorBody: "", + ErrorSource: errorSource, + ErrorOwner: errorOwner, + + CreatedAt: time.Now(), + } + applyOpsLatencyFieldsFromContext(c, entry) + + if apiKey != nil { + entry.APIKeyID = &apiKey.ID + entry.APIKeyPrefix = keyPrefix(apiKey.Key, 8) + if apiKey.User != nil { + entry.UserID = &apiKey.User.ID + } + if apiKey.GroupID != nil { + entry.GroupID = apiKey.GroupID + } + if apiKey.Group != nil && apiKey.Group.Platform != "" { + entry.Platform = apiKey.Group.Platform + } + } + + if clientIP := strings.TrimSpace(ip.GetClientIP(c)); clientIP != "" { + entry.ClientIP = &clientIP + } + + enqueueOpsErrorLog(ops, entry) +} + // isCountTokensRequest checks if the request is a count_tokens request func isCountTokensRequest(c *gin.Context) bool { if c == nil || c.Request == nil || c.Request.URL == nil { diff --git a/backend/internal/handler/ops_error_logger_test.go b/backend/internal/handler/ops_error_logger_test.go index cf1685f2a4..89bdd938c1 100644 --- a/backend/internal/handler/ops_error_logger_test.go +++ b/backend/internal/handler/ops_error_logger_test.go @@ -139,6 +139,97 @@ func TestOpsErrorLoggerMiddleware_DoesNotBreakOuterMiddlewares(t *testing.T) { require.Equal(t, http.StatusNoContent, rec.Code) } +// setupOpsErrorLogTestQueue 阻止 enqueueOpsErrorLog 启动真实 worker,改用可检查的测试队列。 +func setupOpsErrorLogTestQueue(t *testing.T, size int) { + t.Helper() + resetOpsErrorLoggerStateForTest(t) + opsErrorLogOnce.Do(func() {}) + opsErrorLogMu.Lock() + opsErrorLogQueue = make(chan opsErrorLogJob, size) + opsErrorLogMu.Unlock() +} + +// 就地(in-band) SSE 错误挂在已固化的 HTTP 200 流上:wire 状态码为 200, +// 常规 status>=400 采集路径不会触发。logOpsStreamError 必须据 MarkOpsStreamError +// 补记一条错误日志,且用 IntendedStatus(429) 分级、StatusCode 仍记 wire 的 200。 +func TestLogOpsStreamError_RecordsInBandConcurrencyLimit(t *testing.T) { + setupOpsErrorLogTestQueue(t, 4) + + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + c.Set(opsModelKey, "test-model") + + service.MarkOpsStreamError(c, "rate_limit_error", + "Concurrency limit exceeded for account, please retry later", http.StatusTooManyRequests) + + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + logOpsStreamError(c, ops, http.StatusOK) + + require.Equal(t, int64(1), OpsErrorLogEnqueuedTotal()) + require.Equal(t, int64(1), OpsErrorLogQueueLength()) + + job := <-opsErrorLogQueue + require.NotNil(t, job.entry) + require.Equal(t, "rate_limit_error", job.entry.ErrorType) + require.Equal(t, "request", job.entry.ErrorPhase) + require.True(t, job.entry.IsBusinessLimited) + require.True(t, job.entry.Stream) + require.Equal(t, http.StatusOK, job.entry.StatusCode) // wire 状态码保持 200 + require.Equal(t, "P1", job.entry.Severity) // 用 IntendedStatus 429 分级 + require.Equal(t, "test-model", job.entry.Model) + require.Equal(t, "Concurrency limit exceeded for account, please retry later", job.entry.ErrorMessage) +} + +// 未标记流内错误时 logOpsStreamError 必须是 no-op(不误记正常的 200 流)。 +func TestLogOpsStreamError_NoopWhenNotMarked(t *testing.T) { + setupOpsErrorLogTestQueue(t, 4) + + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + logOpsStreamError(c, ops, http.StatusOK) + + require.Equal(t, int64(0), OpsErrorLogEnqueuedTotal()) +} + +// 命中 skip_monitoring=true 透传规则时不落库,与其它采集分支一致。 +func TestLogOpsStreamError_SkipWhenPassthroughSkipMonitoring(t *testing.T) { + setupOpsErrorLogTestQueue(t, 4) + + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + service.MarkOpsStreamError(c, "upstream_error", "Upstream request failed", http.StatusBadGateway) + c.Set(service.OpsSkipPassthroughKey, true) + + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + logOpsStreamError(c, ops, http.StatusOK) + + require.Equal(t, int64(0), OpsErrorLogEnqueuedTotal()) +} + +// MarkOpsStreamError 采用「首个标记生效」:后续的通用兜底帧不得覆盖根因错误。 +func TestMarkOpsStreamError_FirstWins(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + service.MarkOpsStreamError(c, "rate_limit_error", "Concurrency limit exceeded for account", http.StatusTooManyRequests) + service.MarkOpsStreamError(c, "upstream_error", "Upstream request failed", http.StatusBadGateway) + + se, ok := service.GetOpsStreamError(c) + require.True(t, ok) + require.Equal(t, "rate_limit_error", se.ErrType) + require.Equal(t, "Concurrency limit exceeded for account", se.Message) + require.Equal(t, http.StatusTooManyRequests, se.IntendedStatus) +} + func TestIsKnownOpsErrorType(t *testing.T) { known := []string{ "invalid_request_error", diff --git a/backend/internal/handler/request_body_limit.go b/backend/internal/handler/request_body_limit.go index d746673b34..de24551ba9 100644 --- a/backend/internal/handler/request_body_limit.go +++ b/backend/internal/handler/request_body_limit.go @@ -4,6 +4,9 @@ import ( "errors" "fmt" "net/http" + + "github.com/Wei-Shaw/sub2api/internal/config" + pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" ) func extractMaxBytesError(err error) (*http.MaxBytesError, bool) { @@ -25,3 +28,14 @@ func formatBodyLimit(limit int64) string { func buildBodyTooLargeMessage(limit int64) string { return fmt.Sprintf("Request body too large, limit is %s", formatBodyLimit(limit)) } + +func readLenientJSONRequestBodyWithPrealloc(req *http.Request, cfg *config.Config) ([]byte, error) { + return pkghttputil.ReadLenientJSONRequestBodyWithPrealloc(req, gatewayMaxBodySize(cfg)) +} + +func gatewayMaxBodySize(cfg *config.Config) int64 { + if cfg == nil { + return 0 + } + return cfg.Gateway.MaxBodySize +} diff --git a/backend/internal/handler/request_body_parse_log.go b/backend/internal/handler/request_body_parse_log.go new file mode 100644 index 0000000000..c38a7f9bd5 --- /dev/null +++ b/backend/internal/handler/request_body_parse_log.go @@ -0,0 +1,54 @@ +package handler + +import ( + "strconv" + + "github.com/Wei-Shaw/sub2api/internal/service" + "go.uber.org/zap" +) + +// parseFailureSnippetLen bounds the head/tail snippets logged on body parse +// failure. 256 bytes is enough to see the structural context (model field, +// first content block / trailing brace) without dumping user payloads. +const parseFailureSnippetLen = 256 + +// logRequestBodyParseFailure records the real reason a request body failed +// JSON parsing/validation. The client keeps receiving the generic +// "Failed to parse request body"; the sanitized diagnostics (underlying +// error with byte offset, body length, escaped head/tail snippets) land in +// the server log only, so operators can distinguish genuinely invalid JSON +// from a truncated or partially consumed body. +// +// err may be nil for call sites that validate with gjson.ValidBytes directly; +// the diagnostic error is derived from the body in that case. +func logRequestBodyParseFailure(reqLog *zap.Logger, body []byte, err error) { + if reqLog == nil { + return + } + if err == nil { + err = service.DescribeInvalidJSON(body) + } + + head := body + var tail []byte + if len(body) > parseFailureSnippetLen { + head = body[:parseFailureSnippetLen] + tail = body[len(body)-parseFailureSnippetLen:] + } + + fields := []zap.Field{ + zap.Error(err), + zap.Int("body_len", len(body)), + zap.String("body_head", sanitizeBodySnippet(head)), + } + if len(tail) > 0 { + fields = append(fields, zap.String("body_tail", sanitizeBodySnippet(tail))) + } + reqLog.Warn("parse request body failed", fields...) +} + +// sanitizeBodySnippet escapes control characters and invalid UTF-8 so the +// snippet is always a single printable log line. +func sanitizeBodySnippet(b []byte) string { + return strconv.Quote(string(b)) +} diff --git a/backend/internal/handler/request_body_parse_log_test.go b/backend/internal/handler/request_body_parse_log_test.go new file mode 100644 index 0000000000..c1477eb4d7 --- /dev/null +++ b/backend/internal/handler/request_body_parse_log_test.go @@ -0,0 +1,100 @@ +//go:build unit + +package handler + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" +) + +func newObservedLogger(t *testing.T) (*zap.Logger, *observer.ObservedLogs) { + t.Helper() + core, logs := observer.New(zap.WarnLevel) + return zap.New(core), logs +} + +func loggedFields(t *testing.T, logs *observer.ObservedLogs) map[string]any { + t.Helper() + entries := logs.All() + require.Len(t, entries, 1) + fields := map[string]any{} + for _, f := range entries[0].Context { + switch f.Key { + case "body_len": + fields[f.Key] = int(f.Integer) + case "error": + fields[f.Key] = f.Interface.(error).Error() + default: + fields[f.Key] = f.String + } + } + return fields +} + +func TestLogRequestBodyParseFailure_DerivesErrorWhenNil(t *testing.T) { + log, logs := newObservedLogger(t) + body := []byte(`{"model": bad}`) + + logRequestBodyParseFailure(log, body, nil) + + fields := loggedFields(t, logs) + require.Equal(t, len(body), fields["body_len"]) + require.Contains(t, fields["error"], "invalid json") + require.Contains(t, fields["error"], "offset=11") +} + +func TestLogRequestBodyParseFailure_ShortBodyHasNoTail(t *testing.T) { + log, logs := newObservedLogger(t) + body := []byte(`{"broken":`) + + logRequestBodyParseFailure(log, body, nil) + + fields := loggedFields(t, logs) + require.Contains(t, fields, "body_head") + require.NotContains(t, fields, "body_tail") + require.Contains(t, fields["body_head"].(string), `{\"broken\":`) +} + +func TestLogRequestBodyParseFailure_LargeBodyBoundedSnippets(t *testing.T) { + log, logs := newObservedLogger(t) + // ~1MB body: head must show the structural prefix, tail the trailing bytes, + // and neither snippet may exceed the configured bound (plus quoting overhead). + body := []byte(`{"model":"claude-sonnet-4-6","big":"` + strings.Repeat("A", 1<<20) + `"`) + + logRequestBodyParseFailure(log, body, nil) + + fields := loggedFields(t, logs) + require.Equal(t, len(body), fields["body_len"]) + head := fields["body_head"].(string) + tail := fields["body_tail"].(string) + require.Contains(t, head, "claude-sonnet-4-6") + require.Contains(t, tail, "AAA") + require.NotContains(t, tail, "claude-sonnet-4-6") + // strconv.Quote adds surrounding quotes and escapes; 4x is a generous cap. + require.LessOrEqual(t, len(head), parseFailureSnippetLen*4) + require.LessOrEqual(t, len(tail), parseFailureSnippetLen*4) +} + +func TestLogRequestBodyParseFailure_EscapesControlCharacters(t *testing.T) { + log, logs := newObservedLogger(t) + body := []byte("{\"model\":\x01\n\"x\"}") + + logRequestBodyParseFailure(log, body, nil) + + fields := loggedFields(t, logs) + head := fields["body_head"].(string) + require.NotContains(t, head, "\n") + require.NotContains(t, head, "\x01") + require.Contains(t, head, `\n`) + require.Contains(t, head, `\x01`) +} + +func TestLogRequestBodyParseFailure_NilLoggerNoPanic(t *testing.T) { + require.NotPanics(t, func() { + logRequestBodyParseFailure(nil, []byte(`{`), nil) + }) +} diff --git a/backend/internal/handler/stream_error_event.go b/backend/internal/handler/stream_error_event.go index f3a33a8c89..2af098dbed 100644 --- a/backend/internal/handler/stream_error_event.go +++ b/backend/internal/handler/stream_error_event.go @@ -85,22 +85,27 @@ func writeResponsesFailedSSE(c *gin.Context, errType, message string) bool { return true } -// inboundIsResponses 判断当前请求是否落在任何 /responses 路由上。 +// inboundIsResponses 判断当前请求是否落在任意 Responses 路由上 +// (不区分 root 还是 compact 变体)。 // // 不能直接用 GetInboundEndpoint(c) == EndpointResponses 比较,因为 -// NormalizeInboundEndpoint 只识别包含 "/v1/responses" 子串的路径; -// 项目里实际注册了多组路由(gateway_v1、top-level bare、codex direct), -// 其中 r.POST("/responses", ...) 和 codexDirect.POST("/responses", ...) -// 的 c.FullPath() 不含 "/v1/" 前缀,会被归一化为原始路径, -// 导致协议合规终止事件没法发出去。 +// GetInboundEndpoint/NormalizeInboundEndpoint 会把 compact 变体归一化为 +// 单独的 EndpointResponsesCompact(而不是 EndpointResponses), +// 而本函数在这里只关心“是不是 Responses 家族的请求”, +// 不需要区分 root/compact,所以不能用那个等值比较。 // -// 这里用 FullPath 的后缀判断,覆盖所有变体: +// 这里改用 FullPath 的后缀/子串判断,一次性覆盖 root 和 compact 的所有变体: // - /v1/responses // - /v1/responses/compact // - /responses // - /responses/compact // - /backend-api/codex/responses // - /backend-api/codex/responses/compact +// +// 对于通配路由(如 "/v1/responses/*action")注册的 FullPath 本身就带有 +// "/responses/" 子串(例如 "/v1/responses/*action"),所以下面的 +// strings.Contains(p, "/responses/") 分支同样能覆盖这些通配路由, +// 不需要额外处理通配符本身。 func inboundIsResponses(c *gin.Context) bool { if c == nil { return false diff --git a/backend/internal/handler/usage_handler.go b/backend/internal/handler/usage_handler.go index be6dc917bb..45f72ae306 100644 --- a/backend/internal/handler/usage_handler.go +++ b/backend/internal/handler/usage_handler.go @@ -138,7 +138,7 @@ func (h *UsageHandler) parseUserUsageFilters(c *gin.Context, requireRange bool) } billingMode := strings.TrimSpace(c.Query("billing_mode")) - if billingMode != "" && !service.BillingMode(billingMode).IsValid() { + if billingMode != "" && !service.BillingMode(billingMode).IsValidUsageFilter() { response.BadRequest(c, "Invalid billing_mode") return nil, false } diff --git a/backend/internal/handler/usage_handler_request_type_test.go b/backend/internal/handler/usage_handler_request_type_test.go index 1dcb1b83a4..8dc9a8b442 100644 --- a/backend/internal/handler/usage_handler_request_type_test.go +++ b/backend/internal/handler/usage_handler_request_type_test.go @@ -162,6 +162,18 @@ func TestUserUsageListInvalidBillingMode(t *testing.T) { require.Equal(t, http.StatusBadRequest, rec.Code) } +func TestUserUsageListAllowsVideoBillingMode(t *testing.T) { + repo := &userUsageRepoCapture{} + router := newUserUsageRequestTypeTestRouter(repo) + + req := httptest.NewRequest(http.MethodGet, "/usage?billing_mode=video", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "video", repo.listFilters.BillingMode) +} + func TestUserUsageListKeepsUserBillingAndIPWithoutAdminCostFields(t *testing.T) { ipAddress := "203.0.113.10" upstreamModel := "upstream-private-model" diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index 090a734c9f..cfbb72554c 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -115,6 +115,7 @@ func ProvideHandlers( paymentHandler *PaymentHandler, paymentWebhookHandler *PaymentWebhookHandler, availableChannelHandler *AvailableChannelHandler, + batchImageHandler *BatchImageHandler, _ *service.IdempotencyCoordinator, _ *service.IdempotencyCleanupService, ) *Handlers { @@ -135,6 +136,7 @@ func ProvideHandlers( Payment: paymentHandler, PaymentWebhook: paymentWebhookHandler, AvailableChannel: availableChannelHandler, + BatchImage: batchImageHandler, } } @@ -156,6 +158,7 @@ var ProviderSet = wire.NewSet( NewPaymentHandler, NewPaymentWebhookHandler, NewAvailableChannelHandler, + NewBatchImageHandler, // Admin handlers admin.NewDashboardHandler, diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go index 88749207c9..52c93be761 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go @@ -42,6 +42,9 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR if len(out.Tools) > 0 && len(req.ToolChoice) > 0 { out.ToolChoice = responsesToolChoiceToChatToolChoice(req.ToolChoice) } + if req.Text != nil { + out.ResponseFormat = responsesTextFormatToChatResponseFormat(req.Text.Format) + } return out, nil } diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_test.go index 3e55e23a81..b194d88141 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_test.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_test.go @@ -73,6 +73,60 @@ func TestResponsesToChatCompletionsRequest_InstructionsAndInputDeveloperRole(t * assert.JSONEq(t, `"Hello"`, string(out.Messages[2].Content)) } +func TestResponsesToChatCompletionsRequest_TextFormatJsonObject(t *testing.T) { + req := &ResponsesRequest{ + Model: "gpt-4o", + Input: json.RawMessage(`[ + {"role":"user","content":"Return JSON"} + ]`), + Text: &ResponsesText{ + Format: json.RawMessage(`{"type":"json_object"}`), + }, + } + + out, err := ResponsesToChatCompletionsRequest(req) + require.NoError(t, err) + assert.JSONEq(t, `{"type":"json_object"}`, string(out.ResponseFormat)) +} + +func TestResponsesToChatCompletionsRequest_TextFormatJsonSchema(t *testing.T) { + req := &ResponsesRequest{ + Model: "gpt-4o", + Input: json.RawMessage(`[ + {"role":"user","content":"Return structured JSON"} + ]`), + Text: &ResponsesText{ + Format: json.RawMessage(`{ + "type":"json_schema", + "name":"answer", + "schema":{ + "type":"object", + "properties":{"ok":{"type":"boolean"}}, + "required":["ok"], + "additionalProperties":false + }, + "strict":true + }`), + }, + } + + out, err := ResponsesToChatCompletionsRequest(req) + require.NoError(t, err) + assert.JSONEq(t, `{ + "type":"json_schema", + "json_schema":{ + "name":"answer", + "schema":{ + "type":"object", + "properties":{"ok":{"type":"boolean"}}, + "required":["ok"], + "additionalProperties":false + }, + "strict":true + } + }`, string(out.ResponseFormat)) +} + func chatMessageRoles(messages []ChatMessage) []string { roles := make([]string, 0, len(messages)) for _, message := range messages { diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_test.go index fdb008b261..66afcf2270 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_test.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_test.go @@ -242,6 +242,62 @@ func TestChatCompletionsToResponses_ReasoningEffort(t *testing.T) { assert.Equal(t, "auto", resp.Reasoning.Summary) } +func TestChatCompletionsToResponses_ResponseFormatJsonObject(t *testing.T) { + req := &ChatCompletionsRequest{ + Model: "gpt-4o", + Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`"Return JSON"`)}}, + ResponseFormat: json.RawMessage(`{"type":"json_object"}`), + } + + resp, err := ChatCompletionsToResponses(req) + require.NoError(t, err) + require.NotNil(t, resp.Text) + assert.JSONEq(t, `{"type":"json_object"}`, string(resp.Text.Format)) + + payload, err := json.Marshal(resp) + require.NoError(t, err) + var serialized struct { + Text ResponsesText `json:"text"` + } + require.NoError(t, json.Unmarshal(payload, &serialized)) + assert.JSONEq(t, `{"type":"json_object"}`, string(serialized.Text.Format)) +} + +func TestChatCompletionsToResponses_ResponseFormatJsonSchema(t *testing.T) { + req := &ChatCompletionsRequest{ + Model: "gpt-4o", + Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`"Return structured JSON"`)}}, + ResponseFormat: json.RawMessage(`{ + "type":"json_schema", + "json_schema":{ + "name":"answer", + "schema":{ + "type":"object", + "properties":{"ok":{"type":"boolean"}}, + "required":["ok"], + "additionalProperties":false + }, + "strict":true + } + }`), + } + + resp, err := ChatCompletionsToResponses(req) + require.NoError(t, err) + require.NotNil(t, resp.Text) + assert.JSONEq(t, `{ + "type":"json_schema", + "name":"answer", + "schema":{ + "type":"object", + "properties":{"ok":{"type":"boolean"}}, + "required":["ok"], + "additionalProperties":false + }, + "strict":true + }`, string(resp.Text.Format)) +} + func TestChatCompletionsToResponses_ImageURL(t *testing.T) { content := `[{"type":"text","text":"Describe this"},{"type":"image_url","image_url":{"url":"data:image/png;base64,abc123"}}]` req := &ChatCompletionsRequest{ diff --git a/backend/internal/pkg/apicompat/chatcompletions_to_responses.go b/backend/internal/pkg/apicompat/chatcompletions_to_responses.go index 7cbb4f5f20..07c557ab3b 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_to_responses.go +++ b/backend/internal/pkg/apicompat/chatcompletions_to_responses.go @@ -69,6 +69,13 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest, } } + if format := chatResponseFormatToResponsesTextFormat(req.ResponseFormat); len(format) > 0 { + if out.Text == nil { + out.Text = &ResponsesText{} + } + out.Text.Format = format + } + // tools[] and legacy functions[] → ResponsesTool[] if len(req.Tools) > 0 || len(req.Functions) > 0 { out.Tools = convertChatToolsToResponses(req.Tools, req.Functions) diff --git a/backend/internal/pkg/apicompat/response_format.go b/backend/internal/pkg/apicompat/response_format.go new file mode 100644 index 0000000000..afb5c3e2fd --- /dev/null +++ b/backend/internal/pkg/apicompat/response_format.go @@ -0,0 +1,92 @@ +package apicompat + +import "encoding/json" + +func chatResponseFormatToResponsesTextFormat(raw json.RawMessage) json.RawMessage { + raw = normalizedRawJSON(raw) + if len(raw) == 0 { + return nil + } + + obj, ok := rawJSONObject(raw) + if !ok || rawString(obj["type"]) != "json_schema" { + return raw + } + + schemaRaw := normalizedRawJSON(obj["json_schema"]) + if len(schemaRaw) == 0 { + return raw + } + + var schema map[string]json.RawMessage + if err := json.Unmarshal(schemaRaw, &schema); err != nil { + return raw + } + schema["type"] = rawJSONString("json_schema") + + out, err := json.Marshal(schema) + if err != nil { + return raw + } + return out +} + +func responsesTextFormatToChatResponseFormat(raw json.RawMessage) json.RawMessage { + raw = normalizedRawJSON(raw) + if len(raw) == 0 { + return nil + } + + obj, ok := rawJSONObject(raw) + if !ok || rawString(obj["type"]) != "json_schema" { + return raw + } + if _, alreadyChatShape := obj["json_schema"]; alreadyChatShape { + return raw + } + + schema := make(map[string]json.RawMessage, len(obj)) + for key, value := range obj { + if key == "type" { + continue + } + schema[key] = value + } + if len(schema) == 0 { + return raw + } + + schemaRaw, err := json.Marshal(schema) + if err != nil { + return raw + } + out, err := json.Marshal(map[string]json.RawMessage{ + "type": rawJSONString("json_schema"), + "json_schema": schemaRaw, + }) + if err != nil { + return raw + } + return out +} + +func normalizedRawJSON(raw json.RawMessage) json.RawMessage { + raw = bytesTrimSpace(raw) + if len(raw) == 0 || string(raw) == "null" { + return nil + } + return append(json.RawMessage(nil), raw...) +} + +func rawJSONObject(raw json.RawMessage) (map[string]json.RawMessage, bool) { + var obj map[string]json.RawMessage + if err := json.Unmarshal(raw, &obj); err != nil { + return nil, false + } + return obj, true +} + +func rawJSONString(value string) json.RawMessage { + data, _ := json.Marshal(value) + return data +} diff --git a/backend/internal/pkg/apicompat/responses_to_anthropic_instructions_test.go b/backend/internal/pkg/apicompat/responses_to_anthropic_instructions_test.go new file mode 100644 index 0000000000..b63787cd44 --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_to_anthropic_instructions_test.go @@ -0,0 +1,126 @@ +package apicompat + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResponsesToAnthropicRequest_Instructions(t *testing.T) { + t.Run("instructions_becomes_system", func(t *testing.T) { + req := &ResponsesRequest{ + Model: "claude-sonnet-4-20250514", + Instructions: "You are a helpful assistant.", + Input: json.RawMessage(`[{"role":"user","content":"hello"}]`), + } + + result, err := ResponsesToAnthropicRequest(req) + require.NoError(t, err) + + var system string + require.NoError(t, json.Unmarshal(result.System, &system)) + assert.Equal(t, "You are a helpful assistant.", system) + assert.NotEmpty(t, result.Messages) + }) + + t.Run("empty_instructions_no_system", func(t *testing.T) { + req := &ResponsesRequest{ + Model: "claude-sonnet-4-20250514", + Input: json.RawMessage(`[{"role":"user","content":"hello"}]`), + } + + result, err := ResponsesToAnthropicRequest(req) + require.NoError(t, err) + assert.Nil(t, result.System) + }) + + t.Run("instructions_and_system_item_concatenated", func(t *testing.T) { + req := &ResponsesRequest{ + Model: "claude-sonnet-4-20250514", + Instructions: "Top-level instruction.", + Input: json.RawMessage(`[ + {"role":"system","content":"Input-level system prompt."}, + {"role":"user","content":"hello"} + ]`), + } + + result, err := ResponsesToAnthropicRequest(req) + require.NoError(t, err) + + var system string + require.NoError(t, json.Unmarshal(result.System, &system)) + assert.Contains(t, system, "Top-level instruction.") + assert.Contains(t, system, "Input-level system prompt.") + }) + + t.Run("instructions_with_string_input", func(t *testing.T) { + req := &ResponsesRequest{ + Model: "claude-sonnet-4-20250514", + Instructions: "Be concise.", + Input: json.RawMessage(`"What is Go?"`), + } + + result, err := ResponsesToAnthropicRequest(req) + require.NoError(t, err) + + var system string + require.NoError(t, json.Unmarshal(result.System, &system)) + assert.Equal(t, "Be concise.", system) + require.Len(t, result.Messages, 1) + assert.Equal(t, "user", result.Messages[0].Role) + }) +} + +func TestConvertResponsesInputToAnthropic_DeveloperRole(t *testing.T) { + t.Run("developer_becomes_system", func(t *testing.T) { + input := `[ + {"role":"developer","content":[{"type":"input_text","text":"You are a code reviewer."}]}, + {"role":"user","content":"review this code"} + ]` + + system, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input)) + require.NoError(t, err) + + var systemText string + require.NoError(t, json.Unmarshal(system, &systemText)) + assert.Equal(t, "You are a code reviewer.", systemText) + + require.Len(t, messages, 1) + assert.Equal(t, "user", messages[0].Role) + }) + + t.Run("developer_does_not_become_user", func(t *testing.T) { + input := `[ + {"role":"developer","content":[{"type":"input_text","text":"System prompt."}]}, + {"role":"user","content":"hi"} + ]` + + _, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input)) + require.NoError(t, err) + + for _, m := range messages { + if m.Role == "user" { + var s string + if json.Unmarshal(m.Content, &s) == nil { + assert.NotContains(t, s, "System prompt.") + } + } + } + }) + + t.Run("instructions_and_developer_concatenated_in_order", func(t *testing.T) { + input := `[ + {"role":"developer","content":"Extra context."}, + {"role":"user","content":"hello"} + ]` + + system, _, err := convertResponsesInputToAnthropic("Main instruction.", json.RawMessage(input)) + require.NoError(t, err) + + var systemText string + require.NoError(t, json.Unmarshal(system, &systemText)) + assert.Equal(t, "Main instruction.\n\nExtra context.", systemText) + }) +} diff --git a/backend/internal/pkg/apicompat/responses_to_anthropic_request.go b/backend/internal/pkg/apicompat/responses_to_anthropic_request.go index 6da249ed25..46f57d0ca6 100644 --- a/backend/internal/pkg/apicompat/responses_to_anthropic_request.go +++ b/backend/internal/pkg/apicompat/responses_to_anthropic_request.go @@ -11,7 +11,7 @@ import ( // enables Anthropic platform groups to accept OpenAI Responses API requests // by converting them to the native /v1/messages format before forwarding upstream. func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, error) { - system, messages, err := convertResponsesInputToAnthropic(req.Input) + system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input) if err != nil { return nil, err } @@ -98,14 +98,23 @@ func mapResponsesEffortToAnthropic(effort string) string { } // convertResponsesInputToAnthropic extracts system prompt and messages from -// a Responses API input array. Returns the system as raw JSON (for Anthropic's -// polymorphic system field) and a list of Anthropic messages. -func convertResponsesInputToAnthropic(inputRaw json.RawMessage) (json.RawMessage, []AnthropicMessage, error) { +// a Responses API instructions + input array. Returns the system as raw JSON +// (for Anthropic's polymorphic system field) and a list of Anthropic messages. +func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMessage) (json.RawMessage, []AnthropicMessage, error) { + var systemParts []string + if strings.TrimSpace(instructions) != "" { + systemParts = append(systemParts, strings.TrimSpace(instructions)) + } + // Try as plain string input. var inputStr string if err := json.Unmarshal(inputRaw, &inputStr); err == nil { content, _ := json.Marshal(inputStr) - return nil, []AnthropicMessage{{Role: "user", Content: content}}, nil + var system json.RawMessage + if len(systemParts) > 0 { + system, _ = json.Marshal(strings.Join(systemParts, "\n\n")) + } + return system, []AnthropicMessage{{Role: "user", Content: content}}, nil } var items []ResponsesInputItem @@ -113,16 +122,14 @@ func convertResponsesInputToAnthropic(inputRaw json.RawMessage) (json.RawMessage return nil, nil, fmt.Errorf("parse responses input: %w", err) } - var system json.RawMessage var messages []AnthropicMessage for _, item := range items { switch { - case item.Role == "system": - // System prompt → Anthropic system field + case item.Role == "system" || item.Role == "developer": text := extractTextFromContent(item.Content) if text != "" { - system, _ = json.Marshal(text) + systemParts = append(systemParts, text) } case item.Type == "function_call": @@ -201,6 +208,11 @@ func convertResponsesInputToAnthropic(inputRaw json.RawMessage) (json.RawMessage messages = normalizeAnthropicToolPairing(messages) messages = mergeConsecutiveMessages(messages) + var system json.RawMessage + if len(systemParts) > 0 { + system, _ = json.Marshal(strings.Join(systemParts, "\n\n")) + } + return system, messages, nil } diff --git a/backend/internal/pkg/apicompat/responses_to_anthropic_tool_pairing_test.go b/backend/internal/pkg/apicompat/responses_to_anthropic_tool_pairing_test.go index b2522f274b..1a51b5478d 100644 --- a/backend/internal/pkg/apicompat/responses_to_anthropic_tool_pairing_test.go +++ b/backend/internal/pkg/apicompat/responses_to_anthropic_tool_pairing_test.go @@ -58,7 +58,7 @@ func hasToolResult(blocks []AnthropicContentBlock, toolUseID string) bool { func convertAnthropic(t *testing.T, input string) []AnthropicMessage { t.Helper() - _, messages, err := convertResponsesInputToAnthropic(json.RawMessage(input)) + _, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input)) require.NoError(t, err) assertAnthropicPairing(t, messages) return messages diff --git a/backend/internal/pkg/apicompat/types.go b/backend/internal/pkg/apicompat/types.go index 25baa7a080..41969bc024 100644 --- a/backend/internal/pkg/apicompat/types.go +++ b/backend/internal/pkg/apicompat/types.go @@ -216,7 +216,8 @@ type ResponsesReasoning struct { // ResponsesText configures text output options in the Responses API. type ResponsesText struct { - Verbosity string `json:"verbosity,omitempty"` // "low" | "medium" | "high" + Format json.RawMessage `json:"format,omitempty"` + Verbosity string `json:"verbosity,omitempty"` // "low" | "medium" | "high" } // ResponsesInputItem is one item in the Responses API input array. @@ -486,6 +487,7 @@ type ChatCompletionsRequest struct { ReasoningEffort string `json:"reasoning_effort,omitempty"` // "low" | "medium" | "high" | "xhigh" ServiceTier string `json:"service_tier,omitempty"` Stop json.RawMessage `json:"stop,omitempty"` // string or []string + ResponseFormat json.RawMessage `json:"response_format,omitempty"` // Legacy function calling (deprecated but still supported) Functions []ChatFunction `json:"functions,omitempty"` diff --git a/backend/internal/pkg/httputil/body.go b/backend/internal/pkg/httputil/body.go index cee129484c..2bc3b9e753 100644 --- a/backend/internal/pkg/httputil/body.go +++ b/backend/internal/pkg/httputil/body.go @@ -16,6 +16,7 @@ import ( const ( requestBodyReadInitCap = 512 requestBodyReadMaxInitCap = 1 << 20 + jsonUTF8BOMLen = 3 // maxDecompressedBodySize limits the decompressed request body to 64 MB // to prevent decompression bomb attacks. maxDecompressedBodySize = 64 << 20 @@ -64,6 +65,16 @@ func ReadRequestBodyWithPrealloc(req *http.Request) ([]byte, error) { return decoded, nil } +// ReadLenientJSONRequestBodyWithPrealloc reads a request body and normalizes +// JSON string control bytes before strict validation. +func ReadLenientJSONRequestBodyWithPrealloc(req *http.Request, maxNormalizedBytes int64) ([]byte, error) { + body, err := ReadRequestBodyWithPrealloc(req) + if err != nil { + return nil, err + } + return NormalizeLenientJSONRequestBody(body, maxNormalizedBytes) +} + func decompressRequestBody(encoding string, raw []byte) ([]byte, error) { switch encoding { case "zstd": @@ -91,3 +102,77 @@ func decompressRequestBody(encoding string, raw []byte) ([]byte, error) { return nil, errors.New("unsupported Content-Encoding") } } + +// NormalizeLenientJSONRequestBody escapes raw control bytes that broken +// OpenAI-compatible clients sometimes place inside JSON strings. +func NormalizeLenientJSONRequestBody(body []byte, maxNormalizedBytes int64) ([]byte, error) { + if maxNormalizedBytes <= 0 { + maxNormalizedBytes = maxDecompressedBodySize + } + + body = trimUTF8BOM(body) + if len(body) == 0 { + return body, nil + } + if int64(len(body)) > maxNormalizedBytes { + return nil, &http.MaxBytesError{Limit: maxNormalizedBytes} + } + + var out []byte + inString := false + escaped := false + for i, b := range body { + if inString && isJSONControlByte(b) { + if out == nil { + capHint := len(body) + 6 + if int64(capHint) > maxNormalizedBytes { + capHint = int(maxNormalizedBytes) + } + out = make([]byte, 0, capHint) + out = append(out, body[:i]...) + } + if int64(len(out)+6) > maxNormalizedBytes { + return nil, &http.MaxBytesError{Limit: maxNormalizedBytes} + } + out = appendJSONUnicodeEscape(out, b) + escaped = false + continue + } + + switch { + case escaped: + escaped = false + case inString && b == '\\': + escaped = true + case b == '"': + inString = !inString + } + + if out != nil { + if int64(len(out)+1) > maxNormalizedBytes { + return nil, &http.MaxBytesError{Limit: maxNormalizedBytes} + } + out = append(out, b) + } + } + if out != nil { + return out, nil + } + return body, nil +} + +func trimUTF8BOM(body []byte) []byte { + if len(body) >= jsonUTF8BOMLen && body[0] == 0xef && body[1] == 0xbb && body[2] == 0xbf { + return body[jsonUTF8BOMLen:] + } + return body +} + +func isJSONControlByte(b byte) bool { + return b < 0x20 || b == 0x7f +} + +func appendJSONUnicodeEscape(dst []byte, b byte) []byte { + const hex = "0123456789abcdef" + return append(dst, '\\', 'u', '0', '0', hex[b>>4], hex[b&0x0f]) +} diff --git a/backend/internal/pkg/httputil/body_lenient_json_test.go b/backend/internal/pkg/httputil/body_lenient_json_test.go new file mode 100644 index 0000000000..71ffc392b8 --- /dev/null +++ b/backend/internal/pkg/httputil/body_lenient_json_test.go @@ -0,0 +1,184 @@ +package httputil + +import ( + "bytes" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/tidwall/gjson" +) + +func TestNormalizeLenientJSONRequestBody_accepts_client_control_chars_in_strings(t *testing.T) { + tests := []struct { + name string + body []byte + path string + want string + wantRaw string + }{ + { + name: "null byte in message content", + body: []byte("{\"messages\":[{\"content\":\"hello\x00world\"}]}"), + path: "messages.0.content", + want: "hello\x00world", + wantRaw: `"hello\u0000world"`, + }, + { + name: "ansi escape in message content", + body: []byte("{\"messages\":[{\"content\":\"hello\x1b[31mred\x1b[0m\"}]}"), + path: "messages.0.content", + want: "hello\x1b[31mred\x1b[0m", + wantRaw: `"hello\u001b[31mred\u001b[0m"`, + }, + { + name: "leading UTF-8 BOM", + body: []byte("\xef\xbb\xbf{\"input\":\"hello\"}"), + path: "input", + want: "hello", + wantRaw: `"hello"`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Given + if gjson.ValidBytes(tt.body) { + t.Fatalf("test payload should reproduce strict JSON rejection: %q", tt.body) + } + + // When + got, err := NormalizeLenientJSONRequestBody(tt.body, 1024) + if err != nil { + t.Fatalf("NormalizeLenientJSONRequestBody: %v", err) + } + + // Then + if !gjson.ValidBytes(got) { + t.Fatalf("normalized body should be valid JSON: %q", got) + } + result := gjson.GetBytes(got, tt.path) + if result.String() != tt.want { + t.Fatalf("value mismatch: got %q want %q", result.String(), tt.want) + } + if result.Raw != tt.wantRaw { + t.Fatalf("raw value mismatch: got %q want %q", result.Raw, tt.wantRaw) + } + }) + } +} + +func TestNormalizeLenientJSONRequestBody_keeps_invalid_structure_invalid(t *testing.T) { + tests := []struct { + name string + body []byte + }{ + { + name: "truncated JSON", + body: []byte("{\"messages\":[{\"content\":\"hello\"}]"), + }, + { + name: "control character outside string", + body: []byte("{\"input\":\"hello\"}\x00"), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // When + got, err := NormalizeLenientJSONRequestBody(tt.body, 1024) + if err != nil { + t.Fatalf("NormalizeLenientJSONRequestBody: %v", err) + } + + // Then + if gjson.ValidBytes(got) { + t.Fatalf("normalization must not repair invalid JSON structure: %q", got) + } + }) + } +} + +func TestNormalizeLenientJSONRequestBody_allows_http_requests_with_client_control_chars(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Given + body, err := ReadLenientJSONRequestBodyWithPrealloc(r, 1024) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + // When + if !gjson.ValidBytes(body) { + http.Error(w, "Failed to parse request body", http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusAccepted) + })) + defer server.Close() + + tests := []struct { + name string + body []byte + want int + }{ + { + name: "null byte in JSON string", + body: []byte("{\"model\":\"gpt-5.5\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\x00world\"}]}"), + want: http.StatusAccepted, + }, + { + name: "ANSI escape in JSON string", + body: []byte("{\"model\":\"gpt-5.5\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\x1b[31mred\x1b[0m\"}]}"), + want: http.StatusAccepted, + }, + { + name: "leading UTF-8 BOM", + body: []byte("\xef\xbb\xbf{\"model\":\"gpt-5.5\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}"), + want: http.StatusAccepted, + }, + { + name: "truncated JSON", + body: []byte("{\"model\":\"gpt-5.5\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]"), + want: http.StatusBadRequest, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodPost, server.URL+"/v1/chat/completions", bytes.NewReader(tt.body)) + if err != nil { + t.Fatalf("NewRequest: %v", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := server.Client().Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != tt.want { + t.Fatalf("status mismatch: got %d want %d", resp.StatusCode, tt.want) + } + }) + } +} + +func TestNormalizeLenientJSONRequestBody_rejects_expansion_past_limit(t *testing.T) { + // Given + body := []byte("{\"input\":\"\x00\x00\"}") + + // When + _, err := NormalizeLenientJSONRequestBody(body, int64(len(body)+5)) + + // Then + var maxErr *http.MaxBytesError + if !errors.As(err, &maxErr) { + t.Fatalf("expected MaxBytesError, got %T %v", err, err) + } + if maxErr.Limit != int64(len(body)+5) { + t.Fatalf("limit mismatch: got %d want %d", maxErr.Limit, len(body)+5) + } +} diff --git a/backend/internal/pkg/usagestats/usage_log_types.go b/backend/internal/pkg/usagestats/usage_log_types.go index a96ea670c4..29af4d2d50 100644 --- a/backend/internal/pkg/usagestats/usage_log_types.go +++ b/backend/internal/pkg/usagestats/usage_log_types.go @@ -164,13 +164,16 @@ type UserSpendingRankingResponse struct { // UserBreakdownItem represents per-user usage breakdown within a dimension (group, model, endpoint). type UserBreakdownItem struct { - UserID int64 `json:"user_id"` - Email string `json:"email"` - Requests int64 `json:"requests"` - TotalTokens int64 `json:"total_tokens"` - Cost float64 `json:"cost"` // 标准计费 - ActualCost float64 `json:"actual_cost"` // 实际扣除 - AccountCost float64 `json:"account_cost"` // 账号成本 + UserID int64 `json:"user_id"` + Email string `json:"email"` + Requests int64 `json:"requests"` + InputTokens int64 `json:"input_tokens"` // 输入 token 累计 + OutputTokens int64 `json:"output_tokens"` // 输出 token 累计 + CacheTokens int64 `json:"cache_tokens"` // 缓存创建 + 读取 token 累计 + TotalTokens int64 `json:"total_tokens"` // 输入+输出+缓存 token 累计 + Cost float64 `json:"cost"` // 标准计费 + ActualCost float64 `json:"actual_cost"` // 实际扣除 + AccountCost float64 `json:"account_cost"` // 账号成本 } // UserBreakdownDimension specifies the dimension to filter for user breakdown. @@ -187,6 +190,8 @@ type UserBreakdownDimension struct { RequestType *int16 // filter by request_type (non-nil to enable) Stream *bool // filter by stream flag (non-nil to enable) BillingType *int8 // filter by billing_type (non-nil to enable) + // SortBy 指定排序列(空 = 默认按 actual_cost)。合法值由 repo 层 allowlist 校验。 + SortBy string } // APIKeyUsageTrendPoint represents API key usage trend data point diff --git a/backend/internal/pkg/xai/models.go b/backend/internal/pkg/xai/models.go index 4902fcb94f..3a65f32c51 100644 --- a/backend/internal/pkg/xai/models.go +++ b/backend/internal/pkg/xai/models.go @@ -10,8 +10,10 @@ type Model struct { } var defaultModels = []Model{ + {ID: "grok-4.5", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.5"}, {ID: "grok-4.3", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.3"}, {ID: "grok-build-0.1", Object: "model", OwnedBy: "xai", DisplayName: "Grok Build 0.1"}, + {ID: "grok-composer-2.5-fast", Object: "model", OwnedBy: "xai", DisplayName: "Grok Composer 2.5 Fast"}, {ID: "grok-4.20-0309-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Reasoning"}, {ID: "grok-4.20-0309-non-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Non Reasoning"}, {ID: "grok-4.20-multi-agent-0309", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Multi Agent"}, @@ -39,13 +41,17 @@ func DefaultModelIDs() []string { } func DefaultModelMapping() map[string]string { - mapping := make(map[string]string, len(defaultModels)+3) + mapping := make(map[string]string, len(defaultModels)+5) for _, model := range defaultModels { mapping[model.ID] = model.ID } - mapping["grok"] = "grok-4.3" - mapping["grok-latest"] = "grok-4.3" + mapping["grok"] = "grok-4.5" + mapping["grok-latest"] = "grok-4.5" + mapping["grok-4.5-latest"] = "grok-4.5" mapping["grok-build"] = "grok-build-0.1" + mapping["grok-build-latest"] = "grok-4.5" + mapping["grok-composer"] = "grok-composer-2.5-fast" + mapping["composer-2.5"] = "grok-composer-2.5-fast" mapping["grok-4.20-reasoning"] = "grok-4.20-0309-reasoning" mapping["grok-4.20-non-reasoning"] = "grok-4.20-0309-non-reasoning" return mapping diff --git a/backend/internal/pkg/xai/oauth_test.go b/backend/internal/pkg/xai/oauth_test.go index 68a4fea240..d3d3d5cb29 100644 --- a/backend/internal/pkg/xai/oauth_test.go +++ b/backend/internal/pkg/xai/oauth_test.go @@ -207,9 +207,14 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) { t.Parallel() mapping := DefaultModelMapping() - require.Equal(t, "grok-4.3", mapping["grok"]) - require.Equal(t, "grok-4.3", mapping["grok-latest"]) + require.Equal(t, "grok-4.5", mapping["grok"]) + require.Equal(t, "grok-4.5", mapping["grok-latest"]) + require.Equal(t, "grok-4.5", mapping["grok-4.5"]) + require.Equal(t, "grok-4.5", mapping["grok-4.5-latest"]) require.Equal(t, "grok-build-0.1", mapping["grok-build"]) + require.Equal(t, "grok-4.5", mapping["grok-build-latest"]) + require.Equal(t, "grok-composer-2.5-fast", mapping["grok-composer"]) + require.Equal(t, "grok-composer-2.5-fast", mapping["composer-2.5"]) require.Equal(t, "grok-4.20-0309-reasoning", mapping["grok-4.20-reasoning"]) require.Equal(t, "grok-4.20-0309-non-reasoning", mapping["grok-4.20-non-reasoning"]) require.Equal(t, "grok-4.20-multi-agent-0309", mapping["grok-4.20-multi-agent-0309"]) diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index 76cad809c3..d348ef29b7 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -14,9 +14,11 @@ import ( "github.com/Wei-Shaw/sub2api/ent/schema/mixins" "github.com/Wei-Shaw/sub2api/ent/user" "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/lib/pq" "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "entgo.io/ent/dialect" entsql "entgo.io/ent/dialect/sql" ) @@ -177,11 +179,17 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldAllowImageGeneration, + group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, + group.FieldVideoRateIndependent, + group.FieldVideoRateMultiplier, + group.FieldVideoPrice480p, + group.FieldVideoPrice720p, + group.FieldVideoPrice1080p, group.FieldClaudeCodeOnly, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest, @@ -387,10 +395,9 @@ func (r *apiKeyRepository) deleteWithAudit(ctx context.Context, exec *dbent.Clie return nil } -func (r *apiKeyRepository) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) { +func (r *apiKeyRepository) apiKeyListByUserIDQuery(userID int64, filters service.APIKeyListFilters) *dbent.APIKeyQuery { q := r.activeQuery().Where(apikey.UserIDEQ(userID)) - // Apply filters if filters.Search != "" { q = q.Where(apikey.Or( apikey.NameContainsFold(filters.Search), @@ -408,6 +415,12 @@ func (r *apiKeyRepository) ListByUserID(ctx context.Context, userID int64, param } } + return q +} + +func (r *apiKeyRepository) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) { + q := r.apiKeyListByUserIDQuery(userID, filters) + total, err := q.Count(ctx) if err != nil { return nil, nil, err @@ -430,10 +443,119 @@ func (r *apiKeyRepository) ListByUserID(ctx context.Context, userID int64, param for i := range keys { outKeys = append(outKeys, *apiKeyEntityToService(keys[i])) } + if err := r.attachLastUsedIPs(ctx, outKeys); err != nil { + return nil, nil, err + } return outKeys, paginationResultFromTotal(int64(total), params), nil } +func (r *apiKeyRepository) ListAllByUserID(ctx context.Context, userID int64, filters service.APIKeyListFilters) ([]service.APIKey, error) { + keys, err := r.apiKeyListByUserIDQuery(userID, filters). + WithGroup(). + Order(dbent.Asc(apikey.FieldID)). + All(ctx) + if err != nil { + return nil, err + } + + outKeys := make([]service.APIKey, 0, len(keys)) + for i := range keys { + outKeys = append(outKeys, *apiKeyEntityToService(keys[i])) + } + if err := r.attachLastUsedIPs(ctx, outKeys); err != nil { + return nil, err + } + return outKeys, nil +} + +func (r *apiKeyRepository) attachLastUsedIPs(ctx context.Context, keys []service.APIKey) error { + if len(keys) == 0 || r.sql == nil { + return nil + } + + apiKeyIDs := make([]int64, 0, len(keys)) + for i := range keys { + apiKeyIDs = append(apiKeyIDs, keys[i].ID) + } + + lastUsedIPs, err := r.latestUsageLogIPs(ctx, apiKeyIDs) + if err != nil { + return err + } + for i := range keys { + if ip, ok := lastUsedIPs[keys[i].ID]; ok { + keys[i].LastUsedIP = &ip + } + } + return nil +} + +func (r *apiKeyRepository) latestUsageLogIPs(ctx context.Context, apiKeyIDs []int64) (result map[int64]string, err error) { + if len(apiKeyIDs) == 0 || r.sql == nil { + return map[int64]string{}, nil + } + + query, args := latestUsageLogIPsQuery(apiKeyIDs, r.client.Driver().Dialect()) + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + } + }() + + out := make(map[int64]string, len(apiKeyIDs)) + for rows.Next() { + var apiKeyID int64 + var ipAddress string + if err := rows.Scan(&apiKeyID, &ipAddress); err != nil { + return nil, err + } + out[apiKeyID] = ipAddress + } + if err := rows.Err(); err != nil { + return nil, err + } + return out, nil +} + +func latestUsageLogIPsQuery(apiKeyIDs []int64, dialectName string) (string, []any) { + if dialectName == dialect.Postgres { + return ` + SELECT api_key_id, ip_address + FROM ( + SELECT api_key_id, ip_address, + ROW_NUMBER() OVER (PARTITION BY api_key_id ORDER BY created_at DESC, id DESC) AS rn + FROM usage_logs + WHERE api_key_id = ANY($1::bigint[]) + AND ip_address IS NOT NULL + AND ip_address <> '' + ) ranked + WHERE rn = 1`, []any{pq.Array(apiKeyIDs)} + } + + placeholders := make([]string, len(apiKeyIDs)) + args := make([]any, len(apiKeyIDs)) + for i, id := range apiKeyIDs { + placeholders[i] = "?" + args[i] = id + } + return fmt.Sprintf(` + SELECT api_key_id, ip_address + FROM ( + SELECT api_key_id, ip_address, + ROW_NUMBER() OVER (PARTITION BY api_key_id ORDER BY created_at DESC, id DESC) AS rn + FROM usage_logs + WHERE api_key_id IN (%s) + AND ip_address IS NOT NULL + AND ip_address <> '' + ) ranked + WHERE rn = 1`, strings.Join(placeholders, ", ")), args +} + func (r *apiKeyRepository) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) { if len(apiKeyIDs) == 0 { return []int64{}, nil @@ -503,14 +625,24 @@ func apiKeyListOrder(params pagination.PaginationParams) []func(*entsql.Selector field = apikey.FieldLastUsedAt case "created_at": field = apikey.FieldCreatedAt + case "id": + field = apikey.FieldID default: field = apikey.FieldID } if sortOrder == pagination.SortOrderAsc { - return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(apikey.FieldID)} + orders := []func(*entsql.Selector){dbent.Asc(field)} + if field != apikey.FieldID { + orders = append(orders, dbent.Asc(apikey.FieldID)) + } + return orders } - return []func(*entsql.Selector){dbent.Desc(field), dbent.Desc(apikey.FieldID)} + orders := []func(*entsql.Selector){dbent.Desc(field)} + if field != apikey.FieldID { + orders = append(orders, dbent.Desc(apikey.FieldID)) + } + return orders } // SearchAPIKeys searches API keys by user ID and/or keyword (name) @@ -755,6 +887,7 @@ func userEntityToService(u *dbent.User) *service.User { PasswordHash: u.PasswordHash, Role: u.Role, Balance: u.Balance, + FrozenBalance: u.FrozenBalance, Concurrency: u.Concurrency, Status: u.Status, SignupSource: u.SignupSource, @@ -797,11 +930,19 @@ func groupEntityToService(g *dbent.Group) *service.Group { WeeklyLimitUSD: g.WeeklyLimitUsd, MonthlyLimitUSD: g.MonthlyLimitUsd, AllowImageGeneration: g.AllowImageGeneration, + AllowBatchImageGeneration: g.AllowBatchImageGeneration, ImageRateIndependent: g.ImageRateIndependent, ImageRateMultiplier: g.ImageRateMultiplier, ImagePrice1K: g.ImagePrice1k, ImagePrice2K: g.ImagePrice2k, ImagePrice4K: g.ImagePrice4k, + BatchImageDiscountMultiplier: g.BatchImageDiscountMultiplier, + BatchImageHoldMultiplier: g.BatchImageHoldMultiplier, + VideoRateIndependent: g.VideoRateIndependent, + VideoRateMultiplier: g.VideoRateMultiplier, + VideoPrice480P: g.VideoPrice480p, + VideoPrice720P: g.VideoPrice720p, + VideoPrice1080P: g.VideoPrice1080p, DefaultValidityDays: g.DefaultValidityDays, ClaudeCodeOnly: g.ClaudeCodeOnly, FallbackGroupID: g.FallbackGroupID, diff --git a/backend/internal/repository/api_key_repo_last_used_unit_test.go b/backend/internal/repository/api_key_repo_last_used_unit_test.go index 7c6e2850e8..839eda7f75 100644 --- a/backend/internal/repository/api_key_repo_last_used_unit_test.go +++ b/backend/internal/repository/api_key_repo_last_used_unit_test.go @@ -8,6 +8,7 @@ import ( dbent "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/ent/enttest" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/stretchr/testify/require" @@ -30,7 +31,7 @@ func newAPIKeyRepoSQLite(t *testing.T) (*apiKeyRepository, *dbent.Client) { client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(drv))) t.Cleanup(func() { _ = client.Close() }) - return &apiKeyRepository{client: client}, client + return &apiKeyRepository{client: client, sql: db}, client } func mustCreateAPIKeyRepoUser(t *testing.T, ctx context.Context, client *dbent.Client, email string) *service.User { @@ -45,6 +46,85 @@ func mustCreateAPIKeyRepoUser(t *testing.T, ctx context.Context, client *dbent.C return userEntityToService(u) } +func mustCreateAPIKeyRepoAccount(t *testing.T, ctx context.Context, client *dbent.Client, name string) int64 { + t.Helper() + a, err := client.Account.Create(). + SetName(name). + SetPlatform(service.PlatformOpenAI). + SetType(service.AccountTypeAPIKey). + SetStatus(service.StatusActive). + SetCredentials(map[string]any{"api_key": "sk-test"}). + Save(ctx) + require.NoError(t, err) + return a.ID +} + +func mustCreateAPIKeyRepoUsageLog(t *testing.T, ctx context.Context, client *dbent.Client, userID, apiKeyID, accountID int64, requestID string, createdAt time.Time, ipAddress *string) { + t.Helper() + builder := client.UsageLog.Create(). + SetUserID(userID). + SetAPIKeyID(apiKeyID). + SetAccountID(accountID). + SetRequestID(requestID). + SetModel("gpt-5"). + SetCreatedAt(createdAt) + if ipAddress != nil { + builder.SetIPAddress(*ipAddress) + } + _, err := builder.Save(ctx) + require.NoError(t, err) +} + +func TestAPIKeyRepositoryListByUserIDAttachesLastUsedIP(t *testing.T) { + repo, client := newAPIKeyRepoSQLite(t) + ctx := context.Background() + user := mustCreateAPIKeyRepoUser(t, ctx, client, "list-last-used-ip@test.com") + accountID := mustCreateAPIKeyRepoAccount(t, ctx, client, "acc-list-last-used-ip") + + withLogs := &service.APIKey{ + UserID: user.ID, + Key: "sk-list-last-used-ip-logs", + Name: "With Logs", + Status: service.StatusActive, + } + emptyOnly := &service.APIKey{ + UserID: user.ID, + Key: "sk-list-last-used-ip-empty", + Name: "Empty Only", + Status: service.StatusActive, + } + noLogs := &service.APIKey{ + UserID: user.ID, + Key: "sk-list-last-used-ip-none", + Name: "No Logs", + Status: service.StatusActive, + } + require.NoError(t, repo.Create(ctx, withLogs)) + require.NoError(t, repo.Create(ctx, emptyOnly)) + require.NoError(t, repo.Create(ctx, noLogs)) + + olderIP := "198.51.100.10" + newerEmptyIP := "" + newestIP := "203.0.113.20" + base := time.Now().UTC().Add(-3 * time.Hour).Truncate(time.Second) + mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, withLogs.ID, accountID, "req-last-ip-older", base, &olderIP) + mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, withLogs.ID, accountID, "req-last-ip-empty", base.Add(time.Hour), &newerEmptyIP) + mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, withLogs.ID, accountID, "req-last-ip-newest", base.Add(2*time.Hour), &newestIP) + mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, emptyOnly.ID, accountID, "req-empty-ip", base.Add(3*time.Hour), &newerEmptyIP) + + keys, _, err := repo.ListByUserID(ctx, user.ID, pagination.PaginationParams{Page: 1, PageSize: 10}, service.APIKeyListFilters{}) + require.NoError(t, err) + + byID := make(map[int64]service.APIKey, len(keys)) + for _, key := range keys { + byID[key.ID] = key + } + require.NotNil(t, byID[withLogs.ID].LastUsedIP) + require.Equal(t, newestIP, *byID[withLogs.ID].LastUsedIP) + require.Nil(t, byID[emptyOnly.ID].LastUsedIP) + require.Nil(t, byID[noLogs.ID].LastUsedIP) +} + func TestAPIKeyRepository_CreateWithLastUsedAt(t *testing.T) { repo, client := newAPIKeyRepoSQLite(t) ctx := context.Background() diff --git a/backend/internal/repository/api_key_repo_sort_integration_test.go b/backend/internal/repository/api_key_repo_sort_integration_test.go index 69812882fe..e636fd848b 100644 --- a/backend/internal/repository/api_key_repo_sort_integration_test.go +++ b/backend/internal/repository/api_key_repo_sort_integration_test.go @@ -23,3 +23,20 @@ func (s *APIKeyRepoSuite) TestListByUserID_SortByNameAsc() { s.Require().Equal("a-key", keys[0].Name) s.Require().Equal("z-key", keys[1].Name) } + +func (s *APIKeyRepoSuite) TestListByUserID_SortByID() { + user := s.mustCreateUser("sort-id@example.com") + first := s.mustCreateApiKey(user.ID, "sk-id-a", "a-key", nil) + second := s.mustCreateApiKey(user.ID, "sk-id-b", "b-key", nil) + + keys, _, err := s.repo.ListByUserID(s.ctx, user.ID, pagination.PaginationParams{ + Page: 1, + PageSize: 10, + SortBy: "id", + SortOrder: "desc", + }, service.APIKeyListFilters{}) + s.Require().NoError(err) + s.Require().Len(keys, 2) + s.Require().Equal(second.ID, keys[0].ID) + s.Require().Equal(first.ID, keys[1].ID) +} diff --git a/backend/internal/repository/batch_image_download_limiter.go b/backend/internal/repository/batch_image_download_limiter.go new file mode 100644 index 0000000000..ffcd0dfe81 --- /dev/null +++ b/backend/internal/repository/batch_image_download_limiter.go @@ -0,0 +1,102 @@ +package repository + +import ( + "context" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/redis/go-redis/v9" +) + +const ( + defaultBatchImageDownloadActivePrefix = "batch_image:download:active:" + defaultBatchImageDownloadActiveTTL = 10 * time.Minute + defaultBatchImageDownloadConcurrency = 2 +) + +var batchImageDownloadAcquireScript = redis.NewScript(` +local current = tonumber(redis.call("GET", KEYS[1]) or "0") +local max = tonumber(ARGV[1]) +if current >= max then + return 0 +end +redis.call("INCR", KEYS[1]) +redis.call("EXPIRE", KEYS[1], ARGV[2]) +return 1 +`) + +var batchImageDownloadReleaseScript = redis.NewScript(` +local current = tonumber(redis.call("GET", KEYS[1]) or "0") +if current <= 1 then + redis.call("DEL", KEYS[1]) + return 0 +end +return redis.call("DECR", KEYS[1]) +`) + +type batchImageDownloadLimiter struct { + rdb *redis.Client + activePrefix string + maxActive int + ttl time.Duration +} + +func NewBatchImageDownloadLimiter(rdb *redis.Client, cfg *config.Config) service.BatchImageDownloadLimiter { + maxActive := defaultBatchImageDownloadConcurrency + ttl := defaultBatchImageDownloadActiveTTL + if cfg != nil { + if cfg.BatchImage.MaxDownloadConcurrencyPerUser > 0 { + maxActive = cfg.BatchImage.MaxDownloadConcurrencyPerUser + } + if cfg.BatchImage.MaxDownloadDurationSeconds > 0 { + ttl = time.Duration(cfg.BatchImage.MaxDownloadDurationSeconds) * time.Second + } + } + return &batchImageDownloadLimiter{ + rdb: rdb, + activePrefix: defaultBatchImageDownloadActivePrefix, + maxActive: maxActive, + ttl: ttl, + } +} + +func (l *batchImageDownloadLimiter) Acquire(ctx context.Context, userID string, kind string) (service.BatchImageDownloadPermit, error) { + if l == nil || l.rdb == nil { + return nil, service.ErrBatchImageDownloadLimited + } + key := l.activeKey(userID) + ok, err := batchImageDownloadAcquireScript.Run(ctx, l.rdb, []string{key}, l.maxActive, int(l.ttl.Seconds())).Int() + if err != nil { + return nil, err + } + if ok != 1 { + return nil, service.ErrBatchImageDownloadLimited + } + return &batchImageDownloadPermit{rdb: l.rdb, key: key}, nil +} + +func (l *batchImageDownloadLimiter) activeKey(userID string) string { + return l.activePrefix + userID +} + +type batchImageDownloadPermit struct { + rdb *redis.Client + key string + once sync.Once + err error +} + +func (p *batchImageDownloadPermit) Release(ctx context.Context) error { + if p == nil || p.rdb == nil || p.key == "" { + return nil + } + p.once.Do(func() { + _, p.err = batchImageDownloadReleaseScript.Run(ctx, p.rdb, []string{p.key}).Result() + }) + return p.err +} + +var _ service.BatchImageDownloadLimiter = (*batchImageDownloadLimiter)(nil) +var _ service.BatchImageDownloadPermit = (*batchImageDownloadPermit)(nil) diff --git a/backend/internal/repository/batch_image_download_limiter_test.go b/backend/internal/repository/batch_image_download_limiter_test.go new file mode 100644 index 0000000000..f261003aa6 --- /dev/null +++ b/backend/internal/repository/batch_image_download_limiter_test.go @@ -0,0 +1,43 @@ +//go:build unit + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestBatchImageDownloadLimiter_AcquireDenyReleaseAndTTL(t *testing.T) { + ctx := context.Background() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + limiter := &batchImageDownloadLimiter{ + rdb: rdb, + activePrefix: defaultBatchImageDownloadActivePrefix, + maxActive: 1, + ttl: time.Minute, + } + + permit, err := limiter.Acquire(ctx, "11", "zip") + require.NoError(t, err) + require.NotNil(t, permit) + require.True(t, mr.TTL(limiter.activeKey("11")) > 0) + + _, err = limiter.Acquire(ctx, "11", "zip") + require.ErrorIs(t, err, service.ErrBatchImageDownloadLimited) + + require.NoError(t, permit.Release(ctx)) + require.NoError(t, permit.Release(ctx)) + require.False(t, mr.Exists(limiter.activeKey("11"))) + + permit, err = limiter.Acquire(ctx, "11", "zip") + require.NoError(t, err) + require.NotNil(t, permit) +} diff --git a/backend/internal/repository/batch_image_queue.go b/backend/internal/repository/batch_image_queue.go new file mode 100644 index 0000000000..e10bc1584a --- /dev/null +++ b/backend/internal/repository/batch_image_queue.go @@ -0,0 +1,354 @@ +package repository + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/redis/go-redis/v9" +) + +const ( + defaultBatchImageReadyKey = "batch_image:queue:ready" + defaultBatchImageDelayedKey = "batch_image:queue:delayed" + defaultBatchImageActiveKey = "batch_image:queue:active" + defaultBatchImageInflightPrefix = "batch_image:queue:inflight:" + defaultBatchImageLockPrefix = "batch_image:queue:lock:" + defaultBatchImageInflightTTL = 7 * 24 * time.Hour + defaultBatchImageJobLockTTL = 5 * time.Minute + + // batchImageReservePollInterval 是原子 Reserve 脚本空轮询的间隔。 + // 用轮询替代 BRPop 是为了保证 "弹出 + 写 active" 的原子性。 + batchImageReservePollInterval = time.Second +) + +var batchImageMoveDueDelayedScript = redis.NewScript(` +local jobs = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1], "LIMIT", 0, ARGV[2]) +for _, job in ipairs(jobs) do + redis.call("ZREM", KEYS[1], job) + redis.call("LPUSH", KEYS[2], job) +end +return #jobs +`) + +var batchImageRecoverStaleActiveScript = redis.NewScript(` +local jobs = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1], "LIMIT", 0, ARGV[2]) +for _, job in ipairs(jobs) do + redis.call("ZREM", KEYS[1], job) + redis.call("LPUSH", KEYS[2], job) +end +return #jobs +`) + +var batchImageReleaseLockScript = redis.NewScript(` +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) +end +return 0 +`) + +var batchImageRefreshLockScript = redis.NewScript(` +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("PEXPIRE", KEYS[1], ARGV[2]) +end +return 0 +`) + +// batchImageReserveScript 原子地从 ready 弹出并写入 active zset。 +// BRPop + ZAdd 两步方案在两步之间进程崩溃时 job 会脱离所有队列结构, +// 且 inflight 去重键(默认 7 天)会挡住所有重新入队。 +var batchImageReserveScript = redis.NewScript(` +local job = redis.call("RPOP", KEYS[1]) +if not job then + return nil +end +redis.call("ZADD", KEYS[2], ARGV[1], job) +return job +`) + +// batchImageEnqueueScript 原子地设置 inflight 去重键并推入 ready。 +// SetNX + LPush 两步方案在两步之间进程崩溃时,inflight 键(默认 7 天) +// 会挡住所有后续入队,而 job 从未进入 ready。 +var batchImageEnqueueScript = redis.NewScript(` +if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then + redis.call("LPUSH", KEYS[2], ARGV[1]) + return 1 +end +return 0 +`) + +type batchImageQueue struct { + rdb *redis.Client + readyKey string + delayedKey string + activeKey string + inflightPrefix string + lockPrefix string + inflightTTL time.Duration + lockTTL time.Duration +} + +func NewBatchImageQueue(rdb *redis.Client, cfg *config.Config) service.BatchImageQueue { + return newBatchImageQueueWithOptions(rdb, batchImageQueueOptionsFromConfig(cfg)) +} + +type batchImageQueueOptions struct { + ReadyKey string + DelayedKey string + ActiveKey string + InflightPrefix string + LockPrefix string + InflightTTL time.Duration + LockTTL time.Duration +} + +func newBatchImageQueueWithOptions(rdb *redis.Client, opts batchImageQueueOptions) *batchImageQueue { + opts = normalizeBatchImageQueueOptions(opts) + return &batchImageQueue{ + rdb: rdb, + readyKey: opts.ReadyKey, + delayedKey: opts.DelayedKey, + activeKey: opts.ActiveKey, + inflightPrefix: opts.InflightPrefix, + lockPrefix: opts.LockPrefix, + inflightTTL: opts.InflightTTL, + lockTTL: opts.LockTTL, + } +} + +func batchImageQueueOptionsFromConfig(cfg *config.Config) batchImageQueueOptions { + if cfg == nil { + return batchImageQueueOptions{} + } + return batchImageQueueOptions{ + ReadyKey: cfg.BatchImage.QueueReadyKey, + DelayedKey: cfg.BatchImage.QueueDelayedKey, + ActiveKey: cfg.BatchImage.QueueActiveKey, + InflightPrefix: cfg.BatchImage.InflightKeyPrefix, + LockPrefix: cfg.BatchImage.LockKeyPrefix, + InflightTTL: time.Duration(cfg.BatchImage.InflightTTLSeconds) * time.Second, + LockTTL: time.Duration(cfg.BatchImage.JobLockTTLSeconds) * time.Second, + } +} + +func normalizeBatchImageQueueOptions(opts batchImageQueueOptions) batchImageQueueOptions { + if opts.ReadyKey == "" { + opts.ReadyKey = defaultBatchImageReadyKey + } + if opts.DelayedKey == "" { + opts.DelayedKey = defaultBatchImageDelayedKey + } + if opts.ActiveKey == "" { + opts.ActiveKey = defaultBatchImageActiveKey + } + if opts.InflightPrefix == "" { + opts.InflightPrefix = defaultBatchImageInflightPrefix + } + if opts.LockPrefix == "" { + opts.LockPrefix = defaultBatchImageLockPrefix + } + if opts.InflightTTL <= 0 { + opts.InflightTTL = defaultBatchImageInflightTTL + } + if opts.LockTTL <= 0 { + opts.LockTTL = defaultBatchImageJobLockTTL + } + return opts +} + +func (q *batchImageQueue) Enqueue(ctx context.Context, batchID string) error { + if !service.IsValidBatchImageID(batchID) { + return service.ErrInvalidBatchImageQueuePayload + } + + applied, err := batchImageEnqueueScript.Run(ctx, q.rdb, + []string{q.inflightKey(batchID), q.readyKey}, + batchID, q.inflightTTL.Milliseconds(), + ).Int() + if err != nil { + return err + } + if applied == 0 { + return service.ErrBatchImageAlreadyQueued + } + return nil +} + +func (q *batchImageQueue) Reserve(ctx context.Context, blockTimeout time.Duration) (service.ReservedBatchImageJob, error) { + deadline := time.Now().Add(blockTimeout) + for { + batchID, err := q.reserveOnce(ctx) + if err == nil { + return service.ReservedBatchImageJob{BatchID: batchID}, nil + } + if !errors.Is(err, service.ErrBatchImageQueueEmpty) { + return service.ReservedBatchImageJob{}, err + } + remaining := time.Until(deadline) + if remaining <= 0 { + return service.ReservedBatchImageJob{}, service.ErrBatchImageQueueEmpty + } + wait := batchImageReservePollInterval + if remaining < wait { + wait = remaining + } + timer := time.NewTimer(wait) + select { + case <-ctx.Done(): + timer.Stop() + return service.ReservedBatchImageJob{}, ctx.Err() + case <-timer.C: + } + } +} + +func (q *batchImageQueue) reserveOnce(ctx context.Context) (string, error) { + raw, err := batchImageReserveScript.Run(ctx, q.rdb, []string{q.readyKey, q.activeKey}, time.Now().UnixMilli()).Result() + if errors.Is(err, redis.Nil) { + return "", service.ErrBatchImageQueueEmpty + } + if err != nil { + return "", err + } + batchID, ok := raw.(string) + if !ok || !service.IsValidBatchImageID(batchID) { + // 非法 payload 已被脚本写入 active,必须移除,否则 stale 恢复会把它 + // 无限重投回 ready。 + if ok && batchID != "" { + _ = q.rdb.ZRem(ctx, q.activeKey, batchID).Err() + } + return "", service.ErrInvalidBatchImageQueuePayload + } + return batchID, nil +} + +func (q *batchImageQueue) RequeueAfter(ctx context.Context, batchID string, delay time.Duration) error { + if !service.IsValidBatchImageID(batchID) { + return service.ErrInvalidBatchImageQueuePayload + } + pipe := q.rdb.TxPipeline() + pipe.ZRem(ctx, q.activeKey, batchID) + pipe.ZRem(ctx, q.delayedKey, batchID) + if delay <= 0 { + pipe.LPush(ctx, q.readyKey, batchID) + } else { + pipe.ZAdd(ctx, q.delayedKey, redis.Z{ + Score: float64(time.Now().Add(delay).UnixMilli()), + Member: batchID, + }) + } + _, err := pipe.Exec(ctx) + return err +} + +func (q *batchImageQueue) Ack(ctx context.Context, batchID string) error { + if !service.IsValidBatchImageID(batchID) { + return service.ErrInvalidBatchImageQueuePayload + } + pipe := q.rdb.TxPipeline() + pipe.ZRem(ctx, q.activeKey, batchID) + pipe.ZRem(ctx, q.delayedKey, batchID) + pipe.Del(ctx, q.inflightKey(batchID)) + _, err := pipe.Exec(ctx) + return err +} + +func (q *batchImageQueue) Heartbeat(ctx context.Context, batchID string) error { + if !service.IsValidBatchImageID(batchID) { + return service.ErrInvalidBatchImageQueuePayload + } + // XX:只刷新已存在的 active 成员。无条件 ZAdd 会在 Ack/Requeue 之后的 + // 竞态心跳里把幽灵成员塞回 active zset。 + return q.rdb.ZAddXX(ctx, q.activeKey, redis.Z{ + Score: float64(time.Now().UnixMilli()), + Member: batchID, + }).Err() +} + +func (q *batchImageQueue) MoveDueDelayedToReady(ctx context.Context, limit int) (int, error) { + if limit <= 0 { + limit = 100 + } + return batchImageMoveDueDelayedScript.Run(ctx, q.rdb, []string{q.delayedKey, q.readyKey}, time.Now().UnixMilli(), limit).Int() +} + +func (q *batchImageQueue) RecoverStaleActive(ctx context.Context, staleAfter time.Duration, limit int) (int, error) { + if staleAfter <= 0 { + return 0, service.ErrInvalidBatchImageQueuePayload + } + if limit <= 0 { + limit = 100 + } + cutoff := time.Now().Add(-staleAfter).UnixMilli() + return batchImageRecoverStaleActiveScript.Run(ctx, q.rdb, []string{q.activeKey, q.readyKey}, cutoff, limit).Int() +} + +func (q *batchImageQueue) TryAcquireJobLock(ctx context.Context, batchID string, ttl time.Duration) (service.BatchImageJobLock, bool, error) { + if !service.IsValidBatchImageID(batchID) { + return nil, false, service.ErrInvalidBatchImageQueuePayload + } + if ttl <= 0 { + ttl = q.lockTTL + } + token, err := newBatchImageLockToken() + if err != nil { + return nil, false, err + } + key := q.lockKey(batchID) + ok, err := q.rdb.SetNX(ctx, key, token, ttl).Result() + if err != nil { + return nil, false, err + } + if !ok { + return nil, false, nil + } + return &batchImageRedisJobLock{rdb: q.rdb, key: key, token: token}, true, nil +} + +func (q *batchImageQueue) inflightKey(batchID string) string { + return q.inflightPrefix + batchID +} + +func (q *batchImageQueue) lockKey(batchID string) string { + return q.lockPrefix + batchID +} + +type batchImageRedisJobLock struct { + rdb *redis.Client + key string + token string +} + +func (l *batchImageRedisJobLock) Release(ctx context.Context) error { + if l == nil || l.rdb == nil || l.key == "" || l.token == "" { + return nil + } + return batchImageReleaseLockScript.Run(ctx, l.rdb, []string{l.key}, l.token).Err() +} + +// Refresh 在仍持有锁(token 匹配)时续期 TTL,供长处理任务的心跳调用。 +func (l *batchImageRedisJobLock) Refresh(ctx context.Context, ttl time.Duration) error { + if l == nil || l.rdb == nil || l.key == "" || l.token == "" { + return nil + } + if ttl <= 0 { + ttl = defaultBatchImageJobLockTTL + } + return batchImageRefreshLockScript.Run(ctx, l.rdb, []string{l.key}, l.token, ttl.Milliseconds()).Err() +} + +var _ service.BatchImageJobLockRefresher = (*batchImageRedisJobLock)(nil) + +func newBatchImageLockToken() (string, error) { + var b [16]byte + if _, err := rand.Read(b[:]); err != nil { + return "", err + } + return hex.EncodeToString(b[:]), nil +} + +var _ service.BatchImageQueue = (*batchImageQueue)(nil) diff --git a/backend/internal/repository/batch_image_queue_test.go b/backend/internal/repository/batch_image_queue_test.go new file mode 100644 index 0000000000..63412ce54e --- /dev/null +++ b/backend/internal/repository/batch_image_queue_test.go @@ -0,0 +1,199 @@ +//go:build unit + +package repository + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestBatchImageQueue_DuplicateEnqueueReturnsAlreadyQueued(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + batchID := "imgbatch_duplicate" + + require.NoError(t, queue.Enqueue(ctx, batchID)) + err := queue.Enqueue(ctx, batchID) + require.Error(t, err) + require.True(t, errors.Is(err, service.ErrBatchImageAlreadyQueued)) +} + +func TestBatchImageQueue_RequeueAfterMovesJobFromActiveToDelayed(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + batchID := "imgbatch_requeue_after" + require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, redis.Z{ + Score: float64(time.Now().UnixMilli()), + Member: batchID, + }).Err()) + + require.NoError(t, queue.RequeueAfter(ctx, batchID, time.Minute)) + require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, batchID).Err(), redis.Nil) + score, err := queue.rdb.ZScore(ctx, queue.delayedKey, batchID).Result() + require.NoError(t, err) + require.Greater(t, score, float64(time.Now().UnixMilli())) +} + +func TestBatchImageQueue_MoveDueDelayedToReadyMovesDueJobs(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + dueBatchID := "imgbatch_due" + futureBatchID := "imgbatch_future" + now := time.Now() + require.NoError(t, queue.rdb.ZAdd(ctx, queue.delayedKey, + redis.Z{Score: float64(now.Add(-time.Second).UnixMilli()), Member: dueBatchID}, + redis.Z{Score: float64(now.Add(time.Hour).UnixMilli()), Member: futureBatchID}, + ).Err()) + + moved, err := queue.MoveDueDelayedToReady(ctx, 10) + require.NoError(t, err) + require.Equal(t, 1, moved) + require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.delayedKey, dueBatchID).Err(), redis.Nil) + require.NoError(t, queue.rdb.ZScore(ctx, queue.delayedKey, futureBatchID).Err()) + + reserved, err := queue.Reserve(ctx, time.Millisecond) + require.NoError(t, err) + require.Equal(t, dueBatchID, reserved.BatchID) +} + +func TestBatchImageQueue_RecoverStaleActiveMovesStaleJobsToReady(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + staleBatchID := "imgbatch_stale" + recentBatchID := "imgbatch_recent" + now := time.Now() + require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, + redis.Z{Score: float64(now.Add(-time.Hour).UnixMilli()), Member: staleBatchID}, + redis.Z{Score: float64(now.UnixMilli()), Member: recentBatchID}, + ).Err()) + + moved, err := queue.RecoverStaleActive(ctx, 10*time.Minute, 10) + require.NoError(t, err) + require.Equal(t, 1, moved) + require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, staleBatchID).Err(), redis.Nil) + require.NoError(t, queue.rdb.ZScore(ctx, queue.activeKey, recentBatchID).Err()) + + reserved, err := queue.Reserve(ctx, time.Millisecond) + require.NoError(t, err) + require.Equal(t, staleBatchID, reserved.BatchID) +} + +func TestBatchImageQueue_JobLockReleaseOnlyDeletesMatchingToken(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + batchID := "imgbatch_lock" + + lock, ok, err := queue.TryAcquireJobLock(ctx, batchID, time.Minute) + require.NoError(t, err) + require.True(t, ok) + + require.NoError(t, queue.rdb.Set(ctx, queue.lockKey(batchID), "other-token", time.Minute).Err()) + require.NoError(t, lock.Release(ctx)) + got, err := queue.rdb.Get(ctx, queue.lockKey(batchID)).Result() + require.NoError(t, err) + require.Equal(t, "other-token", got) + + require.NoError(t, queue.rdb.Del(ctx, queue.lockKey(batchID)).Err()) + lock, ok, err = queue.TryAcquireJobLock(ctx, batchID, time.Minute) + require.NoError(t, err) + require.True(t, ok) + require.NoError(t, lock.Release(ctx)) + require.ErrorIs(t, queue.rdb.Get(ctx, queue.lockKey(batchID)).Err(), redis.Nil) +} + +func TestBatchImageQueue_ReserveAtomicallyMovesJobToActive(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + batchID := "imgbatch_reserve" + require.NoError(t, queue.Enqueue(ctx, batchID)) + + reserved, err := queue.Reserve(ctx, time.Second) + require.NoError(t, err) + require.Equal(t, batchID, reserved.BatchID) + + // 弹出与写入 active 必须原子完成:ready 已空,active 中有该 job。 + require.Equal(t, int64(0), queue.rdb.LLen(ctx, queue.readyKey).Val()) + score, err := queue.rdb.ZScore(ctx, queue.activeKey, batchID).Result() + require.NoError(t, err) + require.Positive(t, score) +} + +func TestBatchImageQueue_ReserveReturnsEmptyAfterTimeout(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + + start := time.Now() + _, err := queue.Reserve(ctx, 50*time.Millisecond) + require.ErrorIs(t, err, service.ErrBatchImageQueueEmpty) + require.Less(t, time.Since(start), 5*time.Second) +} + +func TestBatchImageQueue_ReserveDropsInvalidPayload(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + require.NoError(t, queue.rdb.LPush(ctx, queue.readyKey, "not-a-batch-id").Err()) + + _, err := queue.Reserve(ctx, 10*time.Millisecond) + require.ErrorIs(t, err, service.ErrInvalidBatchImageQueuePayload) + // 非法 payload 不得残留在 active zset,否则 stale 恢复会无限重投。 + require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, "not-a-batch-id").Err(), redis.Nil) +} + +func TestBatchImageQueue_HeartbeatOnlyRefreshesExistingActiveMember(t *testing.T) { + ctx := context.Background() + queue, _ := newBatchImageQueueTest(t) + batchID := "imgbatch_heartbeat" + + // 不在 active 中:心跳不得创建幽灵成员。 + require.NoError(t, queue.Heartbeat(ctx, batchID)) + require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, batchID).Err(), redis.Nil) + + require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, redis.Z{Score: 1, Member: batchID}).Err()) + require.NoError(t, queue.Heartbeat(ctx, batchID)) + score, err := queue.rdb.ZScore(ctx, queue.activeKey, batchID).Result() + require.NoError(t, err) + require.Greater(t, score, float64(1)) +} + +func TestBatchImageQueue_JobLockRefreshExtendsTTLOnlyForHolder(t *testing.T) { + ctx := context.Background() + queue, mr := newBatchImageQueueTest(t) + batchID := "imgbatch_lock_refresh" + + lock, ok, err := queue.TryAcquireJobLock(ctx, batchID, time.Minute) + require.NoError(t, err) + require.True(t, ok) + refresher, isRefresher := lock.(service.BatchImageJobLockRefresher) + require.True(t, isRefresher) + + require.NoError(t, refresher.Refresh(ctx, 10*time.Minute)) + ttl := mr.TTL(queue.lockKey(batchID)) + require.Greater(t, ttl, 5*time.Minute) + + // token 不匹配时不得续期他人持有的锁。 + require.NoError(t, queue.rdb.Set(ctx, queue.lockKey(batchID), "other-token", time.Minute).Err()) + require.NoError(t, refresher.Refresh(ctx, 10*time.Minute)) + ttl = mr.TTL(queue.lockKey(batchID)) + require.LessOrEqual(t, ttl, time.Minute) +} + +func newBatchImageQueueTest(t *testing.T) (*batchImageQueue, *miniredis.Miniredis) { + t.Helper() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { + _ = rdb.Close() + }) + queue := newBatchImageQueueWithOptions(rdb, batchImageQueueOptions{ + InflightTTL: time.Hour, + LockTTL: time.Minute, + }) + return queue, mr +} diff --git a/backend/internal/repository/batch_image_repo.go b/backend/internal/repository/batch_image_repo.go new file mode 100644 index 0000000000..6531f75fd5 --- /dev/null +++ b/backend/internal/repository/batch_image_repo.go @@ -0,0 +1,986 @@ +package repository + +import ( + "context" + "database/sql" + "encoding/json" + "strconv" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +type batchImageSQLExecutor interface { + ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) + QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) + QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row +} + +type batchImageRepository struct { + db *sql.DB + sql batchImageSQLExecutor +} + +func NewBatchImageRepository(db *sql.DB) service.BatchImageRepository { + return &batchImageRepository{db: db, sql: db} +} + +func (r *batchImageRepository) CreateBatchImageJob(ctx context.Context, params service.CreateBatchImageJobParams) (*service.BatchImageJob, error) { + if !service.IsSupportedBatchImageProvider(params.Provider) { + return nil, service.ErrBatchImageInvalidProvider + } + if params.BatchID == "" { + batchID, err := service.NewBatchImageID() + if err != nil { + return nil, err + } + params.BatchID = batchID + } + if params.Status == "" { + params.Status = service.BatchImageJobStatusCreated + } + if params.Currency == "" { + params.Currency = "USD" + } + + job, err := createBatchImageJobWithSQL(ctx, r.sql, params) + if err != nil { + return nil, translatePersistenceError(err, nil, service.ErrBatchImageJobExists) + } + return job, nil +} + +func (r *batchImageRepository) GetBatchImageJobByBatchID(ctx context.Context, batchID string) (*service.BatchImageJob, error) { + job, err := scanBatchImageJob(r.sql.QueryRowContext(ctx, batchImageJobSelectSQL+" WHERE batch_id = $1", batchID)) + if err != nil { + return nil, translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + return job, nil +} + +func (r *batchImageRepository) GetBatchImageJobByIdempotencyKey(ctx context.Context, userID, apiKeyID int64, key string) (*service.BatchImageJob, error) { + job, err := scanBatchImageJob(r.sql.QueryRowContext(ctx, batchImageJobSelectSQL+` + WHERE user_id = $1 AND api_key_id = $2 AND idempotency_key = $3 + ORDER BY id DESC LIMIT 1`, userID, apiKeyID, key)) + if err != nil { + return nil, translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + return job, nil +} + +func (r *batchImageRepository) GetBatchImageJobByBatchIDForOwner(ctx context.Context, userID, apiKeyID int64, batchID string) (*service.BatchImageJob, error) { + job, err := scanBatchImageJob(r.sql.QueryRowContext(ctx, batchImageJobSelectSQL+` + WHERE batch_id = $1 AND user_id = $2 AND api_key_id = $3 AND user_deleted_at IS NULL`, batchID, userID, apiKeyID)) + if err != nil { + return nil, translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + return job, nil +} + +func (r *batchImageRepository) ListBatchImageJobsForOwner(ctx context.Context, userID, apiKeyID int64, filter service.BatchImageJobFilter) ([]*service.BatchImageJob, error) { + limit := filter.Limit + if limit <= 0 || limit > 100 { + limit = 20 + } + if filter.Offset < 0 { + filter.Offset = 0 + } + + query := batchImageJobSelectSQL + " WHERE user_id = $1 AND api_key_id = $2" + args := []any{userID, apiKeyID} + if filter.ExcludeDeleted { + query += " AND user_deleted_at IS NULL" + } + if filter.Status != "" { + query += " AND status = $" + strconv.Itoa(len(args)+1) + args = append(args, filter.Status) + } + if filter.TaskNameLike != "" { + query += " AND task_name ILIKE $" + strconv.Itoa(len(args)+1) + args = append(args, "%"+filter.TaskNameLike+"%") + } + if filter.Downloaded != nil { + if *filter.Downloaded { + query += " AND downloaded_at IS NOT NULL" + } else { + query += " AND downloaded_at IS NULL" + } + } + if filter.CreatedAfter != nil { + query += " AND created_at >= $" + strconv.Itoa(len(args)+1) + args = append(args, *filter.CreatedAfter) + } + if filter.CreatedBefore != nil { + query += " AND created_at < $" + strconv.Itoa(len(args)+1) + args = append(args, *filter.CreatedBefore) + } + query += " ORDER BY created_at DESC, id DESC LIMIT $" + strconv.Itoa(len(args)+1) + " OFFSET $" + strconv.Itoa(len(args)+2) + args = append(args, limit, filter.Offset) + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + return scanBatchImageJobs(rows) +} + +func (r *batchImageRepository) GetBatchImageJobByID(ctx context.Context, id int64) (*service.BatchImageJob, error) { + job, err := scanBatchImageJob(r.sql.QueryRowContext(ctx, batchImageJobSelectSQL+" WHERE id = $1", id)) + if err != nil { + return nil, translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + return job, nil +} + +func (r *batchImageRepository) TransitionBatchImageJobStatus(ctx context.Context, batchID, toStatus string, opts service.BatchImageTransitionOptions) error { + if r.db == nil { + return r.transitionBatchImageJobStatusWithSQL(ctx, r.sql, batchID, toStatus, opts) + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { + _ = tx.Rollback() + }() + + if err := r.transitionBatchImageJobStatusWithSQL(ctx, tx, batchID, toStatus, opts); err != nil { + return err + } + return tx.Commit() +} + +func (r *batchImageRepository) TouchBatchImageJobSubmitting(ctx context.Context, batchID string) error { + _, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET updated_at = $2 +WHERE batch_id = $1 + AND status IN ('created', 'uploading')`, batchID, time.Now()) + return err +} + +func (r *batchImageRepository) FailStaleUnsubmittedBatchImageJob(ctx context.Context, batchID string, cutoff time.Time, code, message string) (bool, error) { + now := time.Now() + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET status = 'failed', + last_error_code = $2, + last_error_message = $3, + finished_at = CASE WHEN finished_at IS NULL THEN $4 ELSE finished_at END, + updated_at = $4, + version = version + 1 +WHERE batch_id = $1 + AND status IN ('created', 'uploading') + AND provider_job_name IS NULL + AND updated_at <= $5`, batchID, code, message, now, cutoff) + if err != nil { + return false, err + } + affected, err := res.RowsAffected() + if err != nil { + return false, err + } + if affected == 0 { + return false, nil + } + return true, appendBatchImageEventWithSQL(ctx, r.sql, batchID, "billing_hold_recovery_failed_unsubmitted", map[string]any{ + "batch_id": batchID, + "error_code": code, + }) +} + +func (r *batchImageRepository) UpdateBatchImageJobProviderOutputRef(ctx context.Context, batchID, providerOutputRef string) error { + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET provider_output_ref = $2, updated_at = $3 +WHERE batch_id = $1`, batchID, providerOutputRef, time.Now()) + if err != nil { + return translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + if affected, err := res.RowsAffected(); err == nil && affected == 0 { + return service.ErrBatchImageJobNotFound + } + return nil +} + +func (r *batchImageRepository) UpdateBatchImageJobProviderSubmit(ctx context.Context, params service.UpdateBatchImageJobProviderSubmitParams) error { + if r.db == nil { + return r.updateBatchImageJobProviderSubmitWithSQL(ctx, r.sql, params) + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { + _ = tx.Rollback() + }() + if err := r.updateBatchImageJobProviderSubmitWithSQL(ctx, tx, params); err != nil { + return err + } + return tx.Commit() +} + +func (r *batchImageRepository) updateBatchImageJobProviderSubmitWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, params service.UpdateBatchImageJobProviderSubmitParams) error { + var current string + if err := sqlq.QueryRowContext(ctx, `SELECT status FROM batch_image_jobs WHERE batch_id = $1 FOR UPDATE`, params.BatchID).Scan(¤t); err != nil { + return translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + if !service.CanTransitionBatchImageJob(current, service.BatchImageJobStatusSubmitted) { + return service.ErrBatchImageInvalidTransition + } + now := time.Now() + if _, err := sqlq.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET status = 'submitted', + provider_job_name = $2, + provider_input_ref = NULLIF($3, ''), + provider_output_ref = NULLIF($4, ''), + gcs_input_uri = NULLIF($5, ''), + gcs_output_uri = NULLIF($6, ''), + submitted_at = CASE WHEN submitted_at IS NULL THEN $7 ELSE submitted_at END, + updated_at = $7, + version = version + 1 +WHERE batch_id = $1`, params.BatchID, params.ProviderJobName, params.ProviderInputRef, params.ProviderOutputRef, params.GCSInputURI, params.GCSOutputURI, now); err != nil { + return err + } + return appendBatchImageEventWithSQL(ctx, sqlq, params.BatchID, "provider_submitted", params.EventPayload) +} + +func (r *batchImageRepository) RecordBatchImageJobSubmitFailure(ctx context.Context, batchID, code, message string, markFailed bool) error { + now := time.Now() + statusSQL := "status" + if markFailed { + statusSQL = "'failed'" + } + _, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET status = `+statusSQL+`, + last_error_code = $2, + last_error_message = $3, + finished_at = CASE WHEN `+statusSQL+` = 'failed' AND finished_at IS NULL THEN $4 ELSE finished_at END, + updated_at = $4, + version = version + 1 +WHERE batch_id = $1`, batchID, code, message, now) + if err != nil { + return err + } + eventType := "submit_failed" + if !markFailed { + eventType = "queue_failed" + } + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, eventType, map[string]any{"error_code": code}) +} + +func (r *batchImageRepository) MarkBatchImageJobSettled(ctx context.Context, params service.MarkBatchImageJobSettledParams) error { + if r.db == nil { + return r.markBatchImageJobSettledWithSQL(ctx, r.sql, params) + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { + _ = tx.Rollback() + }() + + if err := r.markBatchImageJobSettledWithSQL(ctx, tx, params); err != nil { + return err + } + return tx.Commit() +} + +func (r *batchImageRepository) markBatchImageJobSettledWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, params service.MarkBatchImageJobSettledParams) error { + now := time.Now() + if params.Now != nil { + now = *params.Now + } + outputExpiresAt := params.OutputExpiresAt + + res, err := sqlq.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET status = 'completed', + actual_cost = $2, + manifest_hash = $3, + settled_at = CASE WHEN settled_at IS NULL THEN $4 ELSE settled_at END, + finished_at = CASE WHEN finished_at IS NULL THEN $4 ELSE finished_at END, + output_expires_at = CASE WHEN output_expires_at IS NULL THEN $5 ELSE output_expires_at END, + updated_at = $4, + version = version + 1 +WHERE batch_id = $1 + AND status = 'settling' + AND (manifest_hash IS NULL OR manifest_hash = '' OR manifest_hash = $3)`, params.BatchID, params.ActualCost, params.ManifestHash, now, outputExpiresAt) + if err != nil { + return err + } + affected, err := res.RowsAffected() + if err != nil { + return err + } + if affected == 0 { + job, getErr := scanBatchImageJob(sqlq.QueryRowContext(ctx, batchImageJobSelectSQL+" WHERE batch_id = $1", params.BatchID)) + if getErr != nil { + return translatePersistenceError(getErr, service.ErrBatchImageJobNotFound, nil) + } + if job.Status != service.BatchImageJobStatusSettling { + if job.Status == service.BatchImageJobStatusCompleted { + return service.ErrBatchImageAlreadySettled + } + return service.ErrBatchImageSettlementInvalidStatus + } + return service.ErrBatchImageSettlementManifestConflict + } + return appendBatchImageEventWithSQL(ctx, sqlq, params.BatchID, "settlement_completed", params.EventPayload) +} + +func (r *batchImageRepository) SetBatchImageJobSettlementFailed(ctx context.Context, batchID, code, message string) (int, error) { + var retryCount int + err := r.sql.QueryRowContext(ctx, ` +UPDATE batch_image_jobs +SET last_error_code = $2, + last_error_message = $3, + retry_count = retry_count + 1, + updated_at = $4 +WHERE batch_id = $1 +RETURNING retry_count`, batchID, code, message, time.Now()).Scan(&retryCount) + if err != nil { + return 0, translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + return retryCount, appendBatchImageEventWithSQL(ctx, r.sql, batchID, "settlement_failed", map[string]any{ + "error_code": code, + }) +} + +func (r *batchImageRepository) transitionBatchImageJobStatusWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, batchID, toStatus string, opts service.BatchImageTransitionOptions) error { + var current string + if err := sqlq.QueryRowContext(ctx, `SELECT status FROM batch_image_jobs WHERE batch_id = $1 FOR UPDATE`, batchID).Scan(¤t); err != nil { + return translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + if !service.CanTransitionBatchImageJob(current, toStatus) { + return service.ErrBatchImageInvalidTransition + } + + now := time.Now() + if opts.Now != nil { + now = *opts.Now + } + + if _, err := sqlq.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET + status = $2::varchar, + version = version + 1, + updated_at = $3, + last_error_code = CASE WHEN $2::varchar = 'failed' THEN $4 ELSE last_error_code END, + last_error_message = CASE WHEN $2::varchar = 'failed' THEN $5 ELSE last_error_message END, + submitted_at = CASE WHEN $2::varchar = 'submitted' AND submitted_at IS NULL THEN $3 ELSE submitted_at END, + started_at = CASE WHEN $2::varchar = 'running' AND started_at IS NULL THEN $3 ELSE started_at END, + finished_at = CASE WHEN $2::varchar IN ('completed', 'failed', 'cancelled') AND finished_at IS NULL THEN $3 ELSE finished_at END, + settled_at = CASE WHEN $2::varchar = 'completed' AND settled_at IS NULL THEN $3 ELSE settled_at END, + output_deleted_at = CASE WHEN $2::varchar = 'output_deleted' AND output_deleted_at IS NULL THEN $3 ELSE output_deleted_at END +WHERE batch_id = $1`, batchID, toStatus, now, opts.ErrorCode, opts.ErrorMessage); err != nil { + return err + } + + if opts.EventType != "" { + return appendBatchImageEventWithSQL(ctx, sqlq, batchID, opts.EventType, opts.EventPayload) + } + return nil +} + +func (r *batchImageRepository) CreateBatchImageItem(ctx context.Context, params service.CreateBatchImageItemParams) (*service.BatchImageItem, error) { + item, err := createBatchImageItemWithSQL(ctx, r.sql, params) + if err != nil { + return nil, translatePersistenceError(err, nil, service.ErrBatchImageItemExists) + } + return item, nil +} + +func (r *batchImageRepository) BulkCreateBatchImageItems(ctx context.Context, params []service.CreateBatchImageItemParams) error { + if len(params) == 0 { + return nil + } + if r.db == nil { + for _, param := range params { + if _, err := createBatchImageItemWithSQL(ctx, r.sql, param); err != nil { + return translatePersistenceError(err, nil, service.ErrBatchImageItemExists) + } + } + return nil + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { + _ = tx.Rollback() + }() + + for _, param := range params { + if _, err := createBatchImageItemWithSQL(ctx, tx, param); err != nil { + return translatePersistenceError(err, nil, service.ErrBatchImageItemExists) + } + } + return tx.Commit() +} + +func (r *batchImageRepository) ReplaceBatchImageItemsForJob(ctx context.Context, batchID string, items []service.CreateBatchImageItemParams, counts service.BatchImageCounts) error { + if r.db == nil { + return r.replaceBatchImageItemsForJobWithSQL(ctx, r.sql, batchID, items, counts) + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { + _ = tx.Rollback() + }() + + if err := r.replaceBatchImageItemsForJobWithSQL(ctx, tx, batchID, items, counts); err != nil { + return err + } + return tx.Commit() +} + +func (r *batchImageRepository) replaceBatchImageItemsForJobWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, batchID string, items []service.CreateBatchImageItemParams, counts service.BatchImageCounts) error { + var id int64 + var status string + if err := sqlq.QueryRowContext(ctx, `SELECT id, status FROM batch_image_jobs WHERE batch_id = $1 FOR UPDATE`, batchID).Scan(&id, &status); err != nil { + return translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + // 仅允许 indexing 状态重建 item 表:防止锁过期后掉队的 worker + // 重写已完成/已结算 job 的条目,造成账目与结果漂移。 + if status != service.BatchImageJobStatusIndexing { + return service.ErrBatchImageIndexStateConflict + } + promptPreviews, err := r.batchImageItemPromptPreviews(ctx, sqlq, batchID) + if err != nil { + return err + } + if _, err := sqlq.ExecContext(ctx, `DELETE FROM batch_image_items WHERE job_id = $1`, batchID); err != nil { + return err + } + for _, item := range items { + item.JobID = batchID + if item.PromptPreview == nil { + if preview := promptPreviews[item.CustomID]; preview != "" { + item.PromptPreview = &preview + } + } + if _, err := createBatchImageItemWithSQL(ctx, sqlq, item); err != nil { + return translatePersistenceError(err, nil, service.ErrBatchImageItemExists) + } + } + _, err = sqlq.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET success_count = $2, + fail_count = $3, + updated_at = $4 +WHERE batch_id = $1`, batchID, counts.SuccessCount, counts.FailCount, time.Now()) + return err +} + +func (r *batchImageRepository) batchImageItemPromptPreviews(ctx context.Context, sqlq batchImageSQLExecutor, batchID string) (map[string]string, error) { + rows, err := sqlq.QueryContext(ctx, `SELECT custom_id, prompt_preview FROM batch_image_items WHERE job_id = $1 AND prompt_preview IS NOT NULL`, batchID) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + out := make(map[string]string) + for rows.Next() { + var customID string + var preview sql.NullString + if err := rows.Scan(&customID, &preview); err != nil { + return nil, err + } + if preview.Valid && preview.String != "" { + out[customID] = preview.String + } + } + return out, rows.Err() +} + +func (r *batchImageRepository) ListBatchImageItems(ctx context.Context, batchID string, filter service.BatchImageItemFilter) ([]*service.BatchImageItem, error) { + limit := filter.Limit + if limit <= 0 || limit > 500 { + limit = 100 + } + if filter.Offset < 0 { + filter.Offset = 0 + } + + query := batchImageItemSelectSQL + " WHERE job_id = $1" + args := []any{batchID} + if filter.Status != "" { + query += " AND status = $2" + args = append(args, filter.Status) + } + query += " ORDER BY id ASC LIMIT $" + strconv.Itoa(len(args)+1) + " OFFSET $" + strconv.Itoa(len(args)+2) + args = append(args, limit, filter.Offset) + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + var items []*service.BatchImageItem + for rows.Next() { + item, err := scanBatchImageItem(rows) + if err != nil { + return nil, err + } + items = append(items, item) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +func (r *batchImageRepository) ListBatchImageItemsForOwner(ctx context.Context, userID, apiKeyID int64, batchID string, filter service.BatchImageItemFilter) ([]*service.BatchImageItem, error) { + if _, err := r.GetBatchImageJobByBatchIDForOwner(ctx, userID, apiKeyID, batchID); err != nil { + return nil, err + } + return r.ListBatchImageItems(ctx, batchID, filter) +} + +func (r *batchImageRepository) GetBatchImageJobForDownload(ctx context.Context, userID, apiKeyID int64, batchID string) (*service.BatchImageJob, error) { + return r.GetBatchImageJobByBatchIDForOwner(ctx, userID, apiKeyID, batchID) +} + +func (r *batchImageRepository) GetBatchImageItemForDownload(ctx context.Context, batchID, customID string) (*service.BatchImageItem, error) { + item, err := scanBatchImageItem(r.sql.QueryRowContext(ctx, batchImageItemSelectSQL+` + WHERE job_id = $1 AND custom_id = $2`, batchID, customID)) + if err != nil { + return nil, translatePersistenceError(err, service.ErrBatchImageItemNotFound, nil) + } + return item, nil +} + +func (r *batchImageRepository) ListBatchImageItemsForDownload(ctx context.Context, batchID string, status string, limit int) ([]*service.BatchImageItem, error) { + return r.ListBatchImageItems(ctx, batchID, service.BatchImageItemFilter{Status: status, Limit: limit}) +} + +func (r *batchImageRepository) ListBatchImageJobsDueForInputCleanup(ctx context.Context, cutoff time.Time, limit int) ([]*service.BatchImageJob, error) { + if limit <= 0 || limit > 1000 { + limit = 100 + } + rows, err := r.sql.QueryContext(ctx, batchImageJobSelectSQL+` + WHERE input_deleted_at IS NULL + AND provider_input_ref IS NOT NULL + AND status IN ('completed', 'failed', 'cancelled', 'output_deleted') + AND COALESCE(finished_at, settled_at, updated_at, created_at) <= $1 + ORDER BY id ASC + LIMIT $2`, cutoff, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + return scanBatchImageJobs(rows) +} + +func (r *batchImageRepository) ListBatchImageJobsDueForOutputCleanup(ctx context.Context, now time.Time, limit int) ([]*service.BatchImageJob, error) { + if limit <= 0 || limit > 1000 { + limit = 100 + } + rows, err := r.sql.QueryContext(ctx, batchImageJobSelectSQL+` + WHERE output_deleted_at IS NULL + AND provider_output_ref IS NOT NULL + AND status = 'completed' + AND output_expires_at IS NOT NULL + AND output_expires_at <= $1 + ORDER BY output_expires_at ASC, id ASC + LIMIT $2`, now, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + return scanBatchImageJobs(rows) +} + +func (r *batchImageRepository) ListStaleUnsubmittedBatchImageJobs(ctx context.Context, cutoff time.Time, limit int) ([]*service.BatchImageJob, error) { + if limit <= 0 || limit > 1000 { + limit = 100 + } + rows, err := r.sql.QueryContext(ctx, batchImageJobSelectSQL+` + WHERE status IN ('created', 'uploading') + AND provider_job_name IS NULL + AND COALESCE(hold_amount, estimated_cost, 0) > 0 + AND updated_at <= $1 + ORDER BY updated_at ASC, id ASC + LIMIT $2`, cutoff, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + return scanBatchImageJobs(rows) +} + +func (r *batchImageRepository) MarkBatchImageInputDeleted(ctx context.Context, batchID string, deletedAt time.Time) error { + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET input_deleted_at = CASE WHEN input_deleted_at IS NULL THEN $2 ELSE input_deleted_at END, + updated_at = $2, + version = version + 1 +WHERE batch_id = $1`, batchID, deletedAt) + if err != nil { + return err + } + if affected, err := res.RowsAffected(); err == nil && affected == 0 { + return service.ErrBatchImageJobNotFound + } + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, "input_cleanup_completed", map[string]any{ + "batch_id": batchID, + "cleanup_target": "input", + "deleted_at": deletedAt.UTC().Format(time.RFC3339), + }) +} + +func (r *batchImageRepository) MarkBatchImageOutputDeleted(ctx context.Context, batchID string, deletedAt time.Time) error { + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET status = CASE WHEN status = 'completed' THEN 'output_deleted' ELSE status END, + output_deleted_at = CASE WHEN output_deleted_at IS NULL THEN $2 ELSE output_deleted_at END, + finished_at = CASE WHEN status = 'completed' AND finished_at IS NULL THEN $2 ELSE finished_at END, + updated_at = $2, + version = version + 1 +WHERE batch_id = $1`, batchID, deletedAt) + if err != nil { + return err + } + if affected, err := res.RowsAffected(); err == nil && affected == 0 { + return service.ErrBatchImageJobNotFound + } + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, "output_cleanup_completed", map[string]any{ + "batch_id": batchID, + "cleanup_target": "output", + "deleted_at": deletedAt.UTC().Format(time.RFC3339), + }) +} + +func (r *batchImageRepository) MarkBatchImageDownloaded(ctx context.Context, batchID string, downloadedAt time.Time) error { + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET downloaded_at = CASE WHEN downloaded_at IS NULL THEN $2 ELSE downloaded_at END, + updated_at = $2 +WHERE batch_id = $1`, batchID, downloadedAt) + if err != nil { + return err + } + if affected, err := res.RowsAffected(); err == nil && affected == 0 { + return service.ErrBatchImageJobNotFound + } + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, "download_completed", map[string]any{ + "batch_id": batchID, + "downloaded_at": downloadedAt.UTC().Format(time.RFC3339), + }) +} + +func (r *batchImageRepository) MarkBatchImageJobUserDeleted(ctx context.Context, userID, apiKeyID int64, batchID string, deletedAt time.Time) error { + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET user_deleted_at = CASE WHEN user_deleted_at IS NULL THEN $4 ELSE user_deleted_at END, + updated_at = $4 +WHERE batch_id = $1 + AND user_id = $2 + AND api_key_id = $3 + AND user_deleted_at IS NULL + AND status IN ('completed', 'failed', 'cancelled', 'output_deleted')`, batchID, userID, apiKeyID, deletedAt) + if err != nil { + return err + } + if affected, err := res.RowsAffected(); err == nil && affected == 0 { + return service.ErrBatchImageRecordDeleteNotReady + } + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, "user_record_deleted", map[string]any{ + "batch_id": batchID, + "deleted_at": deletedAt.UTC().Format(time.RFC3339), + "user_id": userID, + "api_key_id": apiKeyID, + }) +} + +func (r *batchImageRepository) SetBatchImageOutputExpiresAt(ctx context.Context, batchID string, expiresAt time.Time) error { + res, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET output_expires_at = CASE WHEN output_expires_at IS NULL THEN $2 ELSE output_expires_at END, + updated_at = $3 +WHERE batch_id = $1`, batchID, expiresAt, time.Now()) + if err != nil { + return err + } + if affected, err := res.RowsAffected(); err == nil && affected == 0 { + return service.ErrBatchImageJobNotFound + } + return nil +} + +func (r *batchImageRepository) RecordBatchImageCleanupFailure(ctx context.Context, batchID, code, message string) error { + _, err := r.sql.ExecContext(ctx, ` +UPDATE batch_image_jobs +SET last_error_code = $2, + last_error_message = $3, + retry_count = retry_count + 1, + updated_at = $4 +WHERE batch_id = $1`, batchID, code, message, time.Now()) + if err != nil { + return err + } + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, "output_cleanup_failed", map[string]any{"error_code": code}) +} + +func (r *batchImageRepository) AppendBatchImageEvent(ctx context.Context, batchID, eventType string, payload any) error { + return appendBatchImageEventWithSQL(ctx, r.sql, batchID, eventType, payload) +} + +func createBatchImageJobWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, params service.CreateBatchImageJobParams) (*service.BatchImageJob, error) { + return scanBatchImageJob(sqlq.QueryRowContext(ctx, ` +INSERT INTO batch_image_jobs ( + batch_id, user_id, api_key_id, account_id, provider, model, task_name, parent_batch_id, status, + provider_job_name, provider_input_ref, provider_output_ref, gcs_input_uri, gcs_output_uri, + item_count, success_count, fail_count, cancelled_count, + estimated_cost, hold_amount, actual_cost, + base_unit_price, group_rate_multiplier, account_rate_multiplier, + batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price, + pricing_snapshot_version, + currency, hold_id, + idempotency_key, request_hash, manifest_hash, retry_count, output_expires_at +) VALUES ( + $1, $2, $3, $4, $5, $6, $7, $8, $9, + $10, $11, $12, $13, $14, + $15, $16, $17, $18, + $19, $20, $21, + $22, $23, $24, + $25, $26, $27, $28, + $29, + $30, $31, + $32, $33, $34, $35, $36 +) +RETURNING `+batchImageJobColumns, + params.BatchID, params.UserID, params.APIKeyID, params.AccountID, params.Provider, params.Model, params.TaskName, params.ParentBatchID, params.Status, + params.ProviderJobName, params.ProviderInputRef, params.ProviderOutputRef, params.GCSInputURI, params.GCSOutputURI, + params.ItemCount, params.SuccessCount, params.FailCount, params.CancelledCount, + params.EstimatedCost, params.HoldAmount, params.ActualCost, + params.BaseUnitPrice, params.GroupRateMultiplier, params.AccountRateMultiplier, + params.BatchDiscountMultiplier, params.HoldMultiplier, params.BillableUnitPrice, params.HoldUnitPrice, + params.PricingSnapshotVersion, + params.Currency, params.HoldID, + params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.OutputExpiresAt, + )) +} + +func createBatchImageItemWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, params service.CreateBatchImageItemParams) (*service.BatchImageItem, error) { + return scanBatchImageItem(sqlq.QueryRowContext(ctx, ` +INSERT INTO batch_image_items ( + job_id, custom_id, status, request_hash, prompt_preview, provider_source_object, + source_line_number, source_byte_offset, source_byte_length, + mime_type, file_extension, image_count, + error_code, error_message, billed_amount, indexed_at +) VALUES ( + $1, $2, $3, $4, $5, $6, + $7, $8, $9, + $10, $11, $12, + $13, $14, $15, $16 +) +RETURNING `+batchImageItemColumns, + params.JobID, params.CustomID, params.Status, params.RequestHash, params.PromptPreview, params.ProviderSourceObject, + params.SourceLineNumber, params.SourceByteOffset, params.SourceByteLength, + params.MimeType, params.FileExtension, params.ImageCount, + params.ErrorCode, params.ErrorMessage, params.BilledAmount, params.IndexedAt, + )) +} + +func appendBatchImageEventWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, batchID, eventType string, payload any) error { + var payloadArg any + if payload != nil { + payloadBytes, err := json.Marshal(payload) + if err != nil { + return err + } + payloadArg = string(payloadBytes) + } + _, err := sqlq.ExecContext(ctx, ` +INSERT INTO batch_image_events (job_id, event_type, payload) +VALUES ($1, $2, $3)`, batchID, eventType, payloadArg) + return err +} + +type rowScanner interface { + Scan(dest ...any) error +} + +const batchImageJobColumns = ` +id, batch_id, user_id, api_key_id, account_id, provider, model, task_name, parent_batch_id, status, +provider_job_name, provider_input_ref, provider_output_ref, gcs_input_uri, gcs_output_uri, +item_count, success_count, fail_count, cancelled_count, +estimated_cost, hold_amount, actual_cost, +base_unit_price, group_rate_multiplier, account_rate_multiplier, +batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price, +pricing_snapshot_version, +currency, hold_id, +idempotency_key, request_hash, manifest_hash, +retry_count, version, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at, +last_error_code, last_error_message, +created_at, updated_at, submitted_at, started_at, finished_at, settled_at` + +const batchImageJobSelectSQL = `SELECT ` + batchImageJobColumns + ` FROM batch_image_jobs` + +func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) { + var job service.BatchImageJob + var apiKeyID, accountID sql.NullInt64 + var providerJobName, providerInputRef, providerOutputRef, gcsInputURI, gcsOutputURI sql.NullString + var parentBatchID sql.NullString + var holdAmount, actualCost sql.NullFloat64 + var holdID, idempotencyKey, requestHash, manifestHash sql.NullString + var outputExpiresAt, inputDeletedAt, outputDeletedAt, downloadedAt, userDeletedAt sql.NullTime + var lastErrorCode, lastErrorMessage sql.NullString + var submittedAt, startedAt, finishedAt, settledAt sql.NullTime + + err := row.Scan( + &job.ID, &job.BatchID, &job.UserID, &apiKeyID, &accountID, &job.Provider, &job.Model, &job.TaskName, &parentBatchID, &job.Status, + &providerJobName, &providerInputRef, &providerOutputRef, &gcsInputURI, &gcsOutputURI, + &job.ItemCount, &job.SuccessCount, &job.FailCount, &job.CancelledCount, + &job.EstimatedCost, &holdAmount, &actualCost, + &job.BaseUnitPrice, &job.GroupRateMultiplier, &job.AccountRateMultiplier, + &job.BatchDiscountMultiplier, &job.HoldMultiplier, &job.BillableUnitPrice, &job.HoldUnitPrice, + &job.PricingSnapshotVersion, + &job.Currency, &holdID, + &idempotencyKey, &requestHash, &manifestHash, + &job.RetryCount, &job.Version, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt, + &lastErrorCode, &lastErrorMessage, + &job.CreatedAt, &job.UpdatedAt, &submittedAt, &startedAt, &finishedAt, &settledAt, + ) + if err != nil { + return nil, err + } + + job.APIKeyID = batchImageNullInt64Ptr(apiKeyID) + job.AccountID = batchImageNullInt64Ptr(accountID) + job.ProviderJobName = batchImageNullStringPtr(providerJobName) + job.ProviderInputRef = batchImageNullStringPtr(providerInputRef) + job.ProviderOutputRef = batchImageNullStringPtr(providerOutputRef) + job.ParentBatchID = batchImageNullStringPtr(parentBatchID) + job.GCSInputURI = batchImageNullStringPtr(gcsInputURI) + job.GCSOutputURI = batchImageNullStringPtr(gcsOutputURI) + job.HoldAmount = batchImageNullFloat64Ptr(holdAmount) + job.ActualCost = batchImageNullFloat64Ptr(actualCost) + job.HoldID = batchImageNullStringPtr(holdID) + job.IdempotencyKey = batchImageNullStringPtr(idempotencyKey) + job.RequestHash = batchImageNullStringPtr(requestHash) + job.ManifestHash = batchImageNullStringPtr(manifestHash) + job.OutputExpiresAt = batchImageNullTimePtr(outputExpiresAt) + job.InputDeletedAt = batchImageNullTimePtr(inputDeletedAt) + job.OutputDeletedAt = batchImageNullTimePtr(outputDeletedAt) + job.DownloadedAt = batchImageNullTimePtr(downloadedAt) + job.UserDeletedAt = batchImageNullTimePtr(userDeletedAt) + job.LastErrorCode = batchImageNullStringPtr(lastErrorCode) + job.LastErrorMessage = batchImageNullStringPtr(lastErrorMessage) + job.SubmittedAt = batchImageNullTimePtr(submittedAt) + job.StartedAt = batchImageNullTimePtr(startedAt) + job.FinishedAt = batchImageNullTimePtr(finishedAt) + job.SettledAt = batchImageNullTimePtr(settledAt) + return &job, nil +} + +func scanBatchImageJobs(rows *sql.Rows) ([]*service.BatchImageJob, error) { + var jobs []*service.BatchImageJob + for rows.Next() { + job, err := scanBatchImageJob(rows) + if err != nil { + return nil, err + } + jobs = append(jobs, job) + } + if err := rows.Err(); err != nil { + return nil, err + } + return jobs, nil +} + +const batchImageItemColumns = ` +id, job_id, custom_id, status, request_hash, prompt_preview, provider_source_object, +source_line_number, source_byte_offset, source_byte_length, +mime_type, file_extension, image_count, +error_code, error_message, billed_amount, +created_at, indexed_at` + +const batchImageItemSelectSQL = `SELECT ` + batchImageItemColumns + ` FROM batch_image_items` + +func scanBatchImageItem(row rowScanner) (*service.BatchImageItem, error) { + var item service.BatchImageItem + var requestHash, promptPreview, providerSourceObject sql.NullString + var sourceLineNumber sql.NullInt64 + var sourceByteOffset, sourceByteLength sql.NullInt64 + var mimeType, fileExtension, errorCode, errorMessage sql.NullString + var billedAmount sql.NullFloat64 + var indexedAt sql.NullTime + + err := row.Scan( + &item.ID, &item.JobID, &item.CustomID, &item.Status, &requestHash, &promptPreview, &providerSourceObject, + &sourceLineNumber, &sourceByteOffset, &sourceByteLength, + &mimeType, &fileExtension, &item.ImageCount, + &errorCode, &errorMessage, &billedAmount, + &item.CreatedAt, &indexedAt, + ) + if err != nil { + return nil, err + } + + item.RequestHash = batchImageNullStringPtr(requestHash) + item.PromptPreview = batchImageNullStringPtr(promptPreview) + item.ProviderSourceObject = batchImageNullStringPtr(providerSourceObject) + item.SourceLineNumber = batchImageNullIntPtr(sourceLineNumber) + item.SourceByteOffset = batchImageNullInt64Ptr(sourceByteOffset) + item.SourceByteLength = batchImageNullInt64Ptr(sourceByteLength) + item.MimeType = batchImageNullStringPtr(mimeType) + item.FileExtension = batchImageNullStringPtr(fileExtension) + item.ErrorCode = batchImageNullStringPtr(errorCode) + item.ErrorMessage = batchImageNullStringPtr(errorMessage) + item.BilledAmount = batchImageNullFloat64Ptr(billedAmount) + item.IndexedAt = batchImageNullTimePtr(indexedAt) + return &item, nil +} + +func batchImageNullStringPtr(v sql.NullString) *string { + if !v.Valid { + return nil + } + return &v.String +} + +func batchImageNullInt64Ptr(v sql.NullInt64) *int64 { + if !v.Valid { + return nil + } + return &v.Int64 +} + +func batchImageNullIntPtr(v sql.NullInt64) *int { + if !v.Valid { + return nil + } + i := int(v.Int64) + return &i +} + +func batchImageNullFloat64Ptr(v sql.NullFloat64) *float64 { + if !v.Valid { + return nil + } + return &v.Float64 +} + +func batchImageNullTimePtr(v sql.NullTime) *time.Time { + if !v.Valid { + return nil + } + return &v.Time +} + +var _ service.BatchImageRepository = (*batchImageRepository)(nil) diff --git a/backend/internal/repository/batch_image_repo_integration_test.go b/backend/internal/repository/batch_image_repo_integration_test.go new file mode 100644 index 0000000000..df6be4a954 --- /dev/null +++ b/backend/internal/repository/batch_image_repo_integration_test.go @@ -0,0 +1,381 @@ +//go:build integration + +package repository + +import ( + "context" + "crypto/sha1" + "encoding/hex" + "errors" + "regexp" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func newBatchImageRepositoryWithSQL(sqlq batchImageSQLExecutor) *batchImageRepository { + return &batchImageRepository{sql: sqlq} +} + +func TestBatchImageRepository_CreateJobAndDuplicates(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "create") + + job, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + ItemCount: 2, + EstimatedCost: 0.02, + }) + require.NoError(t, err) + require.Equal(t, batchID, job.BatchID) + require.Equal(t, service.BatchImageJobStatusCreated, job.Status) + require.Equal(t, "USD", job.Currency) + + _, err = repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + ItemCount: 1, + }) + require.Error(t, err) + require.True(t, errors.Is(err, service.ErrBatchImageJobExists)) +} + +func TestBatchImageRepository_InvalidProvider(t *testing.T) { + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + + _, err := repo.CreateBatchImageJob(context.Background(), service.CreateBatchImageJobParams{ + BatchID: batchImageTestID(t, "provider"), + UserID: 1001, + Provider: "unknown", + Model: "gemini-2.5-flash-image", + ItemCount: 1, + }) + require.Error(t, err) + require.True(t, errors.Is(err, service.ErrBatchImageInvalidProvider)) +} + +func TestBatchImageRepository_TransitionIncrementsVersionAndEvents(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "transition") + now := time.Date(2026, 7, 3, 8, 0, 0, 0, time.UTC) + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderVertex, + Model: "gemini-2.5-flash-image", + ItemCount: 1, + }) + require.NoError(t, err) + + err = repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusUploading, service.BatchImageTransitionOptions{ + EventType: "status_changed", + EventPayload: map[string]any{"to": service.BatchImageJobStatusUploading}, + Now: &now, + }) + require.NoError(t, err) + + job, err := repo.GetBatchImageJobByBatchID(ctx, batchID) + require.NoError(t, err) + require.Equal(t, service.BatchImageJobStatusUploading, job.Status) + require.Equal(t, 1, job.Version) + + var eventCount int + err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM batch_image_events WHERE job_id = $1 AND event_type = 'status_changed'`, batchID).Scan(&eventCount) + require.NoError(t, err) + require.Equal(t, 1, eventCount) +} + +func TestBatchImageRepository_InvalidTransition(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "invalid-transition") + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + ItemCount: 1, + }) + require.NoError(t, err) + + err = repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusRunning, service.BatchImageTransitionOptions{}) + require.Error(t, err) + require.True(t, errors.Is(err, service.ErrBatchImageInvalidTransition)) +} + +func TestBatchImageRepository_TerminalStatusCannotMoveBack(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "terminal") + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: service.BatchImageJobStatusCompleted, + ItemCount: 1, + }) + require.NoError(t, err) + + err = repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusRunning, service.BatchImageTransitionOptions{}) + require.Error(t, err) + require.True(t, errors.Is(err, service.ErrBatchImageInvalidTransition)) +} + +func TestBatchImageRepository_ItemCustomIDUniqueness(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + firstBatchID := batchImageTestID(t, "items-a") + secondBatchID := batchImageTestID(t, "items-b") + + for _, batchID := range []string{firstBatchID, secondBatchID} { + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + ItemCount: 1, + }) + require.NoError(t, err) + } + + _, err := repo.CreateBatchImageItem(ctx, service.CreateBatchImageItemParams{ + JobID: firstBatchID, + CustomID: "line-1", + Status: service.BatchImageItemStatusSuccess, + ImageCount: 1, + }) + require.NoError(t, err) + + _, err = tx.ExecContext(ctx, `SAVEPOINT batch_image_duplicate_item`) + require.NoError(t, err) + _, err = repo.CreateBatchImageItem(ctx, service.CreateBatchImageItemParams{ + JobID: firstBatchID, + CustomID: "line-1", + Status: service.BatchImageItemStatusFailed, + }) + require.Error(t, err) + require.True(t, errors.Is(err, service.ErrBatchImageItemExists)) + _, rollbackErr := tx.ExecContext(ctx, `ROLLBACK TO SAVEPOINT batch_image_duplicate_item`) + require.NoError(t, rollbackErr) + + _, err = repo.CreateBatchImageItem(ctx, service.CreateBatchImageItemParams{ + JobID: secondBatchID, + CustomID: "line-1", + Status: service.BatchImageItemStatusSuccess, + ImageCount: 1, + }) + require.NoError(t, err) + + items, err := repo.ListBatchImageItems(ctx, firstBatchID, service.BatchImageItemFilter{}) + require.NoError(t, err) + require.Len(t, items, 1) +} + +func TestBatchImageRepository_ReplaceBatchImageItemsForJob(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "replace-items") + lineOne := 1 + lineTwo := 2 + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + ItemCount: 2, + }) + require.NoError(t, err) + + // 非 indexing 状态不允许重建 item 表:防止锁过期后掉队的 worker + // 重写已完成/已结算 job 的条目。 + err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{ + {CustomID: "old", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1}, + }, service.BatchImageCounts{SuccessCount: 1}) + require.ErrorIs(t, err, service.ErrBatchImageIndexStateConflict) + + require.NoError(t, repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusSubmitted, service.BatchImageTransitionOptions{})) + require.NoError(t, repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusIndexing, service.BatchImageTransitionOptions{})) + + err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{ + {CustomID: "old", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1}, + }, service.BatchImageCounts{SuccessCount: 1}) + require.NoError(t, err) + + err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{ + {CustomID: "new-ok", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1}, + {CustomID: "new-fail", Status: service.BatchImageItemStatusFailed, SourceLineNumber: &lineTwo, ErrorCode: batchImageTestStringPtr("SAFETY_BLOCKED")}, + }, service.BatchImageCounts{SuccessCount: 1, FailCount: 1}) + require.NoError(t, err) + + items, err := repo.ListBatchImageItems(ctx, batchID, service.BatchImageItemFilter{}) + require.NoError(t, err) + require.Len(t, items, 2) + require.Equal(t, "new-ok", items[0].CustomID) + require.Equal(t, "new-fail", items[1].CustomID) + + job, err := repo.GetBatchImageJobByBatchID(ctx, batchID) + require.NoError(t, err) + require.Equal(t, 1, job.SuccessCount) + require.Equal(t, 1, job.FailCount) +} + +func TestBatchImageRepository_MarkBatchImageJobSettled(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "settled") + apiKeyID := int64(2001) + accountID := int64(3001) + providerJob := "providers/job" + outputRef := "files/output" + now := time.Date(2026, 7, 4, 10, 0, 0, 0, time.UTC) + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-image", + Status: service.BatchImageJobStatusSettling, + ProviderJobName: &providerJob, + ProviderOutputRef: &outputRef, + ItemCount: 3, + SuccessCount: 2, + FailCount: 1, + }) + require.NoError(t, err) + + err = repo.MarkBatchImageJobSettled(ctx, service.MarkBatchImageJobSettledParams{ + BatchID: batchID, + ActualCost: 0.5, + ManifestHash: "manifest-hash", + EventPayload: map[string]any{"request_id": "batch_image_settlement:" + batchID}, + Now: &now, + }) + require.NoError(t, err) + + job, err := repo.GetBatchImageJobByBatchID(ctx, batchID) + require.NoError(t, err) + require.Equal(t, service.BatchImageJobStatusCompleted, job.Status) + require.NotNil(t, job.ActualCost) + require.Equal(t, 0.5, *job.ActualCost) + require.Equal(t, "manifest-hash", batchImageDerefTest(job.ManifestHash)) + require.NotNil(t, job.SettledAt) + require.Equal(t, now, *job.SettledAt) + + var eventCount int + err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM batch_image_events WHERE job_id = $1 AND event_type = 'settlement_completed'`, batchID).Scan(&eventCount) + require.NoError(t, err) + require.Equal(t, 1, eventCount) +} + +func TestBatchImageRepository_SetBatchImageJobSettlementFailed(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "settlement-failed") + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderGeminiAPI, + Model: "gemini-image", + Status: service.BatchImageJobStatusSettling, + ItemCount: 1, + SuccessCount: 1, + }) + require.NoError(t, err) + + retryCount, err := repo.SetBatchImageJobSettlementFailed(ctx, batchID, "SETTLEMENT_BILLING_FAILED", "temporary") + require.NoError(t, err) + require.Equal(t, 1, retryCount) + + job, err := repo.GetBatchImageJobByBatchID(ctx, batchID) + require.NoError(t, err) + require.Equal(t, service.BatchImageJobStatusSettling, job.Status) + require.Equal(t, "SETTLEMENT_BILLING_FAILED", batchImageDerefTest(job.LastErrorCode)) + require.Equal(t, "temporary", batchImageDerefTest(job.LastErrorMessage)) + require.Equal(t, 1, job.RetryCount) +} + +func TestBatchImageRepository_AppendEvent(t *testing.T) { + ctx := context.Background() + tx := testTx(t) + repo := newBatchImageRepositoryWithSQL(tx) + batchID := batchImageTestID(t, "event") + + _, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: batchID, + UserID: 1001, + Provider: service.BatchImageProviderVertex, + Model: "gemini-2.5-flash-image", + ItemCount: 1, + }) + require.NoError(t, err) + + err = repo.AppendBatchImageEvent(ctx, batchID, "job_created", map[string]any{"batch_id": batchID}) + require.NoError(t, err) + + var payload string + err = tx.QueryRowContext(ctx, `SELECT payload::text FROM batch_image_events WHERE job_id = $1 AND event_type = 'job_created'`, batchID).Scan(&payload) + require.NoError(t, err) + require.Contains(t, payload, batchID) +} + +func batchImageTestID(t *testing.T, prefix string) string { + t.Helper() + safePrefix := batchImageSafeTestIDSegment(prefix, 20) + sum := sha1.Sum([]byte(t.Name())) + return "imgbatch_" + safePrefix + "_" + hex.EncodeToString(sum[:])[:16] +} + +func batchImageSafeTestIDSegment(v string, maxLen int) string { + v = strings.ToLower(strings.TrimSpace(v)) + v = regexp.MustCompile(`[^a-z0-9_-]+`).ReplaceAllString(v, "-") + v = strings.Trim(v, "-_") + if v == "" { + v = "job" + } + if len(v) > maxLen { + v = v[:maxLen] + v = strings.Trim(v, "-_") + } + if v == "" { + return "job" + } + return v +} + +func batchImageTestStringPtr(v string) *string { + return &v +} + +func batchImageDerefTest(v *string) string { + if v == nil { + return "" + } + return *v +} diff --git a/backend/internal/repository/concurrency_cache.go b/backend/internal/repository/concurrency_cache.go index 57b18c6903..b657c1ce8f 100644 --- a/backend/internal/repository/concurrency_cache.go +++ b/backend/internal/repository/concurrency_cache.go @@ -6,6 +6,7 @@ import ( "fmt" "strconv" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" ) @@ -45,6 +46,10 @@ const ( // 后台清理只按批处理索引候选,避免单次任务占用 Redis 太久。 activeIndexCleanupBatchSize = 1000 activeIndexPipelineChunkSize = 500 + + // 一次性迁移 marker:活跃索引机制上线前遗留的等待计数键无法被索引发现, + // 且有流量时 TTL 会被不断刷新,必须清扫一次。marker 存在即代表已完成。 + legacyWaitSweepMarkerKey = "concurrency:startup:legacy_wait_sweep:v1" ) var ( @@ -54,6 +59,7 @@ var ( // ARGV[1] = maxConcurrency // ARGV[2] = TTL(秒) // ARGV[3] = requestID + // 返回 {是否成功, Redis 当前秒},Go 侧复用同一时间源写活跃索引,省去额外 TIME 往返。 acquireScript = redis.NewScript(` -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME') -- replicates correctly. No-op on Redis 5.0+ (effects replication is default). @@ -76,7 +82,7 @@ var ( if exists ~= false then redis.call('ZADD', key, now, requestID) redis.call('EXPIRE', key, ttl) - return 1 + return {1, now} end -- 检查是否达到并发上限 @@ -84,10 +90,10 @@ var ( if count < maxConcurrency then redis.call('ZADD', key, now, requestID) redis.call('EXPIRE', key, ttl) - return 1 + return {1, now} end - return 0 + return {0, now} `) // getCountScript 统计有序集合中的槽位数量并清理过期条目 @@ -136,46 +142,56 @@ var ( // KEYS[1] = wait queue key // ARGV[1] = maxWait // ARGV[2] = TTL in seconds + // 返回 {是否成功, Redis 当前秒},供 Go 侧免额外 TIME 往返写活跃索引。 incrementWaitScript = redis.NewScript(` + -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME') + -- replicates correctly. No-op on Redis 5.0+ (effects replication is default). + redis.replicate_commands() local current = redis.call('GET', KEYS[1]) if current == false then current = 0 else current = tonumber(current) end + local now = tonumber(redis.call('TIME')[1]) if current >= tonumber(ARGV[1]) then - return 0 + return {0, now} end - local newVal = redis.call('INCR', KEYS[1]) + redis.call('INCR', KEYS[1]) -- Refresh TTL so long-running traffic doesn't expire active queue counters. redis.call('EXPIRE', KEYS[1], ARGV[2]) - return 1 - `) + return {1, now} + `) // incrementAccountWaitScript - account-level wait queue count (refresh TTL on each increment) + // 返回值同 incrementWaitScript:{是否成功, Redis 当前秒}。 incrementAccountWaitScript = redis.NewScript(` - local current = redis.call('GET', KEYS[1]) - if current == false then - current = 0 - else - current = tonumber(current) - end + -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME') + -- replicates correctly. No-op on Redis 5.0+ (effects replication is default). + redis.replicate_commands() + local current = redis.call('GET', KEYS[1]) + if current == false then + current = 0 + else + current = tonumber(current) + end + local now = tonumber(redis.call('TIME')[1]) - if current >= tonumber(ARGV[1]) then - return 0 - end + if current >= tonumber(ARGV[1]) then + return {0, now} + end - local newVal = redis.call('INCR', KEYS[1]) + redis.call('INCR', KEYS[1]) - -- Refresh TTL so long-running traffic doesn't expire active queue counters. - redis.call('EXPIRE', KEYS[1], ARGV[2]) + -- Refresh TTL so long-running traffic doesn't expire active queue counters. + redis.call('EXPIRE', KEYS[1], ARGV[2]) - return 1 - `) + return {1, now} + `) // decrementWaitScript - same as before decrementWaitScript = redis.NewScript(` @@ -209,6 +225,7 @@ var ( // startupCleanupSlotScript 清理单个槽位 key 中非当前进程前缀的成员,避免 Redis Cluster CROSSSLOT。 // KEYS[1] 是有序集合键,ARGV[1] 是当前进程前缀,ARGV[2] 是槽位 TTL。 + // 返回 {清除数量, 剩余成员数},Go 侧据剩余数决定索引 member 去留,无需再回读槽位。 startupCleanupSlotScript = redis.NewScript(` local key = KEYS[1] local activePrefix = ARGV[1] @@ -220,12 +237,13 @@ var ( removed = removed + redis.call('ZREM', key, member) end end - if redis.call('ZCARD', key) == 0 then + local remaining = redis.call('ZCARD', key) + if remaining == 0 then redis.call('DEL', key) else redis.call('EXPIRE', key, slotTTL) end - return removed + return {removed, remaining} `) ) @@ -282,28 +300,32 @@ func (c *concurrencyCache) redisUnixSeconds(ctx context.Context) (int64, error) return now.Unix(), nil } -func (c *concurrencyCache) touchAccountActiveIndex(ctx context.Context, accountID int64, ttlSeconds int) { - c.touchActiveIndex(ctx, accountActiveIndexKey, accountID, ttlSeconds) +// slotIndexSpec 描述一个活跃索引及其对应的槽位/等待键构造方式。 +// 用具名字段避免把 slotKey/waitKey 两个同签名函数按位置传参时写反。 +type slotIndexSpec struct { + indexKey string + slotKey func(int64) string + waitKey func(int64) string } -func (c *concurrencyCache) touchUserActiveIndex(ctx context.Context, userID int64, ttlSeconds int) { - c.touchActiveIndex(ctx, userActiveIndexKey, userID, ttlSeconds) -} +var ( + accountSlotIndex = slotIndexSpec{indexKey: accountActiveIndexKey, slotKey: accountSlotKey, waitKey: accountWaitKey} + userSlotIndex = slotIndexSpec{indexKey: userActiveIndexKey, slotKey: userSlotKey, waitKey: waitQueueKey} +) -// touchActiveIndex 是写路径上的轻量标记:主操作已成功时,尽力把 ID 放入活跃索引。 -// 索引失败不影响并发槽位/等待队列本身,后续释放或清理会再次校正。 -func (c *concurrencyCache) touchActiveIndex(ctx context.Context, indexKey string, id int64, ttlSeconds int) { - if c == nil || c.rdb == nil || id <= 0 || ttlSeconds <= 0 { +// touchActiveIndexAt 是写路径上的轻量标记:主操作已成功时,尽力把 ID 放入活跃索引, +// score 为给定的绝对过期时间(Redis Unix 秒)。索引失败不影响并发槽位/等待队列本身, +// 后续释放或清理会再次校正,因此只记日志不上抛。 +func (c *concurrencyCache) touchActiveIndexAt(ctx context.Context, indexKey string, id int64, expireAt int64) { + if c == nil || c.rdb == nil || id <= 0 || expireAt <= 0 { return } - now, err := c.redisUnixSeconds(ctx) - if err != nil { - return - } - _ = c.rdb.ZAdd(ctx, indexKey, redis.Z{ - Score: float64(now + int64(ttlSeconds)), + if err := c.rdb.ZAdd(ctx, indexKey, redis.Z{ + Score: float64(expireAt), Member: strconv.FormatInt(id, 10), - }).Err() + }).Err(); err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: touch active index %s for %d failed: %v", indexKey, id, err) + } } func (c *concurrencyCache) refreshAccountActiveIndex(ctx context.Context, accountID int64) { @@ -316,22 +338,27 @@ func (c *concurrencyCache) refreshUserActiveIndex(ctx context.Context, userID in // refreshActiveIndex 以 Redis 中的真实槽位/等待数为准重建索引状态。 // 释放槽位、等待计数减少、清理过期成员后都会调用它,防止索引残留。 +// 索引维护是 best-effort:失败只记日志,不影响主流程。 func (c *concurrencyCache) refreshActiveIndex(ctx context.Context, indexKey string, id int64, slotKey, waitKey string) { if c == nil || c.rdb == nil || id <= 0 { return } now, err := c.redisUnixSeconds(ctx) if err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: refresh active index %s for %d failed: %v", indexKey, id, err) return } load, err := c.readActiveLoadForKey(ctx, id, slotKey, waitKey, now) if err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: refresh active index %s for %d failed: %v", indexKey, id, err) return } member := strconv.FormatInt(id, 10) if load.slotCount == 0 && load.waitCount <= 0 { - _ = c.rdb.ZRem(ctx, indexKey, member).Err() + if err := c.rdb.ZRem(ctx, indexKey, member).Err(); err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: remove active index member %s from %s failed: %v", member, indexKey, err) + } return } @@ -339,10 +366,7 @@ func (c *concurrencyCache) refreshActiveIndex(ctx context.Context, indexKey stri if ttlSeconds <= 0 { return } - _ = c.rdb.ZAdd(ctx, indexKey, redis.Z{ - Score: float64(now + int64(ttlSeconds)), - Member: member, - }).Err() + c.touchActiveIndexAt(ctx, indexKey, id, now+int64(ttlSeconds)) } type activeIndexLoad struct { @@ -388,9 +412,9 @@ func (c *concurrencyCache) readActiveLoadForKey(ctx context.Context, id int64, s }, nil } -// readAccountIndexLoads 批量读取账号索引候选的真实负载。 +// readIndexLoads 批量读取索引候选的真实负载(账号/用户通用)。 // 分块 Pipeline 可以减少 Redis 往返,同时避免一次 Pipeline 塞入过多命令。 -func (c *concurrencyCache) readAccountIndexLoads(ctx context.Context, members []string, now int64) ([]activeIndexLoad, []string, error) { +func (c *concurrencyCache) readIndexLoads(ctx context.Context, spec slotIndexSpec, members []string, now int64) ([]activeIndexLoad, []string, error) { loads := make([]activeIndexLoad, 0, len(members)) staleMembers := make([]string, 0) candidates := make([]activeIndexLoad, 0, len(members)) @@ -412,17 +436,17 @@ func (c *concurrencyCache) readAccountIndexLoads(ctx context.Context, members [] chunk := candidates[start:end] pipe := c.rdb.Pipeline() - type accountCmd struct { + type loadCmd struct { activeIndexLoad zcardCmd *redis.IntCmd getCmd *redis.StringCmd } - cmds := make([]accountCmd, 0, len(chunk)) + cmds := make([]loadCmd, 0, len(chunk)) for _, candidate := range chunk { - slotKey := accountSlotKey(candidate.id) - waitKey := accountWaitKey(candidate.id) + slotKey := spec.slotKey(candidate.id) + waitKey := spec.waitKey(candidate.id) pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10)) - cmds = append(cmds, accountCmd{ + cmds = append(cmds, loadCmd{ activeIndexLoad: candidate, zcardCmd: pipe.ZCard(ctx, slotKey), getCmd: pipe.Get(ctx, waitKey), @@ -457,12 +481,26 @@ func (c *concurrencyCache) removeActiveIndexMembers(ctx context.Context, indexKe for _, member := range members { args = append(args, member) } - _ = c.rdb.ZRem(ctx, indexKey, args...).Err() + if err := c.rdb.ZRem(ctx, indexKey, args...).Err(); err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: remove %d active index members from %s failed: %v", len(members), indexKey, err) + } } -// touchActiveIndexForLoad 根据已读取的真实负载刷新索引过期时间。 -func (c *concurrencyCache) touchActiveIndexForLoad(ctx context.Context, indexKey string, load activeIndexLoad) { - c.touchActiveIndex(ctx, indexKey, load.id, c.activeIndexTTL(load.slotCount, load.waitCount)) +// runScriptInt64Pair 执行返回两元素整数数组的 Lua 脚本并解析(如 {result, now}、{removed, remaining})。 +func runScriptInt64Pair(ctx context.Context, rdb *redis.Client, script *redis.Script, keys []string, args ...any) (int64, int64, error) { + raw, err := script.Run(ctx, rdb, keys, args...).Result() + if err != nil { + return 0, 0, err + } + first, err := redisScriptInt64At(raw, 0) + if err != nil { + return 0, 0, fmt.Errorf("parse script value 0: %w", err) + } + second, err := redisScriptInt64At(raw, 1) + if err != nil { + return 0, 0, fmt.Errorf("parse script value 1: %w", err) + } + return first, second, nil } // Account slot operations @@ -470,13 +508,13 @@ func (c *concurrencyCache) touchActiveIndexForLoad(ctx context.Context, indexKey func (c *concurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) { key := accountSlotKey(accountID) // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致 - result, err := acquireScript.Run(ctx, c.rdb, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID).Int() + result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID) if err != nil { return false, err } if result == 1 { // 成功占槽后标记活跃账号,后台清理即可从索引定位候选账号。 - c.touchAccountActiveIndex(ctx, accountID, c.slotTTLSeconds) + c.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.slotTTLSeconds)) } return result == 1, nil } @@ -543,13 +581,13 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) { key := userSlotKey(userID) // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致 - result, err := acquireScript.Run(ctx, c.rdb, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID).Int() + result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID) if err != nil { return false, err } if result == 1 { // 成功占槽后标记活跃用户,避免启动清理依赖全量 SCAN。 - c.touchUserActiveIndex(ctx, userID, c.slotTTLSeconds) + c.touchActiveIndexAt(ctx, userActiveIndexKey, userID, now+int64(c.slotTTLSeconds)) } return result == 1, nil } @@ -626,13 +664,13 @@ func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKey func (c *concurrencyCache) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) { key := waitQueueKey(userID) - result, err := incrementWaitScript.Run(ctx, c.rdb, []string{key}, maxWait, c.waitQueueTTLSeconds).Int() + result, now, err := runScriptInt64Pair(ctx, c.rdb, incrementWaitScript, []string{key}, maxWait, c.waitQueueTTLSeconds) if err != nil { return false, err } if result == 1 { // 等待队列也会让用户保持“活跃”,否则槽位为 0 时后台任务可能漏看等待计数。 - c.touchUserActiveIndex(ctx, userID, c.waitQueueTTLSeconds) + c.touchActiveIndexAt(ctx, userActiveIndexKey, userID, now+int64(c.waitQueueTTLSeconds)) } return result == 1, nil } @@ -651,13 +689,13 @@ func (c *concurrencyCache) DecrementWaitCount(ctx context.Context, userID int64) func (c *concurrencyCache) IncrementAccountWaitCount(ctx context.Context, accountID int64, maxWait int) (bool, error) { key := accountWaitKey(accountID) - result, err := incrementAccountWaitScript.Run(ctx, c.rdb, []string{key}, maxWait, c.waitQueueTTLSeconds).Int() + result, now, err := runScriptInt64Pair(ctx, c.rdb, incrementAccountWaitScript, []string{key}, maxWait, c.waitQueueTTLSeconds) if err != nil { return false, err } if result == 1 { // 账号级等待队列同样写入账号活跃索引,供负载查询和清理任务使用。 - c.touchAccountActiveIndex(ctx, accountID, c.waitQueueTTLSeconds) + c.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.waitQueueTTLSeconds)) } return result == 1, nil } @@ -815,113 +853,129 @@ func (c *concurrencyCache) CleanupExpiredAccountSlots(ctx context.Context, accou return err } -// GetActiveAccountLoadMap 只读取活跃账号索引中的账号负载。 -// 这是给热路径使用的轻量视图,避免为获取全局账号负载而扫描所有槽位键。 -func (c *concurrencyCache) GetActiveAccountLoadMap(ctx context.Context) (map[int64]*service.AccountLoadInfo, error) { - now, err := c.redisUnixSeconds(ctx) - if err != nil { - return nil, err +// CleanupExpiredAccountSlotKeys 处理账号与用户两个活跃索引中已到期的候选。 +// (方法名中的 Account 是历史遗留,保留以避免接口变更;实际同时回收两个索引, +// 否则 user 索引的过期成员没有任何清理路径,会无界累积。) +func (c *concurrencyCache) CleanupExpiredAccountSlotKeys(ctx context.Context) error { + if err := c.reconcileExpiredIndexCandidates(ctx, accountSlotIndex); err != nil { + return err } - if err := c.rdb.ZRemRangeByScore(ctx, accountActiveIndexKey, "-inf", strconv.FormatInt(now, 10)).Err(); err != nil { - return nil, fmt.Errorf("cleanup account active index: %w", err) - } - members, err := c.rdb.ZRangeByScore(ctx, accountActiveIndexKey, &redis.ZRangeBy{ - Min: strconv.FormatInt(now+1, 10), - Max: "+inf", - }).Result() - if err != nil { - return nil, fmt.Errorf("read account active index: %w", err) - } - - loads, staleMembers, err := c.readAccountIndexLoads(ctx, members, now) - if err != nil { - return nil, err - } - - loadMap := make(map[int64]*service.AccountLoadInfo, len(loads)) - for _, load := range loads { - if load.slotCount == 0 && load.waitCount <= 0 { - // 索引候选已没有实际负载,删除 member 而不是返回空负载。 - staleMembers = append(staleMembers, load.member) - continue - } - loadMap[load.id] = &service.AccountLoadInfo{ - AccountID: load.id, - CurrentConcurrency: load.slotCount, - WaitingCount: load.waitCount, - } - c.touchActiveIndexForLoad(ctx, accountActiveIndexKey, load) - } - c.removeActiveIndexMembers(ctx, accountActiveIndexKey, staleMembers) - return loadMap, nil + return c.reconcileExpiredIndexCandidates(ctx, userSlotIndex) } -// CleanupExpiredAccountSlotKeys 只处理索引中过期的账号候选。 -// 若候选仍有真实负载,则刷新索引;若没有负载,则移除索引 member。 -func (c *concurrencyCache) CleanupExpiredAccountSlotKeys(ctx context.Context) error { +// reconcileExpiredIndexCandidates 处理单个活跃索引中 score 已到期的候选: +// 无真实负载则移除 member;仍有负载则按真实负载批量刷新 score。 +func (c *concurrencyCache) reconcileExpiredIndexCandidates(ctx context.Context, spec slotIndexSpec) error { now, err := c.redisUnixSeconds(ctx) if err != nil { return err } - members, err := c.rdb.ZRangeByScore(ctx, accountActiveIndexKey, &redis.ZRangeBy{ + members, err := c.rdb.ZRangeByScore(ctx, spec.indexKey, &redis.ZRangeBy{ Min: "-inf", Max: strconv.FormatInt(now, 10), Count: activeIndexCleanupBatchSize, }).Result() if err != nil { - return fmt.Errorf("read expired account active index: %w", err) + return fmt.Errorf("read expired index %s: %w", spec.indexKey, err) } - loads, staleMembers, err := c.readAccountIndexLoads(ctx, members, now) + loads, staleMembers, err := c.readIndexLoads(ctx, spec, members, now) if err != nil { return err } + refreshed := make([]redis.Z, 0, len(loads)) for _, load := range loads { if load.slotCount == 0 && load.waitCount <= 0 { // 真实槽位和等待数都为空,说明这个索引 member 已经完成使命。 staleMembers = append(staleMembers, load.member) continue } - c.touchActiveIndexForLoad(ctx, accountActiveIndexKey, load) + refreshed = append(refreshed, redis.Z{ + Score: float64(now + int64(c.activeIndexTTL(load.slotCount, load.waitCount))), + Member: load.member, + }) } - c.removeActiveIndexMembers(ctx, accountActiveIndexKey, staleMembers) + if len(refreshed) > 0 { + if err := c.rdb.ZAdd(ctx, spec.indexKey, refreshed...).Err(); err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: refresh %d active index members in %s failed: %v", len(refreshed), spec.indexKey, err) + } + } + c.removeActiveIndexMembers(ctx, spec.indexKey, staleMembers) return nil } // CleanupStaleProcessSlots 启动时清理非当前进程前缀的槽位。 -// 清理范围来自活跃索引,避免在 Redis 上 SCAN 全部 concurrency:* 键。 +// 清理范围来自活跃索引(含 score 已过期的成员——它们往往正是崩溃进程留下的残留), +// 避免在 Redis 上 SCAN 全部 concurrency:* 键;另有一次性迁移清扫兜底索引机制上线前的遗留等待计数。 // API Key 槽位(concurrency:api_key:*)是 stats-only 数据:每次 Track/读取都会按分数 // 裁剪过期成员,key 自带 TTL,可在一个 slot TTL 内自愈,因此不参与启动清理。 func (c *concurrencyCache) CleanupStaleProcessSlots(ctx context.Context, activeRequestPrefix string) error { if activeRequestPrefix == "" { return nil } + if err := c.sweepLegacyWaitKeysOnce(ctx); err != nil { + return err + } now, err := c.redisUnixSeconds(ctx) if err != nil { return err } - accountMembers, err := c.activeIndexMembers(ctx, accountActiveIndexKey, now) + accountMembers, err := c.allIndexMembers(ctx, accountActiveIndexKey) if err != nil { return err } - if err := c.cleanupStaleProcessSlotsForIndex(ctx, accountActiveIndexKey, accountMembers, activeRequestPrefix, accountSlotKey, accountWaitKey, c.refreshAccountActiveIndex); err != nil { + if err := c.cleanupStaleProcessSlotsForIndex(ctx, accountSlotIndex, accountMembers, activeRequestPrefix, now); err != nil { return err } - userMembers, err := c.activeIndexMembers(ctx, userActiveIndexKey, now) + userMembers, err := c.allIndexMembers(ctx, userActiveIndexKey) if err != nil { return err } - return c.cleanupStaleProcessSlotsForIndex(ctx, userActiveIndexKey, userMembers, activeRequestPrefix, userSlotKey, waitQueueKey, c.refreshUserActiveIndex) + return c.cleanupStaleProcessSlotsForIndex(ctx, userSlotIndex, userMembers, activeRequestPrefix, now) } -// activeIndexMembers 只返回当前仍未过期的索引 member;过期 member 由对应清理任务处理。 -func (c *concurrencyCache) activeIndexMembers(ctx context.Context, indexKey string, now int64) ([]string, error) { - members, err := c.rdb.ZRangeByScore(ctx, indexKey, &redis.ZRangeBy{ - Min: strconv.FormatInt(now+1, 10), - Max: "+inf", - }).Result() +// sweepLegacyWaitKeysOnce 一次性清扫活跃索引机制上线前遗留的等待计数键。 +// 等待计数在有流量时会不断刷新 TTL、无法自然过期,而索引不认识旧键, +// 因此这里例外地做一次 SCAN,用 marker 键保证整个 Redis 数据生命周期内只执行一次。 +// 先清扫后写 marker:清扫失败时下次启动会重试;并发实例重复清扫是幂等的。 +func (c *concurrencyCache) sweepLegacyWaitKeysOnce(ctx context.Context) error { + exists, err := c.rdb.Exists(ctx, legacyWaitSweepMarkerKey).Result() + if err != nil { + return fmt.Errorf("check legacy wait sweep marker: %w", err) + } + if exists > 0 { + return nil + } + for _, pattern := range []string{accountWaitKeyPrefix + "*", waitQueueKeyPrefix + "*"} { + var cursor uint64 + for { + keys, next, err := c.rdb.Scan(ctx, cursor, pattern, 200).Result() + if err != nil { + return fmt.Errorf("scan legacy wait keys %s: %w", pattern, err) + } + if len(keys) > 0 { + if err := c.rdb.Del(ctx, keys...).Err(); err != nil { + return fmt.Errorf("delete legacy wait keys: %w", err) + } + } + cursor = next + if cursor == 0 { + break + } + } + } + if err := c.rdb.Set(ctx, legacyWaitSweepMarkerKey, "1", 0).Err(); err != nil { + return fmt.Errorf("set legacy wait sweep marker: %w", err) + } + return nil +} + +// allIndexMembers 返回索引中全部 member(含 score 已过期的)。 +// 启动清理必须覆盖过期成员:长时间停机后 score 过期的候选恰恰最可能持有死进程残留。 +func (c *concurrencyCache) allIndexMembers(ctx context.Context, indexKey string) ([]string, error) { + members, err := c.rdb.ZRange(ctx, indexKey, 0, -1).Result() if err != nil { return nil, fmt.Errorf("read active index %s: %w", indexKey, err) } @@ -929,17 +983,17 @@ func (c *concurrencyCache) activeIndexMembers(ctx context.Context, indexKey stri } // cleanupStaleProcessSlotsForIndex 逐个处理索引中的账号/用户。 -// Lua 脚本一次只碰一个槽位 key,兼容 Redis Cluster,随后删除重启后已失效的等待计数。 +// Lua 脚本一次只碰一个槽位 key,兼容 Redis Cluster,随后删除重启后已失效的等待计数; +// 索引 member 的去留由脚本返回的剩余槽位数决定,最后批量写回。 func (c *concurrencyCache) cleanupStaleProcessSlotsForIndex( ctx context.Context, - indexKey string, + spec slotIndexSpec, members []string, activeRequestPrefix string, - slotKeyForID func(int64) string, - waitKeyForID func(int64) string, - refreshIndex func(context.Context, int64), + now int64, ) error { staleMembers := make([]string, 0) + refreshed := make([]redis.Z, 0) for _, member := range members { id, err := strconv.ParseInt(member, 10, 64) if err != nil || id <= 0 { @@ -947,14 +1001,28 @@ func (c *concurrencyCache) cleanupStaleProcessSlotsForIndex( continue } - if _, err := startupCleanupSlotScript.Run(ctx, c.rdb, []string{slotKeyForID(id)}, activeRequestPrefix, c.slotTTLSeconds).Result(); err != nil { - return fmt.Errorf("cleanup stale process slots %s: %w", slotKeyForID(id), err) + _, remaining, err := runScriptInt64Pair(ctx, c.rdb, startupCleanupSlotScript, []string{spec.slotKey(id)}, activeRequestPrefix, c.slotTTLSeconds) + if err != nil { + return fmt.Errorf("cleanup stale process slots %s: %w", spec.slotKey(id), err) } - if err := c.rdb.Del(ctx, waitKeyForID(id)).Err(); err != nil { - return fmt.Errorf("delete stale wait key %s: %w", waitKeyForID(id), err) + // 等待计数属于已死进程,直接删除;剩余槽位(当前进程前缀)决定索引 member 去留。 + if err := c.rdb.Del(ctx, spec.waitKey(id)).Err(); err != nil { + return fmt.Errorf("delete stale wait key %s: %w", spec.waitKey(id), err) + } + if remaining > 0 { + refreshed = append(refreshed, redis.Z{ + Score: float64(now + int64(c.slotTTLSeconds)), + Member: member, + }) + } else { + staleMembers = append(staleMembers, member) } - refreshIndex(ctx, id) } - c.removeActiveIndexMembers(ctx, indexKey, staleMembers) + if len(refreshed) > 0 { + if err := c.rdb.ZAdd(ctx, spec.indexKey, refreshed...).Err(); err != nil { + logger.LegacyPrintf("repository.concurrency", "Warning: refresh %d active index members in %s failed: %v", len(refreshed), spec.indexKey, err) + } + } + c.removeActiveIndexMembers(ctx, spec.indexKey, staleMembers) return nil } diff --git a/backend/internal/repository/concurrency_cache_integration_test.go b/backend/internal/repository/concurrency_cache_integration_test.go index 3c831487de..f7e27d1118 100644 --- a/backend/internal/repository/concurrency_cache_integration_test.go +++ b/backend/internal/repository/concurrency_cache_integration_test.go @@ -77,59 +77,61 @@ func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() { require.Equal(s.T(), 1, cur, "expected 1 after release") } -func (s *ConcurrencyCacheSuite) TestActiveAccountLoadMap_AcquireAndRelease() { +func (s *ConcurrencyCacheSuite) TestAccountActiveIndex_AcquireAndRelease() { accountID := int64(610) - reqID := "active-load-req" + member := strconv.FormatInt(accountID, 10) + reqID := "active-index-req" + + now, err := s.rawCache.redisUnixSeconds(s.ctx) + require.NoError(s.T(), err) ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID) require.NoError(s.T(), err) require.True(s.T(), ok) - loadMap, err := s.rawCache.GetActiveAccountLoadMap(s.ctx) + score, err := s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result() require.NoError(s.T(), err) - require.Contains(s.T(), loadMap, accountID) - require.Equal(s.T(), 1, loadMap[accountID].CurrentConcurrency) + require.Greater(s.T(), int64(score), now, "index score should be a future expiry") require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, reqID)) - loadMap, err = s.rawCache.GetActiveAccountLoadMap(s.ctx) - require.NoError(s.T(), err) - require.NotContains(s.T(), loadMap, accountID) + _, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result() + require.ErrorIs(s.T(), err, redis.Nil, "index member should be removed after load drops to zero") } -func (s *ConcurrencyCacheSuite) TestActiveAccountLoadMap_AccountWaitIndexLifecycle() { +func (s *ConcurrencyCacheSuite) TestAccountActiveIndex_WaitLifecycle() { accountID := int64(611) + member := strconv.FormatInt(accountID, 10) ok, err := s.cache.IncrementAccountWaitCount(s.ctx, accountID, 2) require.NoError(s.T(), err) require.True(s.T(), ok) - loadMap, err := s.rawCache.GetActiveAccountLoadMap(s.ctx) - require.NoError(s.T(), err) - require.Contains(s.T(), loadMap, accountID) - require.Equal(s.T(), 1, loadMap[accountID].WaitingCount) + _, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result() + require.NoError(s.T(), err, "wait increment should register index member") require.NoError(s.T(), s.cache.DecrementAccountWaitCount(s.ctx, accountID)) - loadMap, err = s.rawCache.GetActiveAccountLoadMap(s.ctx) - require.NoError(s.T(), err) - require.NotContains(s.T(), loadMap, accountID) + _, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result() + require.ErrorIs(s.T(), err, redis.Nil, "index member should be removed after wait drops to zero") } -func (s *ConcurrencyCacheSuite) TestActiveAccountLoadMap_RemovesInvalidIndexMember() { - now, err := s.rawCache.redisUnixSeconds(s.ctx) - require.NoError(s.T(), err) - require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{ - Score: float64(now + 60), - Member: "not-an-account-id", - }).Err()) +func (s *ConcurrencyCacheSuite) TestUserActiveIndex_AcquireAndRelease() { + userID := int64(612) + member := strconv.FormatInt(userID, 10) + reqID := "user-active-index-req" - loadMap, err := s.rawCache.GetActiveAccountLoadMap(s.ctx) + ok, err := s.cache.AcquireUserSlot(s.ctx, userID, 2, reqID) require.NoError(s.T(), err) - require.Empty(s.T(), loadMap) + require.True(s.T(), ok) - _, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, "not-an-account-id").Result() - require.ErrorIs(s.T(), err, redis.Nil) + _, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, member).Result() + require.NoError(s.T(), err, "acquire should register user index member") + + require.NoError(s.T(), s.cache.ReleaseUserSlot(s.ctx, userID, reqID)) + + _, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, member).Result() + require.ErrorIs(s.T(), err, redis.Nil, "user index member should be removed after release") } func (s *ConcurrencyCacheSuite) TestAccountSlot_TTL() { @@ -351,6 +353,8 @@ func (s *ConcurrencyCacheSuite) TestAccountWaitQueue_IncrementAndDecrement() { } func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() { + // 预置迁移 marker,隔离一次性清扫,只验证索引驱动的清理路径。 + require.NoError(s.T(), s.rdb.Set(s.ctx, legacyWaitSweepMarkerKey, "1", 0).Err()) accountID := int64(901) userID := int64(902) apiKeyID := int64(903) @@ -619,7 +623,113 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() { require.ErrorIs(s.T(), err, redis.Nil) } +func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys_ReapsUserIndex() { + now, err := s.rawCache.redisUnixSeconds(s.ctx) + require.NoError(s.T(), err) + expiredScore := float64(now - 10) + userKeyWithFresh := fmt.Sprintf("%s%d", userSlotKeyPrefix, 401) + + // 401 有真实负载但索引 score 已过期:应刷新而不是删除。 + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKeyWithFresh, + redis.Z{Score: float64(now), Member: "fresh"}, + ).Err()) + // 402 无任何负载:过期索引 member 应被回收。 + // 非法 member 也应随过期候选一并清除。 + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, + redis.Z{Score: expiredScore, Member: "401"}, + redis.Z{Score: expiredScore, Member: "402"}, + redis.Z{Score: expiredScore, Member: "not-a-user-id"}, + ).Err()) + + require.NoError(s.T(), s.cache.CleanupExpiredAccountSlotKeys(s.ctx)) + + score, err := s.rdb.ZScore(s.ctx, userActiveIndexKey, "401").Result() + require.NoError(s.T(), err) + require.Greater(s.T(), int64(score), now, "loaded user should be re-scheduled, not dropped") + + _, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, "402").Result() + require.ErrorIs(s.T(), err, redis.Nil, "idle expired user member should be reaped") + + _, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, "not-a-user-id").Result() + require.ErrorIs(s.T(), err, redis.Nil, "invalid member should be reaped") +} + +func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_LegacyWaitSweepRunsOnce() { + unindexedAccountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, 2901) + unindexedUserWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, 2902) + require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 5, time.Minute).Err()) + require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedUserWaitKey, 3, time.Minute).Err()) + + // 首次运行:marker 不存在,一次性清扫删除所有遗留等待计数(含未入索引的)。 + require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-")) + + _, err := s.rdb.Get(s.ctx, unindexedAccountWaitKey).Result() + require.ErrorIs(s.T(), err, redis.Nil, "legacy account wait key should be swept on first startup") + _, err = s.rdb.Get(s.ctx, unindexedUserWaitKey).Result() + require.ErrorIs(s.T(), err, redis.Nil, "legacy user wait key should be swept on first startup") + + exists, err := s.rdb.Exists(s.ctx, legacyWaitSweepMarkerKey).Result() + require.NoError(s.T(), err) + require.EqualValues(s.T(), 1, exists, "sweep marker should be set after first run") + + // 再次运行:marker 已存在,未入索引的等待计数不再被触碰。 + require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 5, time.Minute).Err()) + require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-")) + val, err := s.rdb.Get(s.ctx, unindexedAccountWaitKey).Int() + require.NoError(s.T(), err, "sweep must not run twice") + require.Equal(s.T(), 5, val) +} + +func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_ProcessesExpiredIndexMembers() { + // score 已过期的索引成员往往正是崩溃进程留下的残留,启动清理必须覆盖它们。 + require.NoError(s.T(), s.rdb.Set(s.ctx, legacyWaitSweepMarkerKey, "1", 0).Err()) + accountID := int64(3901) + userID := int64(3902) + accountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID) + userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID) + accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID) + + now, err := s.rawCache.redisUnixSeconds(s.ctx) + require.NoError(s.T(), err) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKey, + redis.Z{Score: float64(now), Member: "oldproc-1"}, + ).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKey, + redis.Z{Score: float64(now), Member: "oldproc-2"}, + ).Err()) + require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 4, time.Minute).Err()) + // 索引 score 设为过去时刻,模拟长时间停机后索引已“过期”。 + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{ + Score: float64(now - 100), + Member: strconv.FormatInt(accountID, 10), + }).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{ + Score: float64(now - 100), + Member: strconv.FormatInt(userID, 10), + }).Err()) + + require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-")) + + exists, err := s.rdb.Exists(s.ctx, accountKey).Result() + require.NoError(s.T(), err) + require.EqualValues(s.T(), 0, exists, "stale slot key of expired index member should be purged") + + exists, err = s.rdb.Exists(s.ctx, userKey).Result() + require.NoError(s.T(), err) + require.EqualValues(s.T(), 0, exists) + + _, err = s.rdb.Get(s.ctx, accountWaitKey).Result() + require.ErrorIs(s.T(), err, redis.Nil, "wait counter of expired index member should be deleted") + + _, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, strconv.FormatInt(accountID, 10)).Result() + require.ErrorIs(s.T(), err, redis.Nil, "emptied member should be removed from index") + _, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, strconv.FormatInt(userID, 10)).Result() + require.ErrorIs(s.T(), err, redis.Nil) +} + func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesAndWaitCounters() { + // 预置迁移 marker,确保等待计数删除来自索引驱动路径而非一次性清扫。 + require.NoError(s.T(), s.rdb.Set(s.ctx, legacyWaitSweepMarkerKey, "1", 0).Err()) accountID := int64(901) userID := int64(902) accountSlotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID) diff --git a/backend/internal/repository/github_release_service.go b/backend/internal/repository/github_release_service.go index ad1f22e39b..85fab2b344 100644 --- a/backend/internal/repository/github_release_service.go +++ b/backend/internal/repository/github_release_service.go @@ -67,6 +67,10 @@ func (c *githubReleaseClientError) FetchLatestRelease(ctx context.Context, repo return nil, c.err } +func (c *githubReleaseClientError) FetchRecentReleases(ctx context.Context, repo string, perPage int) ([]*service.GitHubRelease, error) { + return nil, c.err +} + func (c *githubReleaseClientError) DownloadFile(ctx context.Context, url, dest string, maxSize int64) error { return c.err } @@ -103,6 +107,40 @@ func (c *githubReleaseClient) FetchLatestRelease(ctx context.Context, repo strin return &release, nil } +func (c *githubReleaseClient) FetchRecentReleases(ctx context.Context, repo string, perPage int) ([]*service.GitHubRelease, error) { + if perPage <= 0 { + perPage = 10 + } + if perPage > 100 { + perPage = 100 // GitHub API hard limit + } + url := fmt.Sprintf("https://api.github.com/repos/%s/releases?per_page=%d", repo, perPage) + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/vnd.github.v3+json") + req.Header.Set("User-Agent", "Sub2API-Updater") + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("GitHub API returned %d", resp.StatusCode) + } + + var releases []*service.GitHubRelease + if err := json.NewDecoder(resp.Body).Decode(&releases); err != nil { + return nil, err + } + + return releases, nil +} + func (c *githubReleaseClient) DownloadFile(ctx context.Context, url, dest string, maxSize int64) error { req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) if err != nil { diff --git a/backend/internal/repository/github_release_service_test.go b/backend/internal/repository/github_release_service_test.go index d375a19300..abc9a604e6 100644 --- a/backend/internal/repository/github_release_service_test.go +++ b/backend/internal/repository/github_release_service_test.go @@ -29,6 +29,9 @@ type testTransport struct { func (t *testTransport) RoundTrip(req *http.Request) (*http.Response, error) { // Rewrite the URL to point to our test server testURL := t.testServerURL + req.URL.Path + if req.URL.RawQuery != "" { + testURL += "?" + req.URL.RawQuery + } newReq, err := http.NewRequestWithContext(req.Context(), req.Method, testURL, req.Body) if err != nil { return nil, err @@ -245,6 +248,76 @@ func (s *GitHubReleaseServiceSuite) TestFetchLatestRelease_Success() { require.Equal(s.T(), "app-linux-amd64.tar.gz", release.Assets[0].Name) } +func (s *GitHubReleaseServiceSuite) TestFetchRecentReleases_Success() { + releasesJSON := `[ + { + "tag_name": "v1.0.1", + "name": "Release 1.0.1", + "html_url": "https://github.com/test/repo/releases/v1.0.1", + "published_at": "2026-07-08T00:00:00Z", + "prerelease": false, + "assets": [ + { + "name": "app-linux-amd64.tar.gz", + "browser_download_url": "https://github.com/test/repo/releases/download/v1.0.1/app-linux-amd64.tar.gz" + } + ] + }, + { + "tag_name": "v1.0.1-rc1", + "name": "Release 1.0.1-rc1", + "prerelease": true + }, + { + "tag_name": "v1.0.0", + "name": "Release 1.0.0" + } + ]` + + s.srv = newLocalTestServer(s.T(), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(s.T(), "/repos/test/repo/releases", r.URL.Path) + require.Equal(s.T(), "15", r.URL.Query().Get("per_page")) + require.Equal(s.T(), "application/vnd.github.v3+json", r.Header.Get("Accept")) + require.Equal(s.T(), "Sub2API-Updater", r.Header.Get("User-Agent")) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(releasesJSON)) + })) + + s.client = &githubReleaseClient{ + httpClient: &http.Client{ + Transport: &testTransport{testServerURL: s.srv.URL}, + }, + downloadHTTPClient: &http.Client{}, + } + + releases, err := s.client.FetchRecentReleases(context.Background(), "test/repo", 15) + require.NoError(s.T(), err) + require.Len(s.T(), releases, 3) + require.Equal(s.T(), "v1.0.1", releases[0].TagName) + require.False(s.T(), releases[0].Prerelease) + require.Len(s.T(), releases[0].Assets, 1) + require.True(s.T(), releases[1].Prerelease) + require.Equal(s.T(), "v1.0.0", releases[2].TagName) +} + +func (s *GitHubReleaseServiceSuite) TestFetchRecentReleases_Non200() { + s.srv = newLocalTestServer(s.T(), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + + s.client = &githubReleaseClient{ + httpClient: &http.Client{ + Transport: &testTransport{testServerURL: s.srv.URL}, + }, + downloadHTTPClient: &http.Client{}, + } + + _, err := s.client.FetchRecentReleases(context.Background(), "test/repo", 15) + require.Error(s.T(), err) + require.Contains(s.T(), err.Error(), "403") +} + func (s *GitHubReleaseServiceSuite) TestFetchLatestRelease_Non200() { s.srv = newLocalTestServer(s.T(), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNotFound) diff --git a/backend/internal/repository/group_repo.go b/backend/internal/repository/group_repo.go index a4e173006e..37529c60be 100644 --- a/backend/internal/repository/group_repo.go +++ b/backend/internal/repository/group_repo.go @@ -50,11 +50,19 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er SetNillableWeeklyLimitUsd(groupIn.WeeklyLimitUSD). SetNillableMonthlyLimitUsd(groupIn.MonthlyLimitUSD). SetAllowImageGeneration(groupIn.AllowImageGeneration). + SetAllowBatchImageGeneration(groupIn.AllowBatchImageGeneration). SetImageRateIndependent(groupIn.ImageRateIndependent). SetImageRateMultiplier(groupIn.ImageRateMultiplier). SetNillableImagePrice1k(groupIn.ImagePrice1K). SetNillableImagePrice2k(groupIn.ImagePrice2K). SetNillableImagePrice4k(groupIn.ImagePrice4K). + SetBatchImageDiscountMultiplier(groupIn.BatchImageDiscountMultiplier). + SetBatchImageHoldMultiplier(groupIn.BatchImageHoldMultiplier). + SetVideoRateIndependent(groupIn.VideoRateIndependent). + SetVideoRateMultiplier(groupIn.VideoRateMultiplier). + SetNillableVideoPrice480p(groupIn.VideoPrice480P). + SetNillableVideoPrice720p(groupIn.VideoPrice720P). + SetNillableVideoPrice1080p(groupIn.VideoPrice1080P). SetDefaultValidityDays(groupIn.DefaultValidityDays). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetNillableFallbackGroupID(groupIn.FallbackGroupID). @@ -132,11 +140,19 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er SetNillableWeeklyLimitUsd(groupIn.WeeklyLimitUSD). SetNillableMonthlyLimitUsd(groupIn.MonthlyLimitUSD). SetAllowImageGeneration(groupIn.AllowImageGeneration). + SetAllowBatchImageGeneration(groupIn.AllowBatchImageGeneration). SetImageRateIndependent(groupIn.ImageRateIndependent). SetImageRateMultiplier(groupIn.ImageRateMultiplier). SetNillableImagePrice1k(groupIn.ImagePrice1K). SetNillableImagePrice2k(groupIn.ImagePrice2K). SetNillableImagePrice4k(groupIn.ImagePrice4K). + SetBatchImageDiscountMultiplier(groupIn.BatchImageDiscountMultiplier). + SetBatchImageHoldMultiplier(groupIn.BatchImageHoldMultiplier). + SetVideoRateIndependent(groupIn.VideoRateIndependent). + SetVideoRateMultiplier(groupIn.VideoRateMultiplier). + SetNillableVideoPrice480p(groupIn.VideoPrice480P). + SetNillableVideoPrice720p(groupIn.VideoPrice720P). + SetNillableVideoPrice1080p(groupIn.VideoPrice1080P). SetDefaultValidityDays(groupIn.DefaultValidityDays). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetModelRoutingEnabled(groupIn.ModelRoutingEnabled). @@ -184,6 +200,21 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er } else { builder = builder.ClearImagePrice4k() } + if groupIn.VideoPrice480P != nil { + builder = builder.SetVideoPrice480p(*groupIn.VideoPrice480P) + } else { + builder = builder.ClearVideoPrice480p() + } + if groupIn.VideoPrice720P != nil { + builder = builder.SetVideoPrice720p(*groupIn.VideoPrice720P) + } else { + builder = builder.ClearVideoPrice720p() + } + if groupIn.VideoPrice1080P != nil { + builder = builder.SetVideoPrice1080p(*groupIn.VideoPrice1080P) + } else { + builder = builder.ClearVideoPrice1080p() + } // 处理 FallbackGroupID:nil 时清除,否则设置 if groupIn.FallbackGroupID != nil { diff --git a/backend/internal/repository/migrations_runner.go b/backend/internal/repository/migrations_runner.go index 285326537d..7c045fea74 100644 --- a/backend/internal/repository/migrations_runner.go +++ b/backend/internal/repository/migrations_runner.go @@ -77,6 +77,8 @@ var migrationChecksumCompatibilityRules = map[string]migrationChecksumCompatibil "119_enforce_payment_orders_out_trade_no_unique.sql": newMigrationChecksumCompatibilityRule("0bbe809ae48a9d811dabda1ba1c74955bd71c4a9cc610f9128816818dfa6c11e", "ebd2c67cce0116393fb4f1b5d5116a67c6aceb73820dfb5133d1ff6f36d72d34"), "120_enforce_payment_orders_out_trade_no_unique_notx.sql": newMigrationChecksumCompatibilityRule("34aadc0db59a4e390f92a12b73bd74642d9724f33124f73638ae00089ea5e074", "e77921f79d539bc24575cb9c16cbe566d2b23ce816190343d0a7568f6a3fcf61", "707431450603e70a43ce9fbd61e0c12fa67da4875158ccefabacea069587ab22", "04b082b5a239c525154fe9185d324ee2b05ff90da9297e10dba19f9be79aa59a"), "123_fix_legacy_auth_source_grant_on_signup_defaults.sql": newMigrationChecksumCompatibilityRule("2ce43c2cd89e9f9e1febd34a407ed9e84d177386c5544b6f02c1f58a21129f57", "6cd33422f215dcd1f486ab6f35c0ea5805d9ca69bb25906d94bc649156657145"), + "159_batch_image_foundation.sql": newMigrationChecksumCompatibilityRule("d902b70982025ec519749faf058aab7631e82c3f48167b9a4ae4db718eb72cce", "82da85b5d98e67a0507647b873a40373e84538e4adafdeed6767c0ac8b6570b2"), + "161_batch_image_pricing_snapshot.sql": newMigrationChecksumCompatibilityRule("4012af3e43636cb6af22e0176d59d1fcc70615c0f310194329461ae462c4fbd6", "96d915c9b7a6941ae99039e0ff3f1a61481eb9bddd933d11c6fadb2274554e87"), } // ApplyMigrations 将嵌入的 SQL 迁移文件应用到指定的数据库。 diff --git a/backend/internal/repository/migrations_schema_integration_test.go b/backend/internal/repository/migrations_schema_integration_test.go index d39ac39cba..7ef98dd700 100644 --- a/backend/internal/repository/migrations_schema_integration_test.go +++ b/backend/internal/repository/migrations_schema_integration_test.go @@ -49,6 +49,9 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) { requireColumn(t, tx, "usage_logs", "image_output_size", "character varying", 32, true) requireColumn(t, tx, "usage_logs", "image_size_source", "character varying", 16, true) requireColumn(t, tx, "usage_logs", "image_size_breakdown", "jsonb", 0, true) + requireColumn(t, tx, "usage_logs", "video_count", "integer", 0, false) + requireColumn(t, tx, "usage_logs", "video_resolution", "character varying", 10, true) + requireColumn(t, tx, "usage_logs", "video_duration_seconds", "integer", 0, true) requireConstraintDefinitionContains( t, tx, @@ -66,6 +69,9 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) { "usage_logs", "usage_logs_image_billing_size_check", "image_count", + "billing_mode", + "'video'", + "video_count", "image_size IS NOT NULL", "'1K'", "'2K'", diff --git a/backend/internal/repository/usage_billing_repo.go b/backend/internal/repository/usage_billing_repo.go index 91ac536eee..ce0adc2772 100644 --- a/backend/internal/repository/usage_billing_repo.go +++ b/backend/internal/repository/usage_billing_repo.go @@ -63,23 +63,27 @@ func (r *usageBillingRepository) Apply(ctx context.Context, cmd *service.UsageBi } func (r *usageBillingRepository) claimUsageBillingKey(ctx context.Context, tx *sql.Tx, cmd *service.UsageBillingCommand) (bool, error) { + return r.claimUsageBillingRequest(ctx, tx, cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint) +} + +func (r *usageBillingRepository) claimUsageBillingRequest(ctx context.Context, tx *sql.Tx, requestID string, apiKeyID int64, requestFingerprint string) (bool, error) { var id int64 err := tx.QueryRowContext(ctx, ` INSERT INTO usage_billing_dedup (request_id, api_key_id, request_fingerprint) VALUES ($1, $2, $3) ON CONFLICT (request_id, api_key_id) DO NOTHING RETURNING id - `, cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint).Scan(&id) + `, requestID, apiKeyID, requestFingerprint).Scan(&id) if errors.Is(err, sql.ErrNoRows) { var existingFingerprint string if err := tx.QueryRowContext(ctx, ` SELECT request_fingerprint FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2 - `, cmd.RequestID, cmd.APIKeyID).Scan(&existingFingerprint); err != nil { + `, requestID, apiKeyID).Scan(&existingFingerprint); err != nil { return false, err } - if strings.TrimSpace(existingFingerprint) != strings.TrimSpace(cmd.RequestFingerprint) { + if strings.TrimSpace(existingFingerprint) != strings.TrimSpace(requestFingerprint) { return false, service.ErrUsageBillingRequestConflict } return false, nil @@ -92,9 +96,9 @@ func (r *usageBillingRepository) claimUsageBillingKey(ctx context.Context, tx *s SELECT request_fingerprint FROM usage_billing_dedup_archive WHERE request_id = $1 AND api_key_id = $2 - `, cmd.RequestID, cmd.APIKeyID).Scan(&archivedFingerprint) + `, requestID, apiKeyID).Scan(&archivedFingerprint) if err == nil { - if strings.TrimSpace(archivedFingerprint) != strings.TrimSpace(cmd.RequestFingerprint) { + if strings.TrimSpace(archivedFingerprint) != strings.TrimSpace(requestFingerprint) { return false, service.ErrUsageBillingRequestConflict } return false, nil @@ -105,6 +109,68 @@ func (r *usageBillingRepository) claimUsageBillingKey(ctx context.Context, tx *s return true, nil } +func (r *usageBillingRepository) ReserveBatchImageBalance(ctx context.Context, cmd *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error) { + return r.applyBatchImageBalanceHold(ctx, cmd, reserveUsageBillingBatchImageBalance) +} + +func (r *usageBillingRepository) CaptureBatchImageBalance(ctx context.Context, cmd *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error) { + return r.applyBatchImageBalanceHold(ctx, cmd, captureUsageBillingBatchImageBalance) +} + +func (r *usageBillingRepository) ReleaseBatchImageBalance(ctx context.Context, cmd *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error) { + return r.applyBatchImageBalanceHold(ctx, cmd, releaseUsageBillingBatchImageBalance) +} + +func (r *usageBillingRepository) applyBatchImageBalanceHold( + ctx context.Context, + cmd *service.BatchImageBalanceHoldCommand, + apply func(context.Context, *sql.Tx, *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error), +) (_ *service.BatchImageBalanceHoldResult, err error) { + if cmd == nil { + return &service.BatchImageBalanceHoldResult{}, nil + } + if r == nil || r.db == nil { + return nil, errors.New("usage billing repository db is nil") + } + cmd.Normalize() + if cmd.RequestID == "" { + return nil, service.ErrUsageBillingRequestIDRequired + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer func() { + if tx != nil { + _ = tx.Rollback() + } + }() + + applied, err := r.claimUsageBillingRequest(ctx, tx, cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint) + if err != nil { + return nil, err + } + if !applied { + return &service.BatchImageBalanceHoldResult{Applied: false}, nil + } + + result, err := apply(ctx, tx, cmd) + if err != nil { + return nil, err + } + if result == nil { + result = &service.BatchImageBalanceHoldResult{} + } + result.Applied = true + + if err := tx.Commit(); err != nil { + return nil, err + } + tx = nil + return result, nil +} + func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, tx *sql.Tx, cmd *service.UsageBillingCommand, result *service.UsageBillingApplyResult) error { if cmd.SubscriptionCost > 0 && cmd.SubscriptionID != nil { if err := incrementUsageBillingSubscription(ctx, tx, *cmd.SubscriptionID, cmd.SubscriptionCost); err != nil { @@ -206,6 +272,147 @@ func deductUsageBillingBalance(ctx context.Context, tx *sql.Tx, userID int64, am return newBalance, false, nil } +func reserveUsageBillingBatchImageBalance(ctx context.Context, tx *sql.Tx, cmd *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error) { + if cmd.HoldAmount <= 0 { + return &service.BatchImageBalanceHoldResult{}, nil + } + var balance, frozen float64 + err := tx.QueryRowContext(ctx, ` + UPDATE users + SET balance = balance - $1, + frozen_balance = COALESCE(frozen_balance, 0) + $1, + updated_at = NOW() + WHERE id = $2 AND deleted_at IS NULL AND balance >= $1 + RETURNING balance, frozen_balance + `, cmd.HoldAmount, cmd.UserID).Scan(&balance, &frozen) + if err == nil { + return &service.BatchImageBalanceHoldResult{NewBalance: &balance, FrozenBalance: &frozen}, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + if exists, existsErr := userExistsForBilling(ctx, tx, cmd.UserID); existsErr != nil { + return nil, existsErr + } else if !exists { + return nil, service.ErrUserNotFound + } + return nil, service.ErrBatchImageInsufficientBalance +} + +func captureUsageBillingBatchImageBalance(ctx context.Context, tx *sql.Tx, cmd *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error) { + if cmd.HoldAmount <= 0 && cmd.ActualAmount <= 0 { + return &service.BatchImageBalanceHoldResult{}, nil + } + if cmd.ActualAmount-cmd.HoldAmount > 0.00000001 { + return nil, service.ErrBatchImageSettlementCostExceedsHold + } + var balance, frozen float64 + err := tx.QueryRowContext(ctx, ` + UPDATE users + SET balance = balance + + CASE WHEN $1 > $2 THEN $1 - $2 ELSE 0 END + - CASE WHEN $2 > $1 THEN $2 - $1 ELSE 0 END, + frozen_balance = COALESCE(frozen_balance, 0) - $1, + updated_at = NOW() + WHERE id = $3 AND deleted_at IS NULL AND COALESCE(frozen_balance, 0) >= $1 + RETURNING balance, frozen_balance + `, cmd.HoldAmount, cmd.ActualAmount, cmd.UserID).Scan(&balance, &frozen) + if err == nil { + return &service.BatchImageBalanceHoldResult{NewBalance: &balance, FrozenBalance: &frozen}, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + if exists, existsErr := userExistsForBilling(ctx, tx, cmd.UserID); existsErr != nil { + return nil, existsErr + } else if !exists { + return nil, service.ErrUserNotFound + } + return nil, errors.New("batch image frozen balance is insufficient") +} + +func releaseUsageBillingBatchImageBalance(ctx context.Context, tx *sql.Tx, cmd *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error) { + if cmd.HoldAmount <= 0 { + return &service.BatchImageBalanceHoldResult{}, nil + } + // 释放前校验该 job 确实预留过 hold(hold request id 已被 claim), + // 防止从未成功冻结的 job 触发"幻影释放",从其他用户的冻结资金池中凭空生成余额。 + held, heldErr := batchImageHoldClaimExists(ctx, tx, service.BatchImageHoldRequestID(cmd.BatchID), cmd.APIKeyID) + if heldErr != nil { + return nil, heldErr + } + if !held { + logger.LegacyPrintf("repository.usage_billing", "[BatchImage] release skipped, hold was never reserved: batch=%s", cmd.BatchID) + return &service.BatchImageBalanceHoldResult{}, nil + } + var balance, frozen float64 + err := tx.QueryRowContext(ctx, ` + UPDATE users + SET balance = balance + $1, + frozen_balance = COALESCE(frozen_balance, 0) - $1, + updated_at = NOW() + WHERE id = $2 AND deleted_at IS NULL AND COALESCE(frozen_balance, 0) >= $1 + RETURNING balance, frozen_balance + `, cmd.HoldAmount, cmd.UserID).Scan(&balance, &frozen) + if err == nil { + return &service.BatchImageBalanceHoldResult{NewBalance: &balance, FrozenBalance: &frozen}, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + if exists, existsErr := userExistsForBilling(ctx, tx, cmd.UserID); existsErr != nil { + return nil, existsErr + } else if !exists { + return nil, service.ErrUserNotFound + } + return nil, errors.New("batch image frozen balance is insufficient") +} + +// batchImageHoldClaimExists 检查 hold request id 是否已在 dedup(或归档)表中被 claim, +// 即该 batch 的冻结操作确实成功提交过。 +func batchImageHoldClaimExists(ctx context.Context, tx *sql.Tx, holdRequestID string, apiKeyID int64) (bool, error) { + var exists int + err := tx.QueryRowContext(ctx, ` + SELECT 1 + FROM usage_billing_dedup + WHERE request_id = $1 AND api_key_id = $2 + `, holdRequestID, apiKeyID).Scan(&exists) + if err == nil { + return true, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return false, err + } + err = tx.QueryRowContext(ctx, ` + SELECT 1 + FROM usage_billing_dedup_archive + WHERE request_id = $1 AND api_key_id = $2 + `, holdRequestID, apiKeyID).Scan(&exists) + if err == nil { + return true, nil + } + if errors.Is(err, sql.ErrNoRows) { + return false, nil + } + return false, err +} + +func userExistsForBilling(ctx context.Context, tx *sql.Tx, userID int64) (bool, error) { + var exists int + err := tx.QueryRowContext(ctx, ` + SELECT 1 + FROM users + WHERE id = $1 AND deleted_at IS NULL + `, userID).Scan(&exists) + if errors.Is(err, sql.ErrNoRows) { + return false, nil + } + if err != nil { + return false, err + } + return true, nil +} + func incrementUsageBillingAPIKeyQuota(ctx context.Context, tx *sql.Tx, apiKeyID int64, amount float64) (bool, error) { var exhausted bool err := tx.QueryRowContext(ctx, ` diff --git a/backend/internal/repository/usage_billing_repo_unit_test.go b/backend/internal/repository/usage_billing_repo_unit_test.go index 8ed5530a8f..6fe4d4fc26 100644 --- a/backend/internal/repository/usage_billing_repo_unit_test.go +++ b/backend/internal/repository/usage_billing_repo_unit_test.go @@ -16,6 +16,10 @@ import ( const ( conditionalBalanceDeductSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND balance >= \$1\s+RETURNING balance` overdraftBalanceDeductSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING balance` + reserveBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) \+ \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND balance >= \$1\s+RETURNING balance, frozen_balance` + captureBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance\s+\+ CASE WHEN \$1 > \$2 THEN \$1 - \$2 ELSE 0 END\s+- CASE WHEN \$2 > \$1 THEN \$2 - \$1 ELSE 0 END,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$3 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance` + releaseBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance \+ \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance` + userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL` ) func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) { @@ -117,3 +121,141 @@ func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testi require.NoError(t, tx.Rollback()) require.NoError(t, mock.ExpectationsWereMet()) } + +func TestReserveUsageBillingBatchImageBalance_MovesAvailableToFrozen(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(reserveBatchImageHoldSQL). + WithArgs(2.5, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(7.5, 2.5)) + mock.ExpectCommit() + + result, err := reserveUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 2.5}) + require.NoError(t, err) + require.NotNil(t, result.NewBalance) + require.NotNil(t, result.FrozenBalance) + require.InDelta(t, 7.5, *result.NewBalance, 0.000001) + require.InDelta(t, 2.5, *result.FrozenBalance, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestReserveUsageBillingBatchImageBalance_InsufficientBalance(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(reserveBatchImageHoldSQL). + WithArgs(10.0, int64(42)). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(userExistsForBillingSQL). + WithArgs(int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"?column?"}).AddRow(1)) + mock.ExpectRollback() + + _, err = reserveUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 10}) + require.ErrorIs(t, err, service.ErrBatchImageInsufficientBalance) + require.NoError(t, tx.Rollback()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestCaptureUsageBillingBatchImageBalance_ReleasesRemainder(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(captureBatchImageHoldSQL). + WithArgs(1.0, 0.25, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(9.75, 0.0)) + mock.ExpectCommit() + + result, err := captureUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 1, ActualAmount: 0.25}) + require.NoError(t, err) + require.InDelta(t, 9.75, *result.NewBalance, 0.000001) + require.InDelta(t, 0.0, *result.FrozenBalance, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestCaptureUsageBillingBatchImageBalance_RejectsActualCostOverHold(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectRollback() + + _, err = captureUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 0.5, ActualAmount: 1}) + require.ErrorIs(t, err, service.ErrBatchImageSettlementCostExceedsHold) + require.NoError(t, tx.Rollback()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestReleaseUsageBillingBatchImageBalance_ReturnsFrozenToAvailable(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup\s+WHERE request_id = \$1 AND api_key_id = \$2`). + WithArgs(service.BatchImageHoldRequestID("imgbatch_release"), int64(7)). + WillReturnRows(sqlmock.NewRows([]string{"?column?"}).AddRow(1)) + mock.ExpectQuery(releaseBatchImageHoldSQL). + WithArgs(1.0, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(10.0, 0.0)) + mock.ExpectCommit() + + result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, APIKeyID: 7, BatchID: "imgbatch_release", HoldAmount: 1}) + require.NoError(t, err) + require.InDelta(t, 10.0, *result.NewBalance, 0.000001) + require.InDelta(t, 0.0, *result.FrozenBalance, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestReleaseUsageBillingBatchImageBalance_SkipsWhenHoldNeverReserved(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + // dedup 与归档表均无 hold claim:说明该 job 从未成功冻结, + // 释放必须跳过,不得从他人冻结资金池中凭空生成余额。 + mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup\s+WHERE request_id = \$1 AND api_key_id = \$2`). + WithArgs(service.BatchImageHoldRequestID("imgbatch_phantom"), int64(7)). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup_archive\s+WHERE request_id = \$1 AND api_key_id = \$2`). + WithArgs(service.BatchImageHoldRequestID("imgbatch_phantom"), int64(7)). + WillReturnError(sql.ErrNoRows) + mock.ExpectCommit() + + result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, APIKeyID: 7, BatchID: "imgbatch_phantom", HoldAmount: 1}) + require.NoError(t, err) + require.Nil(t, result.NewBalance) + require.Nil(t, result.FrozenBalance) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/internal/repository/usage_log_repo.go b/backend/internal/repository/usage_log_repo.go index 24c648b0a5..341bdff57d 100644 --- a/backend/internal/repository/usage_log_repo.go +++ b/backend/internal/repository/usage_log_repo.go @@ -1,97 +1,17 @@ package repository import ( - "context" "database/sql" - "encoding/json" - "errors" "fmt" - "os" - "strconv" "strings" "sync" - "sync/atomic" "time" dbent "github.com/Wei-Shaw/sub2api/ent" - dbaccount "github.com/Wei-Shaw/sub2api/ent/account" - dbapikey "github.com/Wei-Shaw/sub2api/ent/apikey" - dbgroup "github.com/Wei-Shaw/sub2api/ent/group" - "github.com/Wei-Shaw/sub2api/ent/schema/mixins" - dbuser "github.com/Wei-Shaw/sub2api/ent/user" - dbusersub "github.com/Wei-Shaw/sub2api/ent/usersubscription" - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" - "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" - "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" "github.com/Wei-Shaw/sub2api/internal/service" - "github.com/lib/pq" gocache "github.com/patrickmn/go-cache" - "golang.org/x/sync/errgroup" ) -const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" - -// usageLogInsertArgTypes must stay in the same order as: -// 1. prepareUsageLogInsert().args -// 2. every INSERT/CTE VALUES column list in this file -// 3. execUsageLogInsertNoResult placeholder positions -// 4. scanUsageLog selected column order (via usageLogSelectColumns) -// -// When adding a usage_logs column, update all of those call sites together. -var usageLogInsertArgTypes = [...]string{ - "bigint", // user_id - "bigint", // api_key_id - "bigint", // account_id - "text", // request_id - "text", // model - "text", // requested_model - "text", // upstream_model - "bigint", // group_id - "bigint", // subscription_id - "integer", // input_tokens - "integer", // output_tokens - "integer", // cache_creation_tokens - "integer", // cache_read_tokens - "integer", // cache_creation_5m_tokens - "integer", // cache_creation_1h_tokens - "integer", // image_output_tokens - "numeric", // image_output_cost - "numeric", // input_cost - "numeric", // output_cost - "numeric", // cache_creation_cost - "numeric", // cache_read_cost - "numeric", // total_cost - "numeric", // actual_cost - "numeric", // rate_multiplier - "numeric", // account_rate_multiplier - "smallint", // billing_type - "smallint", // request_type - "boolean", // stream - "boolean", // openai_ws_mode - "integer", // duration_ms - "integer", // first_token_ms - "text", // user_agent - "text", // ip_address - "integer", // image_count - "text", // image_size - "text", // image_input_size - "text", // image_output_size - "text", // image_size_source - "jsonb", // image_size_breakdown - "text", // service_tier - "text", // reasoning_effort - "text", // inbound_endpoint - "text", // upstream_endpoint - "boolean", // cache_ttl_overridden - "bigint", // channel_id - "text", // model_mapping_chain - "text", // billing_tier - "text", // billing_mode - "numeric", // account_stats_cost - "timestamptz", // created_at -} - const rawUsageLogModelColumn = "model" // rawUsageLogModelColumn preserves the exact stored usage_logs.model semantics for direct filters. @@ -160,7 +80,9 @@ func appendUsageLogBillingModeWhereConditionWithAlias(conditions []string, args placeholder := fmt.Sprintf("$%d", len(args)+1) switch service.BillingMode(mode) { case service.BillingModeImage: - conditions = append(conditions, fmt.Sprintf("(%s = %s OR COALESCE(%s, 0) > 0)", column("billing_mode"), placeholder, column("image_count"))) + conditions = append(conditions, fmt.Sprintf("(%s = %s OR ((%s IS NULL OR %s = '') AND COALESCE(%s, 0) > 0))", column("billing_mode"), placeholder, column("billing_mode"), column("billing_mode"), column("image_count"))) + case service.BillingModeVideo: + conditions = append(conditions, fmt.Sprintf("%s = %s", column("billing_mode"), placeholder)) case service.BillingModeToken: conditions = append(conditions, fmt.Sprintf("(%s = %s OR ((%s IS NULL OR %s = '') AND COALESCE(%s, 0) <= 0))", column("billing_mode"), placeholder, column("billing_mode"), column("billing_mode"), column("image_count"))) default: @@ -226,68 +148,6 @@ type usageLogRepository struct { bestEffortRecent *gocache.Cache } -const ( - usageLogCreateBatchMaxSize = 64 - usageLogCreateBatchWindow = 3 * time.Millisecond - usageLogCreateBatchQueueCap = 4096 - usageLogCreateCancelWait = 2 * time.Second - - usageLogBestEffortBatchMaxSize = 256 - usageLogBestEffortBatchWindow = 20 * time.Millisecond - usageLogBestEffortBatchQueueCap = 32768 - usageLogBestEffortRecentTTL = 30 * time.Second -) - -type usageLogCreateRequest struct { - log *service.UsageLog - prepared usageLogInsertPrepared - shared *usageLogCreateShared - resultCh chan usageLogCreateResult -} - -type usageLogCreateResult struct { - inserted bool - err error -} - -type usageLogBestEffortRequest struct { - prepared usageLogInsertPrepared - apiKeyID int64 - resultCh chan error -} - -type usageLogInsertPrepared struct { - createdAt time.Time - requestID string - rateMultiplier float64 - requestType int16 - args []any -} - -type usageLogBatchState struct { - ID int64 - CreatedAt time.Time -} - -type usageLogBatchRow struct { - RequestID string `json:"request_id"` - APIKeyID int64 `json:"api_key_id"` - ID int64 `json:"id"` - CreatedAt time.Time `json:"created_at"` - Inserted bool `json:"inserted"` -} - -type usageLogCreateShared struct { - state atomic.Int32 -} - -const ( - usageLogCreateStateQueued int32 = iota - usageLogCreateStateProcessing - usageLogCreateStateCompleted - usageLogCreateStateCanceled -) - func NewUsageLogRepository(client *dbent.Client, sqlDB *sql.DB) service.UsageLogRepository { return newUsageLogRepositoryWithSQL(client, sqlDB) } @@ -302,4284 +162,6 @@ func newUsageLogRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *usage return repo } -// getPerformanceStats 获取 RPM 和 TPM(近5分钟平均值,可选按用户过滤) -func (r *usageLogRepository) getPerformanceStats(ctx context.Context, userID int64) (rpm, tpm int64, err error) { - fiveMinutesAgo := time.Now().Add(-5 * time.Minute) - query := ` - SELECT - COUNT(*) as request_count, - COALESCE(SUM(input_tokens + output_tokens), 0) as token_count - FROM usage_logs - WHERE created_at >= $1` - args := []any{fiveMinutesAgo} - if userID > 0 { - query += " AND user_id = $2" - args = append(args, userID) - } - - var requestCount int64 - var tokenCount int64 - if err := scanSingleRow(ctx, r.sql, query, args, &requestCount, &tokenCount); err != nil { - return 0, 0, err - } - return requestCount / 5, tokenCount / 5, nil -} - -func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog) (bool, error) { - if log == nil { - return false, nil - } - - if tx := dbent.TxFromContext(ctx); tx != nil { - return r.createSingle(ctx, tx.Client(), log) - } - requestID := strings.TrimSpace(log.RequestID) - if requestID == "" { - return r.createSingle(ctx, r.sql, log) - } - log.RequestID = requestID - return r.createBatched(ctx, log) -} - -func (r *usageLogRepository) CreateBestEffort(ctx context.Context, log *service.UsageLog) error { - if log == nil { - return nil - } - - if tx := dbent.TxFromContext(ctx); tx != nil { - _, err := r.createSingle(ctx, tx.Client(), log) - return err - } - if r.db == nil { - _, err := r.createSingle(ctx, r.sql, log) - return err - } - - r.ensureBestEffortBatcher() - if r.bestEffortBatchCh == nil { - _, err := r.createSingle(ctx, r.sql, log) - return err - } - - req := usageLogBestEffortRequest{ - prepared: prepareUsageLogInsert(log), - apiKeyID: log.APIKeyID, - resultCh: make(chan error, 1), - } - if key, ok := r.bestEffortRecentKey(req.prepared.requestID, req.apiKeyID); ok { - if _, exists := r.bestEffortRecent.Get(key); exists { - return nil - } - } - - // 队列满时阻塞等待而非立即丢弃:批处理器持续排空队列,短暂等待即可入队。 - // 立即丢弃会造成“已扣费但无 usage_log”的永久数据缺口(issue #3656); - // 阻塞上限由调用方 ctx 期限约束,超时后由上层同步兜底。 - select { - case r.bestEffortBatchCh <- req: - case <-ctx.Done(): - return service.MarkUsageLogCreateDropped(ctx.Err()) - } - - select { - case err := <-req.resultCh: - return err - case <-ctx.Done(): - return service.MarkUsageLogCreateDropped(ctx.Err()) - } -} - -func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, log *service.UsageLog) (bool, error) { - prepared := prepareUsageLogInsert(log) - if sqlq == nil { - sqlq = r.sql - } - if ctx != nil && ctx.Err() != nil { - return false, service.MarkUsageLogCreateNotPersisted(ctx.Err()) - } - - query := ` - INSERT INTO usage_logs ( - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, $7, - $8, $9, - $10, $11, $12, $13, - $14, $15, $16, $17, - $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50 - ) - ON CONFLICT (request_id, api_key_id) DO NOTHING - RETURNING id, created_at - ` - - if err := scanSingleRow(ctx, sqlq, query, prepared.args, &log.ID, &log.CreatedAt); err != nil { - if errors.Is(err, sql.ErrNoRows) && prepared.requestID != "" { - selectQuery := "SELECT id, created_at FROM usage_logs WHERE request_id = $1 AND api_key_id = $2" - if err := scanSingleRow(ctx, sqlq, selectQuery, []any{prepared.requestID, log.APIKeyID}, &log.ID, &log.CreatedAt); err != nil { - return false, err - } - log.RateMultiplier = prepared.rateMultiplier - return false, nil - } else { - return false, err - } - } - log.RateMultiplier = prepared.rateMultiplier - return true, nil -} - -func (r *usageLogRepository) createBatched(ctx context.Context, log *service.UsageLog) (bool, error) { - if r.db == nil { - return r.createSingle(ctx, r.sql, log) - } - r.ensureCreateBatcher() - if r.createBatchCh == nil { - return r.createSingle(ctx, r.sql, log) - } - - req := usageLogCreateRequest{ - log: log, - prepared: prepareUsageLogInsert(log), - shared: &usageLogCreateShared{}, - resultCh: make(chan usageLogCreateResult, 1), - } - - // 队列满时阻塞等待而非立即报错:本路径是 best-effort 丢弃后的最后兜底, - // 立即失败会让日志永久丢失;阻塞上限由调用方 ctx 期限约束。 - select { - case r.createBatchCh <- req: - case <-ctx.Done(): - return false, service.MarkUsageLogCreateNotPersisted(ctx.Err()) - } - - select { - case res := <-req.resultCh: - return res.inserted, res.err - case <-ctx.Done(): - if req.shared != nil && req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateCanceled) { - return false, service.MarkUsageLogCreateNotPersisted(ctx.Err()) - } - timer := time.NewTimer(usageLogCreateCancelWait) - defer timer.Stop() - select { - case res := <-req.resultCh: - return res.inserted, res.err - case <-timer.C: - return false, ctx.Err() - } - } -} - -func (r *usageLogRepository) ensureCreateBatcher() { - if r == nil || r.db == nil { - return - } - // nil 检查必须在 Once 内部:在外层做无同步快路径读会与 Once 内的写构成数据竞争。 - r.createBatchOnce.Do(func() { - if r.createBatchCh == nil { - r.createBatchCh = make(chan usageLogCreateRequest, usageLogCreateBatchQueueCap) - go r.runCreateBatcher(r.db) - } - }) -} - -func (r *usageLogRepository) ensureBestEffortBatcher() { - if r == nil || r.db == nil { - return - } - // 同 ensureCreateBatcher:nil 检查放在 Once 内部以避免数据竞争。 - r.bestEffortBatchOnce.Do(func() { - if r.bestEffortBatchCh == nil { - r.bestEffortBatchCh = make(chan usageLogBestEffortRequest, usageLogBestEffortBatchQueueCap) - go r.runBestEffortBatcher(r.db) - } - }) -} - -func (r *usageLogRepository) runCreateBatcher(db *sql.DB) { - for { - first, ok := <-r.createBatchCh - if !ok { - return - } - - batch := make([]usageLogCreateRequest, 0, usageLogCreateBatchMaxSize) - batch = append(batch, first) - - timer := time.NewTimer(usageLogCreateBatchWindow) - batchLoop: - for len(batch) < usageLogCreateBatchMaxSize { - select { - case req, ok := <-r.createBatchCh: - if !ok { - break batchLoop - } - batch = append(batch, req) - case <-timer.C: - break batchLoop - } - } - if !timer.Stop() { - select { - case <-timer.C: - default: - } - } - - r.flushCreateBatch(db, batch) - } -} - -func (r *usageLogRepository) runBestEffortBatcher(db *sql.DB) { - for { - first, ok := <-r.bestEffortBatchCh - if !ok { - return - } - - batch := make([]usageLogBestEffortRequest, 0, usageLogBestEffortBatchMaxSize) - batch = append(batch, first) - - timer := time.NewTimer(usageLogBestEffortBatchWindow) - bestEffortLoop: - for len(batch) < usageLogBestEffortBatchMaxSize { - select { - case req, ok := <-r.bestEffortBatchCh: - if !ok { - break bestEffortLoop - } - batch = append(batch, req) - case <-timer.C: - break bestEffortLoop - } - } - if !timer.Stop() { - select { - case <-timer.C: - default: - } - } - - r.flushBestEffortBatch(db, batch) - } -} - -func (r *usageLogRepository) flushCreateBatch(db *sql.DB, batch []usageLogCreateRequest) { - if len(batch) == 0 { - return - } - - uniqueOrder := make([]string, 0, len(batch)) - preparedByKey := make(map[string]usageLogInsertPrepared, len(batch)) - requestsByKey := make(map[string][]usageLogCreateRequest, len(batch)) - fallback := make([]usageLogCreateRequest, 0) - - for _, req := range batch { - if req.log == nil { - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) - continue - } - if req.shared != nil && !req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateProcessing) { - if req.shared.state.Load() == usageLogCreateStateCanceled { - completeUsageLogCreateRequest(req, usageLogCreateResult{ - inserted: false, - err: service.MarkUsageLogCreateNotPersisted(context.Canceled), - }) - continue - } - } - prepared := req.prepared - if prepared.requestID == "" { - fallback = append(fallback, req) - continue - } - key := usageLogBatchKey(prepared.requestID, req.log.APIKeyID) - if _, exists := requestsByKey[key]; !exists { - uniqueOrder = append(uniqueOrder, key) - preparedByKey[key] = prepared - } - requestsByKey[key] = append(requestsByKey[key], req) - } - - if len(uniqueOrder) > 0 { - insertedMap, stateMap, safeFallback, err := r.batchInsertUsageLogs(db, uniqueOrder, preparedByKey) - if err != nil { - if safeFallback { - for _, key := range uniqueOrder { - fallback = append(fallback, requestsByKey[key]...) - } - } else { - for _, key := range uniqueOrder { - reqs := requestsByKey[key] - state, hasState := stateMap[key] - inserted := insertedMap[key] - for idx, req := range reqs { - req.log.RateMultiplier = preparedByKey[key].rateMultiplier - if hasState { - req.log.ID = state.ID - req.log.CreatedAt = state.CreatedAt - } - switch { - case inserted && idx == 0: - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: true, err: nil}) - case inserted: - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) - case hasState: - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) - case idx == 0: - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: err}) - default: - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) - } - } - } - } - } else { - for _, key := range uniqueOrder { - reqs := requestsByKey[key] - state, ok := stateMap[key] - if !ok { - for _, req := range reqs { - completeUsageLogCreateRequest(req, usageLogCreateResult{ - inserted: false, - err: fmt.Errorf("usage log batch state missing for key=%s", key), - }) - } - continue - } - for idx, req := range reqs { - req.log.ID = state.ID - req.log.CreatedAt = state.CreatedAt - req.log.RateMultiplier = preparedByKey[key].rateMultiplier - completeUsageLogCreateRequest(req, usageLogCreateResult{ - inserted: idx == 0 && insertedMap[key], - err: nil, - }) - } - } - } - } - - if len(fallback) == 0 { - return - } - - for _, req := range fallback { - fallbackCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - inserted, err := r.createSingle(fallbackCtx, db, req.log) - cancel() - completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: inserted, err: err}) - } -} - -func (r *usageLogRepository) flushBestEffortBatch(db *sql.DB, batch []usageLogBestEffortRequest) { - if len(batch) == 0 { - return - } - - type bestEffortGroup struct { - prepared usageLogInsertPrepared - apiKeyID int64 - key string - reqs []usageLogBestEffortRequest - } - - groupsByKey := make(map[string]*bestEffortGroup, len(batch)) - groupOrder := make([]*bestEffortGroup, 0, len(batch)) - preparedList := make([]usageLogInsertPrepared, 0, len(batch)) - - for idx, req := range batch { - prepared := req.prepared - key := fmt.Sprintf("__best_effort_%d", idx) - if prepared.requestID != "" { - key = usageLogBatchKey(prepared.requestID, req.apiKeyID) - } - group, exists := groupsByKey[key] - if !exists { - group = &bestEffortGroup{ - prepared: prepared, - apiKeyID: req.apiKeyID, - key: key, - } - groupsByKey[key] = group - groupOrder = append(groupOrder, group) - preparedList = append(preparedList, prepared) - } - group.reqs = append(group.reqs, req) - } - - if len(preparedList) == 0 { - for _, req := range batch { - sendUsageLogBestEffortResult(req.resultCh, nil) - } - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - query, args := buildUsageLogBestEffortInsertQuery(preparedList) - if _, err := db.ExecContext(ctx, query, args...); err != nil { - logger.LegacyPrintf("repository.usage_log", "best-effort batch insert failed: %v", err) - for _, group := range groupOrder { - singleErr := execUsageLogInsertNoResult(ctx, db, group.prepared) - if singleErr != nil { - logger.LegacyPrintf("repository.usage_log", "best-effort single fallback insert failed: %v", singleErr) - } else if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil { - r.bestEffortRecent.SetDefault(group.key, struct{}{}) - } - for _, req := range group.reqs { - sendUsageLogBestEffortResult(req.resultCh, singleErr) - } - } - return - } - for _, group := range groupOrder { - if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil { - r.bestEffortRecent.SetDefault(group.key, struct{}{}) - } - for _, req := range group.reqs { - sendUsageLogBestEffortResult(req.resultCh, nil) - } - } -} - -func sendUsageLogBestEffortResult(ch chan error, err error) { - if ch == nil { - return - } - select { - case ch <- err: - default: - } -} - -func completeUsageLogCreateRequest(req usageLogCreateRequest, res usageLogCreateResult) { - if req.shared != nil { - req.shared.state.Store(usageLogCreateStateCompleted) - } - sendUsageLogCreateResult(req.resultCh, res) -} - -func (r *usageLogRepository) batchInsertUsageLogs(db *sql.DB, keys []string, preparedByKey map[string]usageLogInsertPrepared) (map[string]bool, map[string]usageLogBatchState, bool, error) { - if len(keys) == 0 { - return map[string]bool{}, map[string]usageLogBatchState{}, false, nil - } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - query, args := buildUsageLogBatchInsertQuery(keys, preparedByKey) - var payload []byte - if err := db.QueryRowContext(ctx, query, args...).Scan(&payload); err != nil { - return nil, nil, true, err - } - var rows []usageLogBatchRow - if err := json.Unmarshal(payload, &rows); err != nil { - return nil, nil, false, err - } - insertedMap := make(map[string]bool, len(keys)) - stateMap := make(map[string]usageLogBatchState, len(keys)) - for _, row := range rows { - key := usageLogBatchKey(row.RequestID, row.APIKeyID) - insertedMap[key] = row.Inserted - stateMap[key] = usageLogBatchState{ - ID: row.ID, - CreatedAt: row.CreatedAt, - } - } - if len(stateMap) != len(keys) { - return insertedMap, stateMap, false, fmt.Errorf("usage log batch state count mismatch: got=%d want=%d", len(stateMap), len(keys)) - } - return insertedMap, stateMap, false, nil -} - -func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usageLogInsertPrepared) (string, []any) { - var query strings.Builder - _, _ = query.WriteString(` - WITH input ( - input_idx, - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - ) AS (VALUES `) - - args := make([]any, 0, len(keys)*50) - argPos := 1 - for idx, key := range keys { - if idx > 0 { - _, _ = query.WriteString(",") - } - _, _ = query.WriteString("(") - _, _ = query.WriteString("$") - _, _ = query.WriteString(strconv.Itoa(argPos)) - args = append(args, idx) - argPos++ - prepared := preparedByKey[key] - for i := 0; i < len(prepared.args); i++ { - _, _ = query.WriteString(",") - _, _ = query.WriteString("$") - _, _ = query.WriteString(strconv.Itoa(argPos)) - if i < len(usageLogInsertArgTypes) { - _, _ = query.WriteString("::") - _, _ = query.WriteString(usageLogInsertArgTypes[i]) - } - argPos++ - } - _, _ = query.WriteString(")") - args = append(args, prepared.args...) - } - _, _ = query.WriteString(` - ), - inserted AS ( - INSERT INTO usage_logs ( - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - ) - SELECT - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - FROM input - ON CONFLICT (request_id, api_key_id) DO NOTHING - RETURNING request_id, api_key_id, id, created_at - ), - resolved AS ( - SELECT - input.input_idx, - input.request_id, - input.api_key_id, - COALESCE(inserted.id, existing.id) AS id, - COALESCE(inserted.created_at, existing.created_at) AS created_at, - (inserted.id IS NOT NULL) AS inserted - FROM input - LEFT JOIN inserted - ON inserted.request_id = input.request_id - AND inserted.api_key_id = input.api_key_id - LEFT JOIN usage_logs existing - ON existing.request_id = input.request_id - AND existing.api_key_id = input.api_key_id - ) - SELECT COALESCE( - json_agg( - json_build_object( - 'request_id', resolved.request_id, - 'api_key_id', resolved.api_key_id, - 'id', resolved.id, - 'created_at', resolved.created_at, - 'inserted', resolved.inserted - ) - ORDER BY resolved.input_idx - ), - '[]'::json - ) - FROM resolved - `) - return query.String(), args -} - -func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (string, []any) { - var query strings.Builder - _, _ = query.WriteString(` - WITH input ( - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - ) AS (VALUES `) - - args := make([]any, 0, len(preparedList)*50) - argPos := 1 - for idx, prepared := range preparedList { - if idx > 0 { - _, _ = query.WriteString(",") - } - _, _ = query.WriteString("(") - for i := 0; i < len(prepared.args); i++ { - if i > 0 { - _, _ = query.WriteString(",") - } - _, _ = query.WriteString("$") - _, _ = query.WriteString(strconv.Itoa(argPos)) - if i < len(usageLogInsertArgTypes) { - _, _ = query.WriteString("::") - _, _ = query.WriteString(usageLogInsertArgTypes[i]) - } - argPos++ - } - _, _ = query.WriteString(")") - args = append(args, prepared.args...) - } - - _, _ = query.WriteString(` - ) - INSERT INTO usage_logs ( - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - ) - SELECT - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - FROM input - ON CONFLICT (request_id, api_key_id) DO NOTHING - `) - - return query.String(), args -} - -func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared usageLogInsertPrepared) error { - _, err := sqlq.ExecContext(ctx, ` - INSERT INTO usage_logs ( - user_id, - api_key_id, - account_id, - request_id, - model, - requested_model, - upstream_model, - group_id, - subscription_id, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - cache_creation_5m_tokens, - cache_creation_1h_tokens, - image_output_tokens, - image_output_cost, - input_cost, - output_cost, - cache_creation_cost, - cache_read_cost, - total_cost, - actual_cost, - rate_multiplier, - account_rate_multiplier, - billing_type, - request_type, - stream, - openai_ws_mode, - duration_ms, - first_token_ms, - user_agent, - ip_address, - image_count, - image_size, - image_input_size, - image_output_size, - image_size_source, - image_size_breakdown, - service_tier, - reasoning_effort, - inbound_endpoint, - upstream_endpoint, - cache_ttl_overridden, - channel_id, - model_mapping_chain, - billing_tier, - billing_mode, - account_stats_cost, - created_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, $7, - $8, $9, - $10, $11, $12, $13, - $14, $15, $16, $17, - $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50 - ) - ON CONFLICT (request_id, api_key_id) DO NOTHING - `, prepared.args...) - return err -} - -func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { - createdAt := log.CreatedAt - if createdAt.IsZero() { - createdAt = time.Now() - } - - requestID := strings.TrimSpace(log.RequestID) - log.RequestID = requestID - - rateMultiplier := log.RateMultiplier - log.SyncRequestTypeAndLegacyFields() - requestType := int16(log.RequestType) - - groupID := nullInt64(log.GroupID) - subscriptionID := nullInt64(log.SubscriptionID) - duration := nullInt(log.DurationMs) - firstToken := nullInt(log.FirstTokenMs) - userAgent := nullString(log.UserAgent) - ipAddress := nullString(log.IPAddress) - imageSize := nullString(log.ImageSize) - imageInputSize := nullString(log.ImageInputSize) - imageOutputSize := nullString(log.ImageOutputSize) - imageSizeSource := nullString(log.ImageSizeSource) - imageSizeBreakdown := nullStringIntMapJSON(log.ImageSizeBreakdown) - serviceTier := nullString(log.ServiceTier) - reasoningEffort := nullString(log.ReasoningEffort) - inboundEndpoint := nullString(log.InboundEndpoint) - upstreamEndpoint := nullString(log.UpstreamEndpoint) - channelID := nullInt64(log.ChannelID) - modelMappingChain := nullString(log.ModelMappingChain) - billingTier := nullString(log.BillingTier) - billingMode := nullString(log.BillingMode) - requestedModel := strings.TrimSpace(log.RequestedModel) - if requestedModel == "" { - requestedModel = strings.TrimSpace(log.Model) - } - upstreamModel := nullString(log.UpstreamModel) - - var requestIDArg any - if requestID != "" { - requestIDArg = requestID - } - - return usageLogInsertPrepared{ - createdAt: createdAt, - requestID: requestID, - rateMultiplier: rateMultiplier, - requestType: requestType, - args: []any{ - log.UserID, - log.APIKeyID, - log.AccountID, - requestIDArg, - log.Model, - nullString(&requestedModel), - upstreamModel, - groupID, - subscriptionID, - log.InputTokens, - log.OutputTokens, - log.CacheCreationTokens, - log.CacheReadTokens, - log.CacheCreation5mTokens, - log.CacheCreation1hTokens, - log.ImageOutputTokens, - log.ImageOutputCost, - log.InputCost, - log.OutputCost, - log.CacheCreationCost, - log.CacheReadCost, - log.TotalCost, - log.ActualCost, - rateMultiplier, - log.AccountRateMultiplier, - log.BillingType, - requestType, - log.Stream, - log.OpenAIWSMode, - duration, - firstToken, - userAgent, - ipAddress, - log.ImageCount, - imageSize, - imageInputSize, - imageOutputSize, - imageSizeSource, - imageSizeBreakdown, - serviceTier, - reasoningEffort, - inboundEndpoint, - upstreamEndpoint, - log.CacheTTLOverridden, - channelID, - modelMappingChain, - billingTier, - billingMode, - log.AccountStatsCost, // account_stats_cost - createdAt, - }, - } -} - -func usageLogBatchKey(requestID string, apiKeyID int64) string { - return requestID + "\x1f" + strconv.FormatInt(apiKeyID, 10) -} - -func sendUsageLogCreateResult(ch chan usageLogCreateResult, res usageLogCreateResult) { - if ch == nil { - return - } - select { - case ch <- res: - default: - } -} - -func (r *usageLogRepository) bestEffortRecentKey(requestID string, apiKeyID int64) (string, bool) { - requestID = strings.TrimSpace(requestID) - if requestID == "" || r == nil || r.bestEffortRecent == nil { - return "", false - } - return usageLogBatchKey(requestID, apiKeyID), true -} - -func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) { - query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1" - rows, err := r.sql.QueryContext(ctx, query, id) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - log = nil - } - }() - if !rows.Next() { - if err = rows.Err(); err != nil { - return nil, err - } - return nil, service.ErrUsageLogNotFound - } - log, err = scanUsageLog(rows) - if err != nil { - return nil, err - } - if err = rows.Err(); err != nil { - return nil, err - } - return log, nil -} - -func (r *usageLogRepository) ListByUser(ctx context.Context, userID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { - return r.listUsageLogsWithPagination(ctx, "WHERE user_id = $1", []any{userID}, params) -} - -func (r *usageLogRepository) ListByAPIKey(ctx context.Context, apiKeyID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { - return r.listUsageLogsWithPagination(ctx, "WHERE api_key_id = $1", []any{apiKeyID}, params) -} - -// UserStats 用户使用统计 -type UserStats struct { - TotalRequests int64 `json:"total_requests"` - TotalTokens int64 `json:"total_tokens"` - TotalCost float64 `json:"total_cost"` - InputTokens int64 `json:"input_tokens"` - OutputTokens int64 `json:"output_tokens"` - CacheReadTokens int64 `json:"cache_read_tokens"` -} - -func (r *usageLogRepository) GetUserStats(ctx context.Context, userID int64, startTime, endTime time.Time) (*UserStats, error) { - query := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(actual_cost), 0) as total_cost, - COALESCE(SUM(input_tokens), 0) as input_tokens, - COALESCE(SUM(output_tokens), 0) as output_tokens, - COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens - FROM usage_logs - WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 - ` - - stats := &UserStats{} - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{userID, startTime, endTime}, - &stats.TotalRequests, - &stats.TotalTokens, - &stats.TotalCost, - &stats.InputTokens, - &stats.OutputTokens, - &stats.CacheReadTokens, - ); err != nil { - return nil, err - } - return stats, nil -} - -// DashboardStats 仪表盘统计 -type DashboardStats = usagestats.DashboardStats - -func (r *usageLogRepository) GetDashboardStats(ctx context.Context) (*DashboardStats, error) { - stats := &DashboardStats{} - now := timezone.Now() - todayStart := timezone.Today() - - if err := r.fillDashboardEntityStats(ctx, stats, todayStart, now); err != nil { - return nil, err - } - if err := r.fillDashboardUsageStatsAggregated(ctx, stats, todayStart, now); err != nil { - return nil, err - } - - rpm, tpm, err := r.getPerformanceStats(ctx, 0) - if err != nil { - return nil, err - } - stats.Rpm = rpm - stats.Tpm = tpm - - return stats, nil -} - -func (r *usageLogRepository) GetDashboardStatsWithRange(ctx context.Context, start, end time.Time) (*DashboardStats, error) { - startUTC := start.UTC() - endUTC := end.UTC() - if !endUTC.After(startUTC) { - return nil, errors.New("统计时间范围无效") - } - - stats := &DashboardStats{} - now := timezone.Now() - todayStart := timezone.Today() - - if err := r.fillDashboardEntityStats(ctx, stats, todayStart, now); err != nil { - return nil, err - } - if err := r.fillDashboardUsageStatsFromUsageLogs(ctx, stats, startUTC, endUTC, todayStart, now); err != nil { - return nil, err - } - - rpm, tpm, err := r.getPerformanceStats(ctx, 0) - if err != nil { - return nil, err - } - stats.Rpm = rpm - stats.Tpm = tpm - - return stats, nil -} - -func (r *usageLogRepository) fillDashboardEntityStats(ctx context.Context, stats *DashboardStats, todayUTC, now time.Time) error { - userStatsQuery := ` - SELECT - COUNT(*) as total_users, - COUNT(CASE WHEN created_at >= $1 THEN 1 END) as today_new_users - FROM users - WHERE deleted_at IS NULL - ` - if err := scanSingleRow( - ctx, - r.sql, - userStatsQuery, - []any{todayUTC}, - &stats.TotalUsers, - &stats.TodayNewUsers, - ); err != nil { - return err - } - - apiKeyStatsQuery := ` - SELECT - COUNT(*) as total_api_keys, - COUNT(CASE WHEN status = $1 THEN 1 END) as active_api_keys - FROM api_keys - WHERE deleted_at IS NULL - ` - if err := scanSingleRow( - ctx, - r.sql, - apiKeyStatsQuery, - []any{service.StatusActive}, - &stats.TotalAPIKeys, - &stats.ActiveAPIKeys, - ); err != nil { - return err - } - - accountStatsQuery := ` - SELECT - COUNT(*) as total_accounts, - COUNT(CASE WHEN status = $1 AND schedulable = true THEN 1 END) as normal_accounts, - COUNT(CASE WHEN status = $2 THEN 1 END) as error_accounts, - COUNT(CASE WHEN rate_limited_at IS NOT NULL AND rate_limit_reset_at > $3 THEN 1 END) as ratelimit_accounts, - COUNT(CASE WHEN overload_until IS NOT NULL AND overload_until > $4 THEN 1 END) as overload_accounts - FROM accounts - WHERE deleted_at IS NULL - ` - if err := scanSingleRow( - ctx, - r.sql, - accountStatsQuery, - []any{service.StatusActive, service.StatusError, now, now}, - &stats.TotalAccounts, - &stats.NormalAccounts, - &stats.ErrorAccounts, - &stats.RateLimitAccounts, - &stats.OverloadAccounts, - ); err != nil { - return err - } - - return nil -} - -func (r *usageLogRepository) fillDashboardUsageStatsAggregated(ctx context.Context, stats *DashboardStats, todayUTC, now time.Time) error { - totalStatsQuery := ` - SELECT - COALESCE(SUM(total_requests), 0) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(SUM(account_cost), 0) as total_account_cost, - COALESCE(SUM(total_duration_ms), 0) as total_duration_ms - FROM usage_dashboard_daily - ` - var totalDurationMs int64 - if err := scanSingleRow( - ctx, - r.sql, - totalStatsQuery, - nil, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.TotalAccountCost, - &totalDurationMs, - ); err != nil { - return err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens - if stats.TotalRequests > 0 { - stats.AverageDurationMs = float64(totalDurationMs) / float64(stats.TotalRequests) - } - - todayStatsQuery := ` - SELECT - total_requests as today_requests, - input_tokens as today_input_tokens, - output_tokens as today_output_tokens, - cache_creation_tokens as today_cache_creation_tokens, - cache_read_tokens as today_cache_read_tokens, - total_cost as today_cost, - actual_cost as today_actual_cost, - account_cost as today_account_cost, - active_users as active_users - FROM usage_dashboard_daily - WHERE bucket_date = $1::date - ` - if err := scanSingleRow( - ctx, - r.sql, - todayStatsQuery, - []any{todayUTC}, - &stats.TodayRequests, - &stats.TodayInputTokens, - &stats.TodayOutputTokens, - &stats.TodayCacheCreationTokens, - &stats.TodayCacheReadTokens, - &stats.TodayCost, - &stats.TodayActualCost, - &stats.TodayAccountCost, - &stats.ActiveUsers, - ); err != nil { - if err != sql.ErrNoRows { - return err - } - } - stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens - - hourlyActiveQuery := ` - SELECT active_users - FROM usage_dashboard_hourly - WHERE bucket_start = $1 - ` - hourStart := now.In(timezone.Location()).Truncate(time.Hour) - if err := scanSingleRow(ctx, r.sql, hourlyActiveQuery, []any{hourStart}, &stats.HourlyActiveUsers); err != nil { - if err != sql.ErrNoRows { - return err - } - } - - return nil -} - -func (r *usageLogRepository) fillDashboardUsageStatsFromUsageLogs(ctx context.Context, stats *DashboardStats, startUTC, endUTC, todayUTC, now time.Time) error { - todayEnd := todayUTC.Add(24 * time.Hour) - combinedStatsQuery := ` - WITH scoped AS ( - SELECT - created_at, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost, - actual_cost, - COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1) AS account_cost, - COALESCE(duration_ms, 0) AS duration_ms - FROM usage_logs - WHERE created_at >= LEAST($1::timestamptz, $3::timestamptz) - AND created_at < GREATEST($2::timestamptz, $4::timestamptz) - ) - SELECT - COUNT(*) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz) AS total_requests, - COALESCE(SUM(input_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_input_tokens, - COALESCE(SUM(output_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_output_tokens, - COALESCE(SUM(cache_creation_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_cache_read_tokens, - COALESCE(SUM(total_cost) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_cost, - COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_actual_cost, - COALESCE(SUM(account_cost) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_account_cost, - COALESCE(SUM(duration_ms) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_duration_ms, - COUNT(*) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz) AS today_requests, - COALESCE(SUM(input_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_input_tokens, - COALESCE(SUM(output_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_output_tokens, - COALESCE(SUM(cache_creation_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_cache_read_tokens, - COALESCE(SUM(total_cost) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_cost, - COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_actual_cost, - COALESCE(SUM(account_cost) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_account_cost - FROM scoped - ` - var totalDurationMs int64 - if err := scanSingleRow( - ctx, - r.sql, - combinedStatsQuery, - []any{startUTC, endUTC, todayUTC, todayEnd}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.TotalAccountCost, - &totalDurationMs, - &stats.TodayRequests, - &stats.TodayInputTokens, - &stats.TodayOutputTokens, - &stats.TodayCacheCreationTokens, - &stats.TodayCacheReadTokens, - &stats.TodayCost, - &stats.TodayActualCost, - &stats.TodayAccountCost, - ); err != nil { - return err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens - if stats.TotalRequests > 0 { - stats.AverageDurationMs = float64(totalDurationMs) / float64(stats.TotalRequests) - } - - stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens - - hourStart := now.UTC().Truncate(time.Hour) - hourEnd := hourStart.Add(time.Hour) - activeUsersQuery := ` - WITH scoped AS ( - SELECT user_id, created_at - FROM usage_logs - WHERE created_at >= LEAST($1::timestamptz, $3::timestamptz) - AND created_at < GREATEST($2::timestamptz, $4::timestamptz) - ) - SELECT - COUNT(DISTINCT CASE WHEN created_at >= $1::timestamptz AND created_at < $2::timestamptz THEN user_id END) AS active_users, - COUNT(DISTINCT CASE WHEN created_at >= $3::timestamptz AND created_at < $4::timestamptz THEN user_id END) AS hourly_active_users - FROM scoped - ` - if err := scanSingleRow(ctx, r.sql, activeUsersQuery, []any{todayUTC, todayEnd, hourStart, hourEnd}, &stats.ActiveUsers, &stats.HourlyActiveUsers); err != nil { - return err - } - - return nil -} - -func (r *usageLogRepository) ListByAccount(ctx context.Context, accountID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { - return r.listUsageLogsWithPagination(ctx, "WHERE account_id = $1", []any{accountID}, params) -} - -func (r *usageLogRepository) ListByUserAndTimeRange(ctx context.Context, userID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { - query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000" - logs, err := r.queryUsageLogs(ctx, query, userID, startTime, endTime) - return logs, nil, err -} - -// GetUserStatsAggregated returns aggregated usage statistics for a user using database-level aggregation -func (r *usageLogRepository) GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) { - query := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms - FROM usage_logs - WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 - ` - - var stats usagestats.UsageStats - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{userID, startTime, endTime}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens - return &stats, nil -} - -// GetAPIKeyStatsAggregated returns aggregated usage statistics for an API key using database-level aggregation -func (r *usageLogRepository) GetAPIKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) { - query := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms - FROM usage_logs - WHERE api_key_id = $1 AND created_at >= $2 AND created_at < $3 - ` - - var stats usagestats.UsageStats - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{apiKeyID, startTime, endTime}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens - return &stats, nil -} - -// GetAccountStatsAggregated 使用 SQL 聚合统计账号使用数据 -// -// 性能优化说明: -// 原实现先查询所有日志记录,再在应用层循环计算统计值: -// 1. 需要传输大量数据到应用层 -// 2. 应用层循环计算增加 CPU 和内存开销 -// -// 新实现使用 SQL 聚合函数: -// 1. 在数据库层完成 COUNT/SUM/AVG 计算 -// 2. 只返回单行聚合结果,大幅减少数据传输量 -// 3. 利用数据库索引优化聚合查询性能 -func (r *usageLogRepository) GetAccountStatsAggregated(ctx context.Context, accountID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) { - query := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms - FROM usage_logs - WHERE account_id = $1 AND created_at >= $2 AND created_at < $3 - ` - - var stats usagestats.UsageStats - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{accountID, startTime, endTime}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens - return &stats, nil -} - -// GetModelStatsAggregated 使用 SQL 聚合统计模型使用数据 -// 性能优化:数据库层聚合计算,避免应用层循环统计 -func (r *usageLogRepository) GetModelStatsAggregated(ctx context.Context, modelName string, startTime, endTime time.Time) (*usagestats.UsageStats, error) { - query := fmt.Sprintf(` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms - FROM usage_logs - WHERE %s = $1 AND created_at >= $2 AND created_at < $3 - `, rawUsageLogModelColumn) - - var stats usagestats.UsageStats - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{modelName, startTime, endTime}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens - return &stats, nil -} - -// GetDailyStatsAggregated 使用 SQL 聚合统计用户的每日使用数据 -// 性能优化:使用 GROUP BY 在数据库层按日期分组聚合,避免应用层循环分组统计 -func (r *usageLogRepository) GetDailyStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (result []map[string]any, err error) { - tzName := resolveUsageStatsTimezone() - query := ` - SELECT - -- 使用应用时区分组,避免数据库会话时区导致日边界偏移。 - TO_CHAR(created_at AT TIME ZONE $4, 'YYYY-MM-DD') as date, - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms - FROM usage_logs - WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 - GROUP BY 1 - ORDER BY 1 - ` - - rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime, tzName) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - result = nil - } - }() - - result = make([]map[string]any, 0) - for rows.Next() { - var ( - date string - totalRequests int64 - totalInputTokens int64 - totalOutputTokens int64 - totalCacheTokens int64 - totalCost float64 - totalActualCost float64 - avgDurationMs float64 - ) - if err = rows.Scan( - &date, - &totalRequests, - &totalInputTokens, - &totalOutputTokens, - &totalCacheTokens, - &totalCost, - &totalActualCost, - &avgDurationMs, - ); err != nil { - return nil, err - } - result = append(result, map[string]any{ - "date": date, - "total_requests": totalRequests, - "total_input_tokens": totalInputTokens, - "total_output_tokens": totalOutputTokens, - "total_cache_tokens": totalCacheTokens, - "total_tokens": totalInputTokens + totalOutputTokens + totalCacheTokens, - "total_cost": totalCost, - "total_actual_cost": totalActualCost, - "average_duration_ms": avgDurationMs, - }) - } - - if err = rows.Err(); err != nil { - return nil, err - } - - return result, nil -} - -// resolveUsageStatsTimezone 获取用于 SQL 分组的时区名称。 -// 优先使用应用初始化的时区,其次尝试读取 TZ 环境变量,最后回落为 UTC。 -func resolveUsageStatsTimezone() string { - tzName := timezone.Name() - if tzName != "" && tzName != "Local" { - return tzName - } - if envTZ := strings.TrimSpace(os.Getenv("TZ")); envTZ != "" { - return envTZ - } - return "UTC" -} - -func (r *usageLogRepository) ListByAPIKeyAndTimeRange(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { - query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE api_key_id = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000" - logs, err := r.queryUsageLogs(ctx, query, apiKeyID, startTime, endTime) - return logs, nil, err -} - -func (r *usageLogRepository) ListByAccountAndTimeRange(ctx context.Context, accountID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { - query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE account_id = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000" - logs, err := r.queryUsageLogs(ctx, query, accountID, startTime, endTime) - return logs, nil, err -} - -func (r *usageLogRepository) ListByModelAndTimeRange(ctx context.Context, modelName string, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { - query := fmt.Sprintf("SELECT %s FROM usage_logs WHERE %s = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000", usageLogSelectColumns, rawUsageLogModelColumn) - logs, err := r.queryUsageLogs(ctx, query, modelName, startTime, endTime) - return logs, nil, err -} - -func (r *usageLogRepository) Delete(ctx context.Context, id int64) error { - _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_logs WHERE id = $1", id) - return err -} - -// GetAccountTodayStats 获取账号今日统计 -func (r *usageLogRepository) GetAccountTodayStats(ctx context.Context, accountID int64) (*usagestats.AccountStats, error) { - today := timezone.Today() - - query := ` - SELECT - COUNT(*) as requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, - COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost, - COALESCE(SUM(total_cost), 0) as standard_cost, - COALESCE(SUM(actual_cost), 0) as user_cost - FROM usage_logs - WHERE account_id = $1 AND created_at >= $2 - ` - - stats := &usagestats.AccountStats{} - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{accountID, today}, - &stats.Requests, - &stats.Tokens, - &stats.Cost, - &stats.StandardCost, - &stats.UserCost, - ); err != nil { - return nil, err - } - return stats, nil -} - -// GetAccountWindowStats 获取账号时间窗口内的统计 -func (r *usageLogRepository) GetAccountWindowStats(ctx context.Context, accountID int64, startTime time.Time) (*usagestats.AccountStats, error) { - query := ` - SELECT - COUNT(*) as requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, - COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost, - COALESCE(SUM(total_cost), 0) as standard_cost, - COALESCE(SUM(actual_cost), 0) as user_cost - FROM usage_logs - WHERE account_id = $1 AND created_at >= $2 - ` - - stats := &usagestats.AccountStats{} - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{accountID, startTime}, - &stats.Requests, - &stats.Tokens, - &stats.Cost, - &stats.StandardCost, - &stats.UserCost, - ); err != nil { - return nil, err - } - return stats, nil -} - -// GetAccountWindowStatsBatch 批量获取同一窗口起点下多个账号的统计数据。 -// 返回 map[accountID]*AccountStats,未命中的账号会返回零值统计,便于上层直接复用。 -func (r *usageLogRepository) GetAccountWindowStatsBatch(ctx context.Context, accountIDs []int64, startTime time.Time) (map[int64]*usagestats.AccountStats, error) { - result := make(map[int64]*usagestats.AccountStats, len(accountIDs)) - if len(accountIDs) == 0 { - return result, nil - } - - query := ` - SELECT - account_id, - COUNT(*) as requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, - COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost, - COALESCE(SUM(total_cost), 0) as standard_cost, - COALESCE(SUM(actual_cost), 0) as user_cost - FROM usage_logs - WHERE account_id = ANY($1) AND created_at >= $2 - GROUP BY account_id - ` - rows, err := r.sql.QueryContext(ctx, query, pq.Array(accountIDs), startTime) - if err != nil { - return nil, err - } - defer func() { _ = rows.Close() }() - - for rows.Next() { - var accountID int64 - stats := &usagestats.AccountStats{} - if err := rows.Scan( - &accountID, - &stats.Requests, - &stats.Tokens, - &stats.Cost, - &stats.StandardCost, - &stats.UserCost, - ); err != nil { - return nil, err - } - result[accountID] = stats - } - if err := rows.Err(); err != nil { - return nil, err - } - - for _, accountID := range accountIDs { - if _, ok := result[accountID]; !ok { - result[accountID] = &usagestats.AccountStats{} - } - } - return result, nil -} - -// GetGeminiUsageTotalsBatch 批量聚合 Gemini 账号在窗口内的 Pro/Flash 请求与用量。 -// 模型分类规则与 service.geminiModelClassFromName 一致:model 包含 flash/lite 视为 flash,其余视为 pro。 -func (r *usageLogRepository) GetGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, startTime, endTime time.Time) (map[int64]service.GeminiUsageTotals, error) { - result := make(map[int64]service.GeminiUsageTotals, len(accountIDs)) - if len(accountIDs) == 0 { - return result, nil - } - - query := ` - SELECT - account_id, - COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 1 ELSE 0 END), 0) AS flash_requests, - COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE 1 END), 0) AS pro_requests, - COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) ELSE 0 END), 0) AS flash_tokens, - COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) END), 0) AS pro_tokens, - COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN actual_cost ELSE 0 END), 0) AS flash_cost, - COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE actual_cost END), 0) AS pro_cost - FROM usage_logs - WHERE account_id = ANY($1) AND created_at >= $2 AND created_at < $3 - GROUP BY account_id - ` - rows, err := r.sql.QueryContext(ctx, query, pq.Array(accountIDs), startTime, endTime) - if err != nil { - return nil, err - } - defer func() { _ = rows.Close() }() - - for rows.Next() { - var accountID int64 - var totals service.GeminiUsageTotals - if err := rows.Scan( - &accountID, - &totals.FlashRequests, - &totals.ProRequests, - &totals.FlashTokens, - &totals.ProTokens, - &totals.FlashCost, - &totals.ProCost, - ); err != nil { - return nil, err - } - result[accountID] = totals - } - if err := rows.Err(); err != nil { - return nil, err - } - - for _, accountID := range accountIDs { - if _, ok := result[accountID]; !ok { - result[accountID] = service.GeminiUsageTotals{} - } - } - return result, nil -} - -// TrendDataPoint represents a single point in trend data -type TrendDataPoint = usagestats.TrendDataPoint - -// ModelStat represents usage statistics for a single model -type ModelStat = usagestats.ModelStat - -// UserUsageTrendPoint represents user usage trend data point -type UserUsageTrendPoint = usagestats.UserUsageTrendPoint - -// UserSpendingRankingItem represents a user spending ranking row. -type UserSpendingRankingItem = usagestats.UserSpendingRankingItem -type UserSpendingRankingResponse = usagestats.UserSpendingRankingResponse - -// APIKeyUsageTrendPoint represents API key usage trend data point -type APIKeyUsageTrendPoint = usagestats.APIKeyUsageTrendPoint - -// GetAPIKeyUsageTrend returns usage trend data grouped by API key and date -func (r *usageLogRepository) GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) (results []APIKeyUsageTrendPoint, err error) { - dateFormat := safeDateFormat(granularity) - - query := fmt.Sprintf(` - WITH top_keys AS ( - SELECT api_key_id - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - GROUP BY api_key_id - ORDER BY SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) DESC - LIMIT $3 - ) - SELECT - TO_CHAR(u.created_at, '%s') as date, - u.api_key_id, - COALESCE(k.name, '') as key_name, - COUNT(*) as requests, - COALESCE(SUM(u.input_tokens + u.output_tokens + u.cache_creation_tokens + u.cache_read_tokens), 0) as tokens - FROM usage_logs u - LEFT JOIN api_keys k ON u.api_key_id = k.id - WHERE u.api_key_id IN (SELECT api_key_id FROM top_keys) - AND u.created_at >= $4 AND u.created_at < $5 - GROUP BY date, u.api_key_id, k.name - ORDER BY date ASC, tokens DESC - `, dateFormat) - - rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit, startTime, endTime) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results = make([]APIKeyUsageTrendPoint, 0) - for rows.Next() { - var row APIKeyUsageTrendPoint - if err = rows.Scan(&row.Date, &row.APIKeyID, &row.KeyName, &row.Requests, &row.Tokens); err != nil { - return nil, err - } - results = append(results, row) - } - if err = rows.Err(); err != nil { - return nil, err - } - - return results, nil -} - -// GetUserUsageTrend returns usage trend data grouped by user and date -func (r *usageLogRepository) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) (results []UserUsageTrendPoint, err error) { - dateFormat := safeDateFormat(granularity) - - query := fmt.Sprintf(` - WITH top_users AS ( - SELECT user_id - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - GROUP BY user_id - ORDER BY SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) DESC - LIMIT $3 - ) - SELECT - TO_CHAR(u.created_at, '%s') as date, - u.user_id, - COALESCE(us.email, '') as email, - COALESCE(us.username, '') as username, - COUNT(*) as requests, - COALESCE(SUM(u.input_tokens + u.output_tokens + u.cache_creation_tokens + u.cache_read_tokens), 0) as tokens, - COALESCE(SUM(u.total_cost), 0) as cost, - COALESCE(SUM(u.actual_cost), 0) as actual_cost - FROM usage_logs u - LEFT JOIN users us ON u.user_id = us.id - WHERE u.user_id IN (SELECT user_id FROM top_users) - AND u.created_at >= $4 AND u.created_at < $5 - GROUP BY date, u.user_id, us.email, us.username - ORDER BY date ASC, tokens DESC - `, dateFormat) - - rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit, startTime, endTime) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results = make([]UserUsageTrendPoint, 0) - for rows.Next() { - var row UserUsageTrendPoint - if err = rows.Scan(&row.Date, &row.UserID, &row.Email, &row.Username, &row.Requests, &row.Tokens, &row.Cost, &row.ActualCost); err != nil { - return nil, err - } - results = append(results, row) - } - if err = rows.Err(); err != nil { - return nil, err - } - - return results, nil -} - -// GetUserSpendingRanking returns user spending ranking aggregated within the time range. -func (r *usageLogRepository) GetUserSpendingRanking(ctx context.Context, startTime, endTime time.Time, limit int) (result *UserSpendingRankingResponse, err error) { - if limit <= 0 { - limit = 12 - } - - query := ` - WITH user_spend AS ( - SELECT - u.user_id, - COALESCE(us.email, '') as email, - COALESCE(SUM(u.actual_cost), 0) as actual_cost, - COUNT(*) as requests, - COALESCE(SUM(u.input_tokens + u.output_tokens + u.cache_creation_tokens + u.cache_read_tokens), 0) as tokens - FROM usage_logs u - LEFT JOIN users us ON u.user_id = us.id - WHERE u.created_at >= $1 AND u.created_at < $2 - GROUP BY u.user_id, us.email - ), - ranked AS ( - SELECT - user_id, - email, - actual_cost, - requests, - tokens, - COALESCE(SUM(actual_cost) OVER (), 0) as total_actual_cost, - COALESCE(SUM(requests) OVER (), 0) as total_requests, - COALESCE(SUM(tokens) OVER (), 0) as total_tokens - FROM user_spend - ORDER BY actual_cost DESC, tokens DESC, user_id ASC - LIMIT $3 - ) - SELECT - user_id, - email, - actual_cost, - requests, - tokens, - total_actual_cost, - total_requests, - total_tokens - FROM ranked - ORDER BY actual_cost DESC, tokens DESC, user_id ASC - ` - - rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - result = nil - } - }() - - ranking := make([]UserSpendingRankingItem, 0) - totalActualCost := 0.0 - totalRequests := int64(0) - totalTokens := int64(0) - for rows.Next() { - var row UserSpendingRankingItem - if err = rows.Scan(&row.UserID, &row.Email, &row.ActualCost, &row.Requests, &row.Tokens, &totalActualCost, &totalRequests, &totalTokens); err != nil { - return nil, err - } - ranking = append(ranking, row) - } - if err = rows.Err(); err != nil { - return nil, err - } - - return &UserSpendingRankingResponse{ - Ranking: ranking, - TotalActualCost: totalActualCost, - TotalRequests: totalRequests, - TotalTokens: totalTokens, - }, nil -} - -// UserDashboardStats 用户仪表盘统计 -type UserDashboardStats = usagestats.UserDashboardStats - -// PlatformDashboardStats 单平台用量明细 -type PlatformDashboardStats = usagestats.PlatformDashboardStats - -// GetUserDashboardStats 获取用户专属的仪表盘统计 -func (r *usageLogRepository) GetUserDashboardStats(ctx context.Context, userID int64) (*UserDashboardStats, error) { - stats := &UserDashboardStats{} - today := timezone.Today() - - // API Key 统计 - if err := scanSingleRow( - ctx, - r.sql, - "SELECT COUNT(*) FROM api_keys WHERE user_id = $1 AND deleted_at IS NULL", - []any{userID}, - &stats.TotalAPIKeys, - ); err != nil { - return nil, err - } - if err := scanSingleRow( - ctx, - r.sql, - "SELECT COUNT(*) FROM api_keys WHERE user_id = $1 AND status = $2 AND deleted_at IS NULL", - []any{userID, service.StatusActive}, - &stats.ActiveAPIKeys, - ); err != nil { - return nil, err - } - - // 累计 Token 统计 - totalStatsQuery := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(duration_ms), 0) as avg_duration_ms - FROM usage_logs - WHERE user_id = $1 - ` - if err := scanSingleRow( - ctx, - r.sql, - totalStatsQuery, - []any{userID}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens - - // 今日 Token 统计 - todayStatsQuery := ` - SELECT - COUNT(*) as today_requests, - COALESCE(SUM(input_tokens), 0) as today_input_tokens, - COALESCE(SUM(output_tokens), 0) as today_output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as today_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as today_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as today_cost, - COALESCE(SUM(actual_cost), 0) as today_actual_cost - FROM usage_logs - WHERE user_id = $1 AND created_at >= $2 - ` - if err := scanSingleRow( - ctx, - r.sql, - todayStatsQuery, - []any{userID, today}, - &stats.TodayRequests, - &stats.TodayInputTokens, - &stats.TodayOutputTokens, - &stats.TodayCacheCreationTokens, - &stats.TodayCacheReadTokens, - &stats.TodayCost, - &stats.TodayActualCost, - ); err != nil { - return nil, err - } - stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens - - // 性能指标:RPM 和 TPM(最近1分钟,仅统计该用户的请求) - rpm, tpm, err := r.getPerformanceStats(ctx, userID) - if err != nil { - return nil, err - } - stats.Rpm = rpm - stats.Tpm = tpm - - // 按"有效平台"维度拆分(group.platform 优先,否则 account.platform)。 - // 与 ops 路径口径一致;HAVING 过滤掉无法确定平台的行(避免出现空字符串平台)。 - // 与上面 totalStatsQuery/todayStatsQuery 的总值可能略微差异,原因有二: - // 1) 无平台归属的极少数行(group/account 都没 platform)会被 HAVING 排除; - // 2) usageLogSuccessFilterUL 会把 actual_cost = 0 的失败 placeholder 行排除, - // 而 totalStatsQuery/todayStatsQuery 没有这层过滤、会把这些行的 request 计数算进去。 - platformQuery := ` - SELECT - ` + usageLogEffectivePlatformExpr + ` as platform, - COUNT(*) as total_requests, - COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(ul.actual_cost), 0) as total_actual_cost, - COUNT(*) FILTER (WHERE ul.created_at >= $2) as today_requests, - COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens) FILTER (WHERE ul.created_at >= $2), 0) as today_tokens, - COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $2), 0) as today_actual_cost - FROM usage_logs ul - LEFT JOIN groups g ON g.id = ul.group_id - LEFT JOIN accounts a ON a.id = ul.account_id - WHERE ul.user_id = $1 - AND ` + usageLogSuccessFilterUL + ` - GROUP BY ` + usageLogEffectivePlatformExpr + ` - HAVING ` + usageLogEffectivePlatformExpr + ` IS NOT NULL AND ` + usageLogEffectivePlatformExpr + ` <> '' - ORDER BY total_actual_cost DESC - ` - rows, err := r.sql.QueryContext(ctx, platformQuery, userID, today) - if err != nil { - return nil, err - } - for rows.Next() { - var p PlatformDashboardStats - if err := rows.Scan( - &p.Platform, - &p.TotalRequests, - &p.TotalTokens, - &p.TotalActualCost, - &p.TodayRequests, - &p.TodayTokens, - &p.TodayActualCost, - ); err != nil { - _ = rows.Close() - return nil, err - } - stats.ByPlatform = append(stats.ByPlatform, p) - } - if err := rows.Close(); err != nil { - return nil, err - } - if err := rows.Err(); err != nil { - return nil, err - } - - return stats, nil -} - -// getPerformanceStatsByAPIKey 获取指定 API Key 的 RPM 和 TPM(近5分钟平均值) -func (r *usageLogRepository) getPerformanceStatsByAPIKey(ctx context.Context, apiKeyID int64) (rpm, tpm int64, err error) { - fiveMinutesAgo := time.Now().Add(-5 * time.Minute) - query := ` - SELECT - COUNT(*) as request_count, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as token_count - FROM usage_logs - WHERE created_at >= $1 AND api_key_id = $2` - args := []any{fiveMinutesAgo, apiKeyID} - - var requestCount int64 - var tokenCount int64 - if err := scanSingleRow(ctx, r.sql, query, args, &requestCount, &tokenCount); err != nil { - return 0, 0, err - } - return requestCount / 5, tokenCount / 5, nil -} - -// GetAPIKeyDashboardStats 获取指定 API Key 的仪表盘统计(按 api_key_id 过滤) -func (r *usageLogRepository) GetAPIKeyDashboardStats(ctx context.Context, apiKeyID int64) (*UserDashboardStats, error) { - stats := &UserDashboardStats{} - today := timezone.Today() - - // API Key 维度不需要统计 key 数量,设为 1 - stats.TotalAPIKeys = 1 - stats.ActiveAPIKeys = 1 - - // 累计 Token 统计 - totalStatsQuery := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(duration_ms), 0) as avg_duration_ms - FROM usage_logs - WHERE api_key_id = $1 - ` - if err := scanSingleRow( - ctx, - r.sql, - totalStatsQuery, - []any{apiKeyID}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens - - // 今日 Token 统计 - todayStatsQuery := ` - SELECT - COUNT(*) as today_requests, - COALESCE(SUM(input_tokens), 0) as today_input_tokens, - COALESCE(SUM(output_tokens), 0) as today_output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as today_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as today_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as today_cost, - COALESCE(SUM(actual_cost), 0) as today_actual_cost - FROM usage_logs - WHERE api_key_id = $1 AND created_at >= $2 - ` - if err := scanSingleRow( - ctx, - r.sql, - todayStatsQuery, - []any{apiKeyID, today}, - &stats.TodayRequests, - &stats.TodayInputTokens, - &stats.TodayOutputTokens, - &stats.TodayCacheCreationTokens, - &stats.TodayCacheReadTokens, - &stats.TodayCost, - &stats.TodayActualCost, - ); err != nil { - return nil, err - } - stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens - - // 性能指标:RPM 和 TPM(最近5分钟,按 API Key 过滤) - rpm, tpm, err := r.getPerformanceStatsByAPIKey(ctx, apiKeyID) - if err != nil { - return nil, err - } - stats.Rpm = rpm - stats.Tpm = tpm - - return stats, nil -} - -// GetUserUsageTrendByUserID 获取指定用户的使用趋势 -func (r *usageLogRepository) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (results []TrendDataPoint, err error) { - dateFormat := safeDateFormat(granularity) - - query := fmt.Sprintf(` - SELECT - TO_CHAR(created_at, '%s') as date, - COUNT(*) as requests, - COALESCE(SUM(input_tokens), 0) as input_tokens, - COALESCE(SUM(output_tokens), 0) as output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(total_cost), 0) as cost, - COALESCE(SUM(actual_cost), 0) as actual_cost - FROM usage_logs - WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 - GROUP BY date - ORDER BY date ASC - `, dateFormat) - - rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results, err = scanTrendRows(rows) - if err != nil { - return nil, err - } - return results, nil -} - -// GetUserModelStats 获取指定用户的模型统计 -func (r *usageLogRepository) GetUserModelStats(ctx context.Context, userID int64, startTime, endTime time.Time) (results []ModelStat, err error) { - return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, 0, 0, 0, "", nil, nil, nil, usagestats.ModelSourceRequested, "") -} - -// UsageLogFilters represents filters for usage log queries -type UsageLogFilters = usagestats.UsageLogFilters - -// ListWithFilters lists usage logs with optional filters (for admin) -func (r *usageLogRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters UsageLogFilters) ([]service.UsageLog, *pagination.PaginationResult, error) { - conditions := make([]string, 0, 9) - args := make([]any, 0, 9) - - if filters.UserID > 0 { - conditions = append(conditions, fmt.Sprintf("user_id = $%d", len(args)+1)) - args = append(args, filters.UserID) - } - if filters.APIKeyID > 0 { - conditions = append(conditions, fmt.Sprintf("api_key_id = $%d", len(args)+1)) - args = append(args, filters.APIKeyID) - } - if filters.AccountID > 0 { - conditions = append(conditions, fmt.Sprintf("account_id = $%d", len(args)+1)) - args = append(args, filters.AccountID) - } - if filters.GroupID > 0 { - conditions = append(conditions, fmt.Sprintf("group_id = $%d", len(args)+1)) - args = append(args, filters.GroupID) - } - conditions, args = appendUsageLogModelWhereCondition(conditions, args, filters.Model, filters.ModelFilterSource) - conditions, args = appendRequestTypeOrStreamWhereCondition(conditions, args, filters.RequestType, filters.Stream) - if filters.BillingType != nil { - conditions = append(conditions, fmt.Sprintf("billing_type = $%d", len(args)+1)) - args = append(args, int16(*filters.BillingType)) - } - conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode) - if filters.StartTime != nil { - conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1)) - args = append(args, *filters.StartTime) - } - if filters.EndTime != nil { - conditions = append(conditions, fmt.Sprintf("created_at < $%d", len(args)+1)) - args = append(args, *filters.EndTime) - } - - whereClause := buildWhere(conditions) - var ( - logs []service.UsageLog - page *pagination.PaginationResult - err error - ) - if shouldUseFastUsageLogTotal(filters) { - logs, page, err = r.listUsageLogsWithFastPagination(ctx, whereClause, args, params) - } else { - logs, page, err = r.listUsageLogsWithPagination(ctx, whereClause, args, params) - } - if err != nil { - return nil, nil, err - } - - if err := r.hydrateUsageLogAssociations(ctx, logs); err != nil { - return nil, nil, err - } - return logs, page, nil -} - -func shouldUseFastUsageLogTotal(filters UsageLogFilters) bool { - if filters.ExactTotal { - return false - } - // 强选择过滤下记录集通常较小,保留精确总数。 - return filters.UserID == 0 && filters.APIKeyID == 0 && filters.AccountID == 0 -} - -// UsageStats represents usage statistics -type UsageStats = usagestats.UsageStats - -// BatchUserUsageStats represents usage stats for a single user -type BatchUserUsageStats = usagestats.BatchUserUsageStats - -// PlatformUsage represents per-platform usage breakdown -type PlatformUsage = usagestats.PlatformUsage - -func normalizePositiveInt64IDs(ids []int64) []int64 { - if len(ids) == 0 { - return nil - } - seen := make(map[int64]struct{}, len(ids)) - out := make([]int64, 0, len(ids)) - for _, id := range ids { - if id <= 0 { - continue - } - if _, ok := seen[id]; ok { - continue - } - seen[id] = struct{}{} - out = append(out, id) - } - return out -} - -// GetBatchUserUsageStats gets today and total actual_cost for multiple users within a time range. -// If startTime is zero, defaults to 30 days ago. -func (r *usageLogRepository) GetBatchUserUsageStats(ctx context.Context, userIDs []int64, startTime, endTime time.Time) (map[int64]*BatchUserUsageStats, error) { - result := make(map[int64]*BatchUserUsageStats) - normalizedUserIDs := normalizePositiveInt64IDs(userIDs) - if len(normalizedUserIDs) == 0 { - return result, nil - } - - // 默认最近 30 天 - if startTime.IsZero() { - startTime = time.Now().AddDate(0, 0, -30) - } - if endTime.IsZero() { - endTime = time.Now() - } - - for _, id := range normalizedUserIDs { - result[id] = &BatchUserUsageStats{UserID: id} - } - - // GROUP BY (user_id, effective_platform) 一次查询同时得到总值与按平台拆分。 - // 应用层把同一 user_id 的多行累加为总值,并把非空 platform 行收集到 ByPlatform。 - query := ` - SELECT - ul.user_id, - ` + usageLogEffectivePlatformExpr + ` as platform, - COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $2 AND ul.created_at < $3), 0) as total_cost, - COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $4), 0) as today_cost - FROM usage_logs ul - LEFT JOIN groups g ON g.id = ul.group_id - LEFT JOIN accounts a ON a.id = ul.account_id - WHERE ul.user_id = ANY($1) - AND ul.created_at >= LEAST($2, $4) - AND ` + usageLogSuccessFilterUL + ` - GROUP BY ul.user_id, ` + usageLogEffectivePlatformExpr + ` - ` - today := timezone.Today() - rows, err := r.sql.QueryContext(ctx, query, pq.Array(normalizedUserIDs), startTime, endTime, today) - if err != nil { - return nil, err - } - for rows.Next() { - var userID int64 - var platform sql.NullString - var total float64 - var todayTotal float64 - if err := rows.Scan(&userID, &platform, &total, &todayTotal); err != nil { - _ = rows.Close() - return nil, err - } - stats, ok := result[userID] - if !ok { - continue - } - stats.TotalActualCost += total - stats.TodayActualCost += todayTotal - if platform.Valid && platform.String != "" { - stats.ByPlatform = append(stats.ByPlatform, PlatformUsage{ - Platform: platform.String, - TotalActualCost: total, - TodayActualCost: todayTotal, - }) - } - } - if err := rows.Close(); err != nil { - return nil, err - } - if err := rows.Err(); err != nil { - return nil, err - } - - return result, nil -} - -// BatchAPIKeyUsageStats represents usage stats for a single API key -type BatchAPIKeyUsageStats = usagestats.BatchAPIKeyUsageStats - -// GetBatchAPIKeyUsageStats gets today and total actual_cost for multiple API keys within a time range. -// If startTime is zero, defaults to 30 days ago. -func (r *usageLogRepository) GetBatchAPIKeyUsageStats(ctx context.Context, apiKeyIDs []int64, startTime, endTime time.Time) (map[int64]*BatchAPIKeyUsageStats, error) { - result := make(map[int64]*BatchAPIKeyUsageStats) - normalizedAPIKeyIDs := normalizePositiveInt64IDs(apiKeyIDs) - if len(normalizedAPIKeyIDs) == 0 { - return result, nil - } - - // 默认最近 30 天 - if startTime.IsZero() { - startTime = time.Now().AddDate(0, 0, -30) - } - if endTime.IsZero() { - endTime = time.Now() - } - - for _, id := range normalizedAPIKeyIDs { - result[id] = &BatchAPIKeyUsageStats{APIKeyID: id} - } - - query := ` - SELECT - api_key_id, - COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $2 AND created_at < $3), 0) as total_cost, - COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $4), 0) as today_cost - FROM usage_logs - WHERE api_key_id = ANY($1) - AND created_at >= LEAST($2, $4) - GROUP BY api_key_id - ` - today := timezone.Today() - rows, err := r.sql.QueryContext(ctx, query, pq.Array(normalizedAPIKeyIDs), startTime, endTime, today) - if err != nil { - return nil, err - } - for rows.Next() { - var apiKeyID int64 - var total float64 - var todayTotal float64 - if err := rows.Scan(&apiKeyID, &total, &todayTotal); err != nil { - _ = rows.Close() - return nil, err - } - if stats, ok := result[apiKeyID]; ok { - stats.TotalActualCost = total - stats.TodayActualCost = todayTotal - } - } - if err := rows.Close(); err != nil { - return nil, err - } - if err := rows.Err(); err != nil { - return nil, err - } - - return result, nil -} - -// GetUsageTrendWithFilters returns usage trend data with optional filters -func (r *usageLogRepository) GetUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) (results []TrendDataPoint, err error) { - return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "") -} - -func (r *usageLogRepository) GetUsageTrendWithUsageFilters(ctx context.Context, startTime, endTime time.Time, granularity string, filters UsageLogFilters) (results []TrendDataPoint, err error) { - return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) -} - -func (r *usageLogRepository) getUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []TrendDataPoint, err error) { - if shouldUsePreaggregatedTrend(granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType, billingMode) { - aggregated, aggregatedErr := r.getUsageTrendFromAggregates(ctx, startTime, endTime, granularity) - if aggregatedErr == nil && len(aggregated) > 0 { - return aggregated, nil - } - } - - dateFormat := safeDateFormat(granularity) - - query := fmt.Sprintf(` - SELECT - TO_CHAR(created_at, '%s') as date, - COUNT(*) as requests, - COALESCE(SUM(input_tokens), 0) as input_tokens, - COALESCE(SUM(output_tokens), 0) as output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(total_cost), 0) as cost, - COALESCE(SUM(actual_cost), 0) as actual_cost - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - `, dateFormat) - - args := []any{startTime, endTime} - if userID > 0 { - query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) - args = append(args, userID) - } - if apiKeyID > 0 { - query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) - args = append(args, apiKeyID) - } - if accountID > 0 { - query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) - args = append(args, accountID) - } - if groupID > 0 { - query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) - args = append(args, groupID) - } - query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource) - query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) - if billingType != nil { - query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) - args = append(args, int16(*billingType)) - } - query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") - query += " GROUP BY date ORDER BY date ASC" - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results, err = scanTrendRows(rows) - if err != nil { - return nil, err - } - return results, nil -} - -func shouldUsePreaggregatedTrend(granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string) bool { - if granularity != "day" && granularity != "hour" { - return false - } - return userID == 0 && - apiKeyID == 0 && - accountID == 0 && - groupID == 0 && - model == "" && - requestType == nil && - stream == nil && - billingType == nil && - billingMode == "" -} - -func (r *usageLogRepository) getUsageTrendFromAggregates(ctx context.Context, startTime, endTime time.Time, granularity string) (results []TrendDataPoint, err error) { - dateFormat := safeDateFormat(granularity) - query := "" - args := []any{startTime, endTime} - - switch granularity { - case "hour": - query = fmt.Sprintf(` - SELECT - TO_CHAR(bucket_start, '%s') as date, - total_requests as requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) as total_tokens, - total_cost as cost, - actual_cost - FROM usage_dashboard_hourly - WHERE bucket_start >= $1 AND bucket_start < $2 - ORDER BY bucket_start ASC - `, dateFormat) - case "day": - query = fmt.Sprintf(` - SELECT - TO_CHAR(bucket_date::timestamp, '%s') as date, - total_requests as requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) as total_tokens, - total_cost as cost, - actual_cost - FROM usage_dashboard_daily - WHERE bucket_date >= $1::date AND bucket_date < $2::date - ORDER BY bucket_date ASC - `, dateFormat) - default: - return nil, nil - } - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results, err = scanTrendRows(rows) - if err != nil { - return nil, err - } - return results, nil -} - -// GetModelStatsWithFilters returns model statistics with optional filters -func (r *usageLogRepository) GetModelStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []ModelStat, err error) { - return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, usagestats.ModelSourceRequested, "") -} - -// GetModelStatsWithFiltersBySource returns model statistics with optional filters and model source dimension. -// source: requested | upstream | mapping. -func (r *usageLogRepository) GetModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8, source string) (results []ModelStat, err error) { - return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, source, "") -} - -func (r *usageLogRepository) GetModelStatsWithUsageFiltersBySource(ctx context.Context, startTime, endTime time.Time, filters UsageLogFilters, source string) (results []ModelStat, err error) { - return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, source, filters.BillingMode) -} - -func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, source string, billingMode string) (results []ModelStat, err error) { - actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost" - // 当仅按 account_id 聚合时,实际费用使用账号倍率(total_cost * account_rate_multiplier)。 - if accountID > 0 && userID == 0 && apiKeyID == 0 { - actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost" - } - accountCostExpr := "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as account_cost" - modelExpr := resolveModelDimensionExpression(source) - - query := fmt.Sprintf(` - SELECT - %s as model, - COUNT(*) as requests, - COALESCE(SUM(input_tokens), 0) as input_tokens, - COALESCE(SUM(output_tokens), 0) as output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(total_cost), 0) as cost, - %s, - %s - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - `, modelExpr, actualCostExpr, accountCostExpr) - - args := []any{startTime, endTime} - if userID > 0 { - query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) - args = append(args, userID) - } - if apiKeyID > 0 { - query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) - args = append(args, apiKeyID) - } - if accountID > 0 { - query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) - args = append(args, accountID) - } - if groupID > 0 { - query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) - args = append(args, groupID) - } - if strings.TrimSpace(model) != "" { - query += fmt.Sprintf(" AND %s = $%d", modelExpr, len(args)+1) - args = append(args, model) - } - query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) - if billingType != nil { - query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) - args = append(args, int16(*billingType)) - } - query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") - query += fmt.Sprintf(" GROUP BY %s ORDER BY total_tokens DESC", modelExpr) - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results, err = scanModelStatsRows(rows) - if err != nil { - return nil, err - } - return results, nil -} - -// GetGroupStatsWithFilters returns group usage statistics with optional filters -func (r *usageLogRepository) GetGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []usagestats.GroupStat, err error) { - return r.getGroupStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, "") -} - -func (r *usageLogRepository) GetGroupStatsWithUsageFilters(ctx context.Context, startTime, endTime time.Time, filters UsageLogFilters) (results []usagestats.GroupStat, err error) { - return r.getGroupStatsWithFilters(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) -} - -func (r *usageLogRepository) getGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []usagestats.GroupStat, err error) { - query := ` - SELECT - COALESCE(ul.group_id, 0) as group_id, - COALESCE(g.name, '') as group_name, - COUNT(*) as requests, - COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(ul.total_cost), 0) as cost, - COALESCE(SUM(ul.actual_cost), 0) as actual_cost, - COALESCE(SUM(COALESCE(ul.account_stats_cost, ul.total_cost) * COALESCE(ul.account_rate_multiplier, 1)), 0) as account_cost - FROM usage_logs ul - LEFT JOIN groups g ON g.id = ul.group_id - WHERE ul.created_at >= $1 AND ul.created_at < $2 - ` - - args := []any{startTime, endTime} - if userID > 0 { - query += fmt.Sprintf(" AND ul.user_id = $%d", len(args)+1) - args = append(args, userID) - } - if apiKeyID > 0 { - query += fmt.Sprintf(" AND ul.api_key_id = $%d", len(args)+1) - args = append(args, apiKeyID) - } - if accountID > 0 { - query += fmt.Sprintf(" AND ul.account_id = $%d", len(args)+1) - args = append(args, accountID) - } - if groupID > 0 { - query += fmt.Sprintf(" AND ul.group_id = $%d", len(args)+1) - args = append(args, groupID) - } - if strings.TrimSpace(model) != "" { - modelExpr := resolveModelDimensionExpressionWithAlias(usagestats.ModelSourceRequested, "ul") - query += fmt.Sprintf(" AND %s = $%d", modelExpr, len(args)+1) - args = append(args, model) - } - query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) - if billingType != nil { - query += fmt.Sprintf(" AND ul.billing_type = $%d", len(args)+1) - args = append(args, int16(*billingType)) - } - query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "ul") - query += " GROUP BY ul.group_id, g.name ORDER BY total_tokens DESC" - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results = make([]usagestats.GroupStat, 0) - for rows.Next() { - var row usagestats.GroupStat - if err := rows.Scan( - &row.GroupID, - &row.GroupName, - &row.Requests, - &row.TotalTokens, - &row.Cost, - &row.ActualCost, - &row.AccountCost, - ); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - -// GetUserBreakdownStats returns per-user usage breakdown within a specific dimension. -func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) (results []usagestats.UserBreakdownItem, err error) { - query := ` - SELECT - COALESCE(ul.user_id, 0) as user_id, - COALESCE(u.email, '') as email, - COUNT(*) as requests, - COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(ul.total_cost), 0) as cost, - COALESCE(SUM(ul.actual_cost), 0) as actual_cost, - COALESCE(SUM(COALESCE(ul.account_stats_cost, ul.total_cost) * COALESCE(ul.account_rate_multiplier, 1)), 0) as account_cost - FROM usage_logs ul - LEFT JOIN users u ON u.id = ul.user_id - WHERE ul.created_at >= $1 AND ul.created_at < $2 - ` - args := []any{startTime, endTime} - - if dim.GroupID > 0 { - query += fmt.Sprintf(" AND ul.group_id = $%d", len(args)+1) - args = append(args, dim.GroupID) - } - if dim.Model != "" { - query += fmt.Sprintf(" AND %s = $%d", resolveModelDimensionExpression(dim.ModelType), len(args)+1) - args = append(args, dim.Model) - } - if dim.Endpoint != "" { - col := resolveEndpointColumn(dim.EndpointType) - query += fmt.Sprintf(" AND %s = $%d", col, len(args)+1) - args = append(args, dim.Endpoint) - } - if dim.UserID > 0 { - query += fmt.Sprintf(" AND ul.user_id = $%d", len(args)+1) - args = append(args, dim.UserID) - } - if dim.APIKeyID > 0 { - query += fmt.Sprintf(" AND ul.api_key_id = $%d", len(args)+1) - args = append(args, dim.APIKeyID) - } - if dim.AccountID > 0 { - query += fmt.Sprintf(" AND ul.account_id = $%d", len(args)+1) - args = append(args, dim.AccountID) - } - if dim.RequestType != nil { - query += fmt.Sprintf(" AND ul.request_type = $%d", len(args)+1) - args = append(args, *dim.RequestType) - } - if dim.Stream != nil { - query += fmt.Sprintf(" AND ul.stream = $%d", len(args)+1) - args = append(args, *dim.Stream) - } - if dim.BillingType != nil { - query += fmt.Sprintf(" AND ul.billing_type = $%d", len(args)+1) - args = append(args, *dim.BillingType) - } - - query += " GROUP BY ul.user_id, u.email ORDER BY actual_cost DESC" - if limit > 0 { - query += fmt.Sprintf(" LIMIT %d", limit) - } - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results = make([]usagestats.UserBreakdownItem, 0) - for rows.Next() { - var row usagestats.UserBreakdownItem - if err := rows.Scan( - &row.UserID, - &row.Email, - &row.Requests, - &row.TotalTokens, - &row.Cost, - &row.ActualCost, - &row.AccountCost, - ); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - -// GetAllGroupUsageSummary returns today's and cumulative actual_cost for every group. -// todayStart is the start-of-day in the caller's timezone (UTC-based). -// TODO(perf): This query scans ALL usage_logs rows for total_cost aggregation. -// When usage_logs exceeds ~1M rows, consider adding a short-lived cache (30s) -// or a materialized view / pre-aggregation table for cumulative costs. -func (r *usageLogRepository) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) { - query := ` - SELECT - g.id AS group_id, - COALESCE(SUM(ul.actual_cost), 0) AS total_cost, - COALESCE(SUM(CASE WHEN ul.created_at >= $1 THEN ul.actual_cost ELSE 0 END), 0) AS today_cost - FROM groups g - LEFT JOIN usage_logs ul ON ul.group_id = g.id - GROUP BY g.id - ` - - rows, err := r.sql.QueryContext(ctx, query, todayStart) - if err != nil { - return nil, err - } - defer func() { _ = rows.Close() }() - var results []usagestats.GroupUsageSummary - for rows.Next() { - var row usagestats.GroupUsageSummary - if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - -// resolveModelDimensionExpression maps model source type to a safe SQL expression. -func resolveModelDimensionExpression(modelType string) string { - return resolveModelDimensionExpressionWithAlias(modelType, "") -} - -func resolveModelDimensionExpressionWithAlias(modelType, alias string) string { - column := func(name string) string { - if alias == "" { - return name - } - return alias + "." + name - } - requestedExpr := fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("requested_model"), column("model")) - switch usagestats.NormalizeModelSource(modelType) { - case usagestats.ModelSourceUpstream: - return fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("upstream_model"), requestedExpr) - case usagestats.ModelSourceMapping: - return fmt.Sprintf("(%s || ' -> ' || COALESCE(NULLIF(TRIM(%s), ''), %s))", requestedExpr, column("upstream_model"), requestedExpr) - default: - return requestedExpr - } -} - -// resolveEndpointColumn maps endpoint type to the corresponding DB column name. -func resolveEndpointColumn(endpointType string) string { - switch endpointType { - case "upstream": - return "ul.upstream_endpoint" - case "path": - return "ul.inbound_endpoint || ' -> ' || ul.upstream_endpoint" - default: - return "ul.inbound_endpoint" - } -} - -// GetGlobalStats gets usage statistics for all users within a time range -func (r *usageLogRepository) GetGlobalStats(ctx context.Context, startTime, endTime time.Time) (*UsageStats, error) { - query := ` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(AVG(duration_ms), 0) as avg_duration_ms - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - ` - - stats := &UsageStats{} - if err := scanSingleRow( - ctx, - r.sql, - query, - []any{startTime, endTime}, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens - return stats, nil -} - -// GetStatsWithFilters gets usage statistics with optional filters -func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters UsageLogFilters) (*UsageStats, error) { - conditions := make([]string, 0, 9) - args := make([]any, 0, 9) - - if filters.UserID > 0 { - conditions = append(conditions, fmt.Sprintf("user_id = $%d", len(args)+1)) - args = append(args, filters.UserID) - } - if filters.APIKeyID > 0 { - conditions = append(conditions, fmt.Sprintf("api_key_id = $%d", len(args)+1)) - args = append(args, filters.APIKeyID) - } - if filters.AccountID > 0 { - conditions = append(conditions, fmt.Sprintf("account_id = $%d", len(args)+1)) - args = append(args, filters.AccountID) - } - if filters.GroupID > 0 { - conditions = append(conditions, fmt.Sprintf("group_id = $%d", len(args)+1)) - args = append(args, filters.GroupID) - } - conditions, args = appendUsageLogModelWhereCondition(conditions, args, filters.Model, filters.ModelFilterSource) - conditions, args = appendRequestTypeOrStreamWhereCondition(conditions, args, filters.RequestType, filters.Stream) - if filters.BillingType != nil { - conditions = append(conditions, fmt.Sprintf("billing_type = $%d", len(args)+1)) - args = append(args, int16(*filters.BillingType)) - } - conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode) - if filters.StartTime != nil { - conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1)) - args = append(args, *filters.StartTime) - } - if filters.EndTime != nil { - conditions = append(conditions, fmt.Sprintf("created_at < $%d", len(args)+1)) - args = append(args, *filters.EndTime) - } - - query := fmt.Sprintf(` - SELECT - COUNT(*) as total_requests, - COALESCE(SUM(input_tokens), 0) as total_input_tokens, - COALESCE(SUM(output_tokens), 0) as total_output_tokens, - COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, - COALESCE(SUM(total_cost), 0) as total_cost, - COALESCE(SUM(actual_cost), 0) as total_actual_cost, - COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as total_account_cost, - COALESCE(AVG(duration_ms), 0) as avg_duration_ms - FROM usage_logs - %s - `, buildWhere(conditions)) - - stats := &UsageStats{} - var totalAccountCost float64 - - start := time.Unix(0, 0).UTC() - if filters.StartTime != nil { - start = *filters.StartTime - } - end := time.Now().UTC() - if filters.EndTime != nil { - end = *filters.EndTime - } - - var endpoints, upstreamEndpoints, endpointPaths []EndpointStat - - // 汇总查询:失败即致命。 - runSummary := func(c context.Context) error { - return scanSingleRow( - c, r.sql, query, args, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalCacheCreationTokens, - &stats.TotalCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &totalAccountCost, - &stats.AverageDurationMs, - ) - } - // endpoint 明细:best-effort(失败 log + 返空),不致命。 - runEndpoints := func(c context.Context) { - res, err := r.getEndpointStatsByColumnWithFilters(c, "inbound_endpoint", start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) - if err != nil { - if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { - logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) - } - res = []EndpointStat{} - } - endpoints = res - } - runUpstream := func(c context.Context) { - res, err := r.getEndpointStatsByColumnWithFilters(c, "upstream_endpoint", start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) - if err != nil { - if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { - logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) - } - res = []EndpointStat{} - } - upstreamEndpoints = res - } - runPaths := func(c context.Context) { - res, err := r.getEndpointPathStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) - if err != nil { - if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { - logger.LegacyPrintf("repository.usage_log", "getEndpointPathStatsWithFilters failed in GetStatsWithFilters: %v", err) - } - res = []EndpointStat{} - } - endpointPaths = res - } - - if r.db != nil { - // 生产路径:r.sql 是 *sql.DB 连接池,可并发。4 条查询并行,延迟取最大值。 - g, gctx := errgroup.WithContext(ctx) - g.Go(func() error { return runSummary(gctx) }) - g.Go(func() error { runEndpoints(gctx); return nil }) - g.Go(func() error { runUpstream(gctx); return nil }) - g.Go(func() error { runPaths(gctx); return nil }) - if err := g.Wait(); err != nil { - return nil, err - } - } else { - // 事务路径(ent.Tx 不能并发查询):顺序执行,行为与重构前一致。 - if err := runSummary(ctx); err != nil { - return nil, err - } - runEndpoints(ctx) - runUpstream(ctx) - runPaths(ctx) - } - - stats.TotalAccountCost = &totalAccountCost - stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens - stats.Endpoints = endpoints - stats.UpstreamEndpoints = upstreamEndpoints - stats.EndpointPaths = endpointPaths - - return stats, nil -} - -// AccountUsageHistory represents daily usage history for an account -type AccountUsageHistory = usagestats.AccountUsageHistory - -// AccountUsageSummary represents summary statistics for an account -type AccountUsageSummary = usagestats.AccountUsageSummary - -// AccountUsageStatsResponse represents the full usage statistics response for an account -type AccountUsageStatsResponse = usagestats.AccountUsageStatsResponse - -// EndpointStat represents endpoint usage statistics row. -type EndpointStat = usagestats.EndpointStat - -func (r *usageLogRepository) getEndpointStatsByColumnWithFilters(ctx context.Context, endpointColumn string, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []EndpointStat, err error) { - actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost" - if accountID > 0 && userID == 0 && apiKeyID == 0 { - actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost" - } - - query := fmt.Sprintf(` - SELECT - COALESCE(NULLIF(TRIM(%s), ''), 'unknown') AS endpoint, - COUNT(*) AS requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, - COALESCE(SUM(total_cost), 0) as cost, - %s - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - `, endpointColumn, actualCostExpr) - - args := []any{startTime, endTime} - if userID > 0 { - query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) - args = append(args, userID) - } - if apiKeyID > 0 { - query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) - args = append(args, apiKeyID) - } - if accountID > 0 { - query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) - args = append(args, accountID) - } - if groupID > 0 { - query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) - args = append(args, groupID) - } - query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource) - query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) - if billingType != nil { - query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) - args = append(args, int16(*billingType)) - } - query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") - query += " GROUP BY endpoint ORDER BY requests DESC" - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results = make([]EndpointStat, 0) - for rows.Next() { - var row EndpointStat - if err := rows.Scan(&row.Endpoint, &row.Requests, &row.TotalTokens, &row.Cost, &row.ActualCost); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - -func (r *usageLogRepository) getEndpointPathStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []EndpointStat, err error) { - actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost" - if accountID > 0 && userID == 0 && apiKeyID == 0 { - actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost" - } - - query := fmt.Sprintf(` - SELECT - CONCAT( - COALESCE(NULLIF(TRIM(inbound_endpoint), ''), 'unknown'), - ' -> ', - COALESCE(NULLIF(TRIM(upstream_endpoint), ''), 'unknown') - ) AS endpoint, - COUNT(*) AS requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, - COALESCE(SUM(total_cost), 0) as cost, - %s - FROM usage_logs - WHERE created_at >= $1 AND created_at < $2 - `, actualCostExpr) - - args := []any{startTime, endTime} - if userID > 0 { - query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) - args = append(args, userID) - } - if apiKeyID > 0 { - query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) - args = append(args, apiKeyID) - } - if accountID > 0 { - query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) - args = append(args, accountID) - } - if groupID > 0 { - query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) - args = append(args, groupID) - } - query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource) - query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) - if billingType != nil { - query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) - args = append(args, int16(*billingType)) - } - query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") - query += " GROUP BY endpoint ORDER BY requests DESC" - - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - results = nil - } - }() - - results = make([]EndpointStat, 0) - for rows.Next() { - var row EndpointStat - if err := rows.Scan(&row.Endpoint, &row.Requests, &row.TotalTokens, &row.Cost, &row.ActualCost); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - -// GetEndpointStatsWithFilters returns inbound endpoint statistics with optional filters. -func (r *usageLogRepository) GetEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]EndpointStat, error) { - return r.getEndpointStatsByColumnWithFilters(ctx, "inbound_endpoint", startTime, endTime, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "") -} - -// GetUpstreamEndpointStatsWithFilters returns upstream endpoint statistics with optional filters. -func (r *usageLogRepository) GetUpstreamEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]EndpointStat, error) { - return r.getEndpointStatsByColumnWithFilters(ctx, "upstream_endpoint", startTime, endTime, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "") -} - -// GetAccountUsageStats returns comprehensive usage statistics for an account over a time range -func (r *usageLogRepository) GetAccountUsageStats(ctx context.Context, accountID int64, startTime, endTime time.Time) (resp *AccountUsageStatsResponse, err error) { - daysCount := int(endTime.Sub(startTime).Hours()/24) + 1 - if daysCount <= 0 { - daysCount = 30 - } - - query := ` - SELECT - TO_CHAR(created_at, 'YYYY-MM-DD') as date, - COUNT(*) as requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, - COALESCE(SUM(total_cost), 0) as cost, - COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost, - COALESCE(SUM(actual_cost), 0) as user_cost - FROM usage_logs - WHERE account_id = $1 AND created_at >= $2 AND created_at < $3 - GROUP BY date - ORDER BY date ASC - ` - - rows, err := r.sql.QueryContext(ctx, query, accountID, startTime, endTime) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - resp = nil - } - }() - - history := make([]AccountUsageHistory, 0) - for rows.Next() { - var date string - var requests int64 - var tokens int64 - var cost float64 - var actualCost float64 - var userCost float64 - if err = rows.Scan(&date, &requests, &tokens, &cost, &actualCost, &userCost); err != nil { - return nil, err - } - t, _ := time.Parse("2006-01-02", date) - history = append(history, AccountUsageHistory{ - Date: date, - Label: t.Format("01/02"), - Requests: requests, - Tokens: tokens, - Cost: cost, - ActualCost: actualCost, - UserCost: userCost, - }) - } - if err = rows.Err(); err != nil { - return nil, err - } - - var totalAccountCost, totalUserCost, totalStandardCost float64 - var totalRequests, totalTokens int64 - var highestCostDay, highestRequestDay *AccountUsageHistory - - for i := range history { - h := &history[i] - totalAccountCost += h.ActualCost - totalUserCost += h.UserCost - totalStandardCost += h.Cost - totalRequests += h.Requests - totalTokens += h.Tokens - - if highestCostDay == nil || h.ActualCost > highestCostDay.ActualCost { - highestCostDay = h - } - if highestRequestDay == nil || h.Requests > highestRequestDay.Requests { - highestRequestDay = h - } - } - - actualDaysUsed := len(history) - if actualDaysUsed == 0 { - actualDaysUsed = 1 - } - - avgQuery := "SELECT COALESCE(AVG(duration_ms), 0) as avg_duration_ms FROM usage_logs WHERE account_id = $1 AND created_at >= $2 AND created_at < $3" - var avgDuration float64 - if err := scanSingleRow(ctx, r.sql, avgQuery, []any{accountID, startTime, endTime}, &avgDuration); err != nil { - return nil, err - } - - summary := AccountUsageSummary{ - Days: daysCount, - ActualDaysUsed: actualDaysUsed, - TotalCost: totalAccountCost, - TotalUserCost: totalUserCost, - TotalStandardCost: totalStandardCost, - TotalRequests: totalRequests, - TotalTokens: totalTokens, - AvgDailyCost: totalAccountCost / float64(actualDaysUsed), - AvgDailyUserCost: totalUserCost / float64(actualDaysUsed), - AvgDailyRequests: float64(totalRequests) / float64(actualDaysUsed), - AvgDailyTokens: float64(totalTokens) / float64(actualDaysUsed), - AvgDurationMs: avgDuration, - } - - todayStr := timezone.Now().Format("2006-01-02") - for i := range history { - if history[i].Date == todayStr { - summary.Today = &struct { - Date string `json:"date"` - Cost float64 `json:"cost"` - UserCost float64 `json:"user_cost"` - Requests int64 `json:"requests"` - Tokens int64 `json:"tokens"` - }{ - Date: history[i].Date, - Cost: history[i].ActualCost, - UserCost: history[i].UserCost, - Requests: history[i].Requests, - Tokens: history[i].Tokens, - } - break - } - } - - if highestCostDay != nil { - summary.HighestCostDay = &struct { - Date string `json:"date"` - Label string `json:"label"` - Cost float64 `json:"cost"` - UserCost float64 `json:"user_cost"` - Requests int64 `json:"requests"` - }{ - Date: highestCostDay.Date, - Label: highestCostDay.Label, - Cost: highestCostDay.ActualCost, - UserCost: highestCostDay.UserCost, - Requests: highestCostDay.Requests, - } - } - - if highestRequestDay != nil { - summary.HighestRequestDay = &struct { - Date string `json:"date"` - Label string `json:"label"` - Requests int64 `json:"requests"` - Cost float64 `json:"cost"` - UserCost float64 `json:"user_cost"` - }{ - Date: highestRequestDay.Date, - Label: highestRequestDay.Label, - Requests: highestRequestDay.Requests, - Cost: highestRequestDay.ActualCost, - UserCost: highestRequestDay.UserCost, - } - } - - models, err := r.GetModelStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, nil, nil, nil) - if err != nil { - models = []ModelStat{} - } - endpoints, endpointErr := r.GetEndpointStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, "", nil, nil, nil) - if endpointErr != nil { - logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetAccountUsageStats: %v", endpointErr) - endpoints = []EndpointStat{} - } - upstreamEndpoints, upstreamEndpointErr := r.GetUpstreamEndpointStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, "", nil, nil, nil) - if upstreamEndpointErr != nil { - logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetAccountUsageStats: %v", upstreamEndpointErr) - upstreamEndpoints = []EndpointStat{} - } - - resp = &AccountUsageStatsResponse{ - History: history, - Summary: summary, - Models: models, - Endpoints: endpoints, - UpstreamEndpoints: upstreamEndpoints, - } - return resp, nil -} - -func (r *usageLogRepository) listUsageLogsWithPagination(ctx context.Context, whereClause string, args []any, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { - countQuery := "SELECT COUNT(*) FROM usage_logs " + whereClause - var total int64 - if err := scanSingleRow(ctx, r.sql, countQuery, args, &total); err != nil { - return nil, nil, err - } - - limitPos := len(args) + 1 - offsetPos := len(args) + 2 - listArgs := append(append([]any{}, args...), params.Limit(), params.Offset()) - query := fmt.Sprintf("SELECT %s FROM usage_logs %s ORDER BY %s LIMIT $%d OFFSET $%d", usageLogSelectColumns, whereClause, usageLogOrderBy(params), limitPos, offsetPos) - logs, err := r.queryUsageLogs(ctx, query, listArgs...) - if err != nil { - return nil, nil, err - } - return logs, paginationResultFromTotal(total, params), nil -} - -func (r *usageLogRepository) listUsageLogsWithFastPagination(ctx context.Context, whereClause string, args []any, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { - limit := params.Limit() - offset := params.Offset() - - limitPos := len(args) + 1 - offsetPos := len(args) + 2 - listArgs := append(append([]any{}, args...), limit+1, offset) - query := fmt.Sprintf("SELECT %s FROM usage_logs %s ORDER BY %s LIMIT $%d OFFSET $%d", usageLogSelectColumns, whereClause, usageLogOrderBy(params), limitPos, offsetPos) - - logs, err := r.queryUsageLogs(ctx, query, listArgs...) - if err != nil { - return nil, nil, err - } - - hasMore := false - if len(logs) > limit { - hasMore = true - logs = logs[:limit] - } - - total := int64(offset) + int64(len(logs)) - if hasMore { - // 只保证“还有下一页”,避免对超大表做全量 COUNT(*)。 - total = int64(offset) + int64(limit) + 1 - } - - return logs, paginationResultFromTotal(total, params), nil -} - -func usageLogOrderBy(params pagination.PaginationParams) string { - sortBy := strings.ToLower(strings.TrimSpace(params.SortBy)) - sortOrder := strings.ToUpper(params.NormalizedSortOrder(pagination.SortOrderDesc)) - - var column string - switch sortBy { - case "model": - column = "COALESCE(NULLIF(TRIM(requested_model), ''), model)" - case "created_at": - column = "created_at" - default: - column = "id" - } - - if column == "id" { - return fmt.Sprintf("id %s", sortOrder) - } - return fmt.Sprintf("%s %s, id %s", column, sortOrder, sortOrder) -} - -func (r *usageLogRepository) queryUsageLogs(ctx context.Context, query string, args ...any) (logs []service.UsageLog, err error) { - rows, err := r.sql.QueryContext(ctx, query, args...) - if err != nil { - return nil, err - } - defer func() { - // 保持主错误优先;仅在无错误时回传 Close 失败。 - // 同时清空返回值,避免误用不完整结果。 - if closeErr := rows.Close(); closeErr != nil && err == nil { - err = closeErr - logs = nil - } - }() - - logs = make([]service.UsageLog, 0) - for rows.Next() { - var log *service.UsageLog - log, err = scanUsageLog(rows) - if err != nil { - return nil, err - } - logs = append(logs, *log) - } - if err = rows.Err(); err != nil { - return nil, err - } - return logs, nil -} - -func (r *usageLogRepository) hydrateUsageLogAssociations(ctx context.Context, logs []service.UsageLog) error { - // 关联数据使用 Ent 批量加载,避免把复杂 SQL 继续膨胀。 - if len(logs) == 0 { - return nil - } - - ids := collectUsageLogIDs(logs) - users, err := r.loadUsers(ctx, ids.userIDs) - if err != nil { - return err - } - apiKeys, err := r.loadAPIKeys(ctx, ids.apiKeyIDs) - if err != nil { - return err - } - accounts, err := r.loadAccounts(ctx, ids.accountIDs) - if err != nil { - return err - } - groups, err := r.loadGroups(ctx, ids.groupIDs) - if err != nil { - return err - } - subs, err := r.loadSubscriptions(ctx, ids.subscriptionIDs) - if err != nil { - return err - } - - for i := range logs { - if user, ok := users[logs[i].UserID]; ok { - logs[i].User = user - } - if key, ok := apiKeys[logs[i].APIKeyID]; ok { - logs[i].APIKey = key - } - if acc, ok := accounts[logs[i].AccountID]; ok { - logs[i].Account = acc - } - if logs[i].GroupID != nil { - if group, ok := groups[*logs[i].GroupID]; ok { - logs[i].Group = group - } - } - if logs[i].SubscriptionID != nil { - if sub, ok := subs[*logs[i].SubscriptionID]; ok { - logs[i].Subscription = sub - } - } - } - return nil -} - -type usageLogIDs struct { - userIDs []int64 - apiKeyIDs []int64 - accountIDs []int64 - groupIDs []int64 - subscriptionIDs []int64 -} - -func collectUsageLogIDs(logs []service.UsageLog) usageLogIDs { - idSet := func() map[int64]struct{} { return make(map[int64]struct{}) } - - userIDs := idSet() - apiKeyIDs := idSet() - accountIDs := idSet() - groupIDs := idSet() - subscriptionIDs := idSet() - - for i := range logs { - userIDs[logs[i].UserID] = struct{}{} - apiKeyIDs[logs[i].APIKeyID] = struct{}{} - accountIDs[logs[i].AccountID] = struct{}{} - if logs[i].GroupID != nil { - groupIDs[*logs[i].GroupID] = struct{}{} - } - if logs[i].SubscriptionID != nil { - subscriptionIDs[*logs[i].SubscriptionID] = struct{}{} - } - } - - return usageLogIDs{ - userIDs: setToSlice(userIDs), - apiKeyIDs: setToSlice(apiKeyIDs), - accountIDs: setToSlice(accountIDs), - groupIDs: setToSlice(groupIDs), - subscriptionIDs: setToSlice(subscriptionIDs), - } -} - -func (r *usageLogRepository) loadUsers(ctx context.Context, ids []int64) (map[int64]*service.User, error) { - out := make(map[int64]*service.User) - if len(ids) == 0 { - return out, nil - } - // 无条件穿透软删除:ids 来自调用方已按 user_id 筛选的日志行;普通用户路径强制 UserID=本人(本人必为活跃用户),不会借此解析他人已删身份;仅 admin 路径可借此显示已删用户。 - models, err := r.client.User.Query().Where(dbuser.IDIn(ids...)).All(mixins.SkipSoftDelete(ctx)) - if err != nil { - return nil, err - } - for _, m := range models { - out[m.ID] = userEntityToService(m) - } - return out, nil -} - -func (r *usageLogRepository) loadAPIKeys(ctx context.Context, ids []int64) (map[int64]*service.APIKey, error) { - out := make(map[int64]*service.APIKey) - if len(ids) == 0 { - return out, nil - } - models, err := r.client.APIKey.Query().Where(dbapikey.IDIn(ids...)).All(ctx) - if err != nil { - return nil, err - } - for _, m := range models { - out[m.ID] = apiKeyEntityToService(m) - } - return out, nil -} - -func (r *usageLogRepository) loadAccounts(ctx context.Context, ids []int64) (map[int64]*service.Account, error) { - out := make(map[int64]*service.Account) - if len(ids) == 0 { - return out, nil - } - models, err := r.client.Account.Query().Where(dbaccount.IDIn(ids...)).All(ctx) - if err != nil { - return nil, err - } - for _, m := range models { - out[m.ID] = accountEntityToService(m) - } - return out, nil -} - -func (r *usageLogRepository) loadGroups(ctx context.Context, ids []int64) (map[int64]*service.Group, error) { - out := make(map[int64]*service.Group) - if len(ids) == 0 { - return out, nil - } - models, err := r.client.Group.Query().Where(dbgroup.IDIn(ids...)).All(ctx) - if err != nil { - return nil, err - } - for _, m := range models { - out[m.ID] = groupEntityToService(m) - } - return out, nil -} - -func (r *usageLogRepository) loadSubscriptions(ctx context.Context, ids []int64) (map[int64]*service.UserSubscription, error) { - out := make(map[int64]*service.UserSubscription) - if len(ids) == 0 { - return out, nil - } - models, err := r.client.UserSubscription.Query().Where(dbusersub.IDIn(ids...)).All(ctx) - if err != nil { - return nil, err - } - for _, m := range models { - out[m.ID] = userSubscriptionEntityToService(m) - } - return out, nil -} - -func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, error) { - var ( - id int64 - userID int64 - apiKeyID int64 - accountID int64 - requestID sql.NullString - model string - requestedModel sql.NullString - upstreamModel sql.NullString - groupID sql.NullInt64 - subscriptionID sql.NullInt64 - inputTokens int - outputTokens int - cacheCreationTokens int - cacheReadTokens int - cacheCreation5m int - cacheCreation1h int - imageOutputTokens int - imageOutputCost float64 - inputCost float64 - outputCost float64 - cacheCreationCost float64 - cacheReadCost float64 - totalCost float64 - actualCost float64 - rateMultiplier float64 - accountRateMultiplier sql.NullFloat64 - billingType int16 - requestTypeRaw int16 - stream bool - openaiWSMode bool - durationMs sql.NullInt64 - firstTokenMs sql.NullInt64 - userAgent sql.NullString - ipAddress sql.NullString - imageCount int - imageSize sql.NullString - imageInputSize sql.NullString - imageOutputSize sql.NullString - imageSizeSource sql.NullString - imageSizeBreakdown sql.NullString - serviceTier sql.NullString - reasoningEffort sql.NullString - inboundEndpoint sql.NullString - upstreamEndpoint sql.NullString - cacheTTLOverridden bool - channelID sql.NullInt64 - modelMappingChain sql.NullString - billingTier sql.NullString - billingMode sql.NullString - accountStatsCost sql.NullFloat64 - createdAt time.Time - ) - - if err := scanner.Scan( - &id, - &userID, - &apiKeyID, - &accountID, - &requestID, - &model, - &requestedModel, - &upstreamModel, - &groupID, - &subscriptionID, - &inputTokens, - &outputTokens, - &cacheCreationTokens, - &cacheReadTokens, - &cacheCreation5m, - &cacheCreation1h, - &imageOutputTokens, - &imageOutputCost, - &inputCost, - &outputCost, - &cacheCreationCost, - &cacheReadCost, - &totalCost, - &actualCost, - &rateMultiplier, - &accountRateMultiplier, - &billingType, - &requestTypeRaw, - &stream, - &openaiWSMode, - &durationMs, - &firstTokenMs, - &userAgent, - &ipAddress, - &imageCount, - &imageSize, - &imageInputSize, - &imageOutputSize, - &imageSizeSource, - &imageSizeBreakdown, - &serviceTier, - &reasoningEffort, - &inboundEndpoint, - &upstreamEndpoint, - &cacheTTLOverridden, - &channelID, - &modelMappingChain, - &billingTier, - &billingMode, - &accountStatsCost, - &createdAt, - ); err != nil { - return nil, err - } - - log := &service.UsageLog{ - ID: id, - UserID: userID, - APIKeyID: apiKeyID, - AccountID: accountID, - Model: model, - RequestedModel: coalesceTrimmedString(requestedModel, model), - InputTokens: inputTokens, - OutputTokens: outputTokens, - CacheCreationTokens: cacheCreationTokens, - CacheReadTokens: cacheReadTokens, - CacheCreation5mTokens: cacheCreation5m, - CacheCreation1hTokens: cacheCreation1h, - ImageOutputTokens: imageOutputTokens, - ImageOutputCost: imageOutputCost, - InputCost: inputCost, - OutputCost: outputCost, - CacheCreationCost: cacheCreationCost, - CacheReadCost: cacheReadCost, - TotalCost: totalCost, - ActualCost: actualCost, - RateMultiplier: rateMultiplier, - AccountRateMultiplier: nullFloat64Ptr(accountRateMultiplier), - BillingType: int8(billingType), - RequestType: service.RequestTypeFromInt16(requestTypeRaw), - ImageCount: imageCount, - CacheTTLOverridden: cacheTTLOverridden, - CreatedAt: createdAt, - } - // 先回填 legacy 字段,再基于 legacy + request_type 计算最终请求类型,保证历史数据兼容。 - log.Stream = stream - log.OpenAIWSMode = openaiWSMode - log.RequestType = log.EffectiveRequestType() - log.Stream, log.OpenAIWSMode = service.ApplyLegacyRequestFields(log.RequestType, stream, openaiWSMode) - - if requestID.Valid { - log.RequestID = requestID.String - } - if groupID.Valid { - value := groupID.Int64 - log.GroupID = &value - } - if subscriptionID.Valid { - value := subscriptionID.Int64 - log.SubscriptionID = &value - } - if durationMs.Valid { - value := int(durationMs.Int64) - log.DurationMs = &value - } - if firstTokenMs.Valid { - value := int(firstTokenMs.Int64) - log.FirstTokenMs = &value - } - if userAgent.Valid { - log.UserAgent = &userAgent.String - } - if ipAddress.Valid { - log.IPAddress = &ipAddress.String - } - if imageSize.Valid { - log.ImageSize = &imageSize.String - } - if imageInputSize.Valid { - log.ImageInputSize = &imageInputSize.String - } - if imageOutputSize.Valid { - log.ImageOutputSize = &imageOutputSize.String - } - if imageSizeSource.Valid { - log.ImageSizeSource = &imageSizeSource.String - } - log.ImageSizeBreakdown = stringIntMapFromNullJSON(imageSizeBreakdown) - if serviceTier.Valid { - log.ServiceTier = &serviceTier.String - } - if reasoningEffort.Valid { - log.ReasoningEffort = &reasoningEffort.String - } - if inboundEndpoint.Valid { - log.InboundEndpoint = &inboundEndpoint.String - } - if upstreamEndpoint.Valid { - log.UpstreamEndpoint = &upstreamEndpoint.String - } - if upstreamModel.Valid { - log.UpstreamModel = &upstreamModel.String - } - if channelID.Valid { - value := channelID.Int64 - log.ChannelID = &value - } - if modelMappingChain.Valid { - log.ModelMappingChain = &modelMappingChain.String - } - if billingTier.Valid { - log.BillingTier = &billingTier.String - } - if billingMode.Valid { - log.BillingMode = &billingMode.String - } - if accountStatsCost.Valid { - log.AccountStatsCost = &accountStatsCost.Float64 - } - - return log, nil -} - -func scanTrendRows(rows *sql.Rows) ([]TrendDataPoint, error) { - results := make([]TrendDataPoint, 0) - for rows.Next() { - var row TrendDataPoint - if err := rows.Scan( - &row.Date, - &row.Requests, - &row.InputTokens, - &row.OutputTokens, - &row.CacheCreationTokens, - &row.CacheReadTokens, - &row.TotalTokens, - &row.Cost, - &row.ActualCost, - ); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - -func scanModelStatsRows(rows *sql.Rows) ([]ModelStat, error) { - results := make([]ModelStat, 0) - for rows.Next() { - var row ModelStat - if err := rows.Scan( - &row.Model, - &row.Requests, - &row.InputTokens, - &row.OutputTokens, - &row.CacheCreationTokens, - &row.CacheReadTokens, - &row.TotalTokens, - &row.Cost, - &row.ActualCost, - &row.AccountCost, - ); err != nil { - return nil, err - } - results = append(results, row) - } - if err := rows.Err(); err != nil { - return nil, err - } - return results, nil -} - func buildWhere(conditions []string) string { if len(conditions) == 0 { return "" @@ -4630,72 +212,3 @@ func buildRequestTypeFilterCondition(startArgIndex int, requestType int16) (stri return fmt.Sprintf("request_type = $%d", startArgIndex), []any{requestTypeArg} } } - -func nullInt64(v *int64) sql.NullInt64 { - if v == nil { - return sql.NullInt64{} - } - return sql.NullInt64{Int64: *v, Valid: true} -} - -func nullInt(v *int) sql.NullInt64 { - if v == nil { - return sql.NullInt64{} - } - return sql.NullInt64{Int64: int64(*v), Valid: true} -} - -func nullFloat64Ptr(v sql.NullFloat64) *float64 { - if !v.Valid { - return nil - } - out := v.Float64 - return &out -} - -func nullString(v *string) sql.NullString { - if v == nil || *v == "" { - return sql.NullString{} - } - return sql.NullString{String: *v, Valid: true} -} - -func nullStringIntMapJSON(v map[string]int) any { - if len(v) == 0 { - return nil - } - payload, err := json.Marshal(v) - if err != nil { - return nil - } - return string(payload) -} - -func stringIntMapFromNullJSON(v sql.NullString) map[string]int { - if !v.Valid || strings.TrimSpace(v.String) == "" { - return nil - } - var out map[string]int - if err := json.Unmarshal([]byte(v.String), &out); err != nil { - return nil - } - if len(out) == 0 { - return nil - } - return out -} - -func coalesceTrimmedString(v sql.NullString, fallback string) string { - if v.Valid && strings.TrimSpace(v.String) != "" { - return v.String - } - return fallback -} - -func setToSlice(set map[int64]struct{}) []int64 { - out := make([]int64, 0, len(set)) - for id := range set { - out = append(out, id) - } - return out -} diff --git a/backend/internal/repository/usage_log_repo_dashboard.go b/backend/internal/repository/usage_log_repo_dashboard.go new file mode 100644 index 0000000000..a56213357e --- /dev/null +++ b/backend/internal/repository/usage_log_repo_dashboard.go @@ -0,0 +1,628 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/Wei-Shaw/sub2api/internal/service" +) + +// getPerformanceStats 获取 RPM 和 TPM(近5分钟平均值,可选按用户过滤) +func (r *usageLogRepository) getPerformanceStats(ctx context.Context, userID int64) (rpm, tpm int64, err error) { + fiveMinutesAgo := time.Now().Add(-5 * time.Minute) + query := ` + SELECT + COUNT(*) as request_count, + COALESCE(SUM(input_tokens + output_tokens), 0) as token_count + FROM usage_logs + WHERE created_at >= $1` + args := []any{fiveMinutesAgo} + if userID > 0 { + query += " AND user_id = $2" + args = append(args, userID) + } + + var requestCount int64 + var tokenCount int64 + if err := scanSingleRow(ctx, r.sql, query, args, &requestCount, &tokenCount); err != nil { + return 0, 0, err + } + return requestCount / 5, tokenCount / 5, nil +} + +// UserStats 用户使用统计 +type UserStats struct { + TotalRequests int64 `json:"total_requests"` + TotalTokens int64 `json:"total_tokens"` + TotalCost float64 `json:"total_cost"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + CacheReadTokens int64 `json:"cache_read_tokens"` +} + +func (r *usageLogRepository) GetUserStats(ctx context.Context, userID int64, startTime, endTime time.Time) (*UserStats, error) { + query := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(actual_cost), 0) as total_cost, + COALESCE(SUM(input_tokens), 0) as input_tokens, + COALESCE(SUM(output_tokens), 0) as output_tokens, + COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens + FROM usage_logs + WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 + ` + + stats := &UserStats{} + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{userID, startTime, endTime}, + &stats.TotalRequests, + &stats.TotalTokens, + &stats.TotalCost, + &stats.InputTokens, + &stats.OutputTokens, + &stats.CacheReadTokens, + ); err != nil { + return nil, err + } + return stats, nil +} + +// DashboardStats 仪表盘统计 +type DashboardStats = usagestats.DashboardStats + +func (r *usageLogRepository) GetDashboardStats(ctx context.Context) (*DashboardStats, error) { + stats := &DashboardStats{} + now := timezone.Now() + todayStart := timezone.Today() + + if err := r.fillDashboardEntityStats(ctx, stats, todayStart, now); err != nil { + return nil, err + } + if err := r.fillDashboardUsageStatsAggregated(ctx, stats, todayStart, now); err != nil { + return nil, err + } + + rpm, tpm, err := r.getPerformanceStats(ctx, 0) + if err != nil { + return nil, err + } + stats.Rpm = rpm + stats.Tpm = tpm + + return stats, nil +} + +func (r *usageLogRepository) GetDashboardStatsWithRange(ctx context.Context, start, end time.Time) (*DashboardStats, error) { + startUTC := start.UTC() + endUTC := end.UTC() + if !endUTC.After(startUTC) { + return nil, errors.New("统计时间范围无效") + } + + stats := &DashboardStats{} + now := timezone.Now() + todayStart := timezone.Today() + + if err := r.fillDashboardEntityStats(ctx, stats, todayStart, now); err != nil { + return nil, err + } + if err := r.fillDashboardUsageStatsFromUsageLogs(ctx, stats, startUTC, endUTC, todayStart, now); err != nil { + return nil, err + } + + rpm, tpm, err := r.getPerformanceStats(ctx, 0) + if err != nil { + return nil, err + } + stats.Rpm = rpm + stats.Tpm = tpm + + return stats, nil +} + +func (r *usageLogRepository) fillDashboardEntityStats(ctx context.Context, stats *DashboardStats, todayUTC, now time.Time) error { + userStatsQuery := ` + SELECT + COUNT(*) as total_users, + COUNT(CASE WHEN created_at >= $1 THEN 1 END) as today_new_users + FROM users + WHERE deleted_at IS NULL + ` + if err := scanSingleRow( + ctx, + r.sql, + userStatsQuery, + []any{todayUTC}, + &stats.TotalUsers, + &stats.TodayNewUsers, + ); err != nil { + return err + } + + apiKeyStatsQuery := ` + SELECT + COUNT(*) as total_api_keys, + COUNT(CASE WHEN status = $1 THEN 1 END) as active_api_keys + FROM api_keys + WHERE deleted_at IS NULL + ` + if err := scanSingleRow( + ctx, + r.sql, + apiKeyStatsQuery, + []any{service.StatusActive}, + &stats.TotalAPIKeys, + &stats.ActiveAPIKeys, + ); err != nil { + return err + } + + accountStatsQuery := ` + SELECT + COUNT(*) as total_accounts, + COUNT(CASE WHEN status = $1 AND schedulable = true THEN 1 END) as normal_accounts, + COUNT(CASE WHEN status = $2 THEN 1 END) as error_accounts, + COUNT(CASE WHEN rate_limited_at IS NOT NULL AND rate_limit_reset_at > $3 THEN 1 END) as ratelimit_accounts, + COUNT(CASE WHEN overload_until IS NOT NULL AND overload_until > $4 THEN 1 END) as overload_accounts + FROM accounts + WHERE deleted_at IS NULL + ` + if err := scanSingleRow( + ctx, + r.sql, + accountStatsQuery, + []any{service.StatusActive, service.StatusError, now, now}, + &stats.TotalAccounts, + &stats.NormalAccounts, + &stats.ErrorAccounts, + &stats.RateLimitAccounts, + &stats.OverloadAccounts, + ); err != nil { + return err + } + + return nil +} + +func (r *usageLogRepository) fillDashboardUsageStatsAggregated(ctx context.Context, stats *DashboardStats, todayUTC, now time.Time) error { + totalStatsQuery := ` + SELECT + COALESCE(SUM(total_requests), 0) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(SUM(account_cost), 0) as total_account_cost, + COALESCE(SUM(total_duration_ms), 0) as total_duration_ms + FROM usage_dashboard_daily + ` + var totalDurationMs int64 + if err := scanSingleRow( + ctx, + r.sql, + totalStatsQuery, + nil, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.TotalAccountCost, + &totalDurationMs, + ); err != nil { + return err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens + if stats.TotalRequests > 0 { + stats.AverageDurationMs = float64(totalDurationMs) / float64(stats.TotalRequests) + } + + todayStatsQuery := ` + SELECT + total_requests as today_requests, + input_tokens as today_input_tokens, + output_tokens as today_output_tokens, + cache_creation_tokens as today_cache_creation_tokens, + cache_read_tokens as today_cache_read_tokens, + total_cost as today_cost, + actual_cost as today_actual_cost, + account_cost as today_account_cost, + active_users as active_users + FROM usage_dashboard_daily + WHERE bucket_date = $1::date + ` + if err := scanSingleRow( + ctx, + r.sql, + todayStatsQuery, + []any{todayUTC}, + &stats.TodayRequests, + &stats.TodayInputTokens, + &stats.TodayOutputTokens, + &stats.TodayCacheCreationTokens, + &stats.TodayCacheReadTokens, + &stats.TodayCost, + &stats.TodayActualCost, + &stats.TodayAccountCost, + &stats.ActiveUsers, + ); err != nil { + if err != sql.ErrNoRows { + return err + } + } + stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens + + hourlyActiveQuery := ` + SELECT active_users + FROM usage_dashboard_hourly + WHERE bucket_start = $1 + ` + hourStart := now.In(timezone.Location()).Truncate(time.Hour) + if err := scanSingleRow(ctx, r.sql, hourlyActiveQuery, []any{hourStart}, &stats.HourlyActiveUsers); err != nil { + if err != sql.ErrNoRows { + return err + } + } + + return nil +} + +func (r *usageLogRepository) fillDashboardUsageStatsFromUsageLogs(ctx context.Context, stats *DashboardStats, startUTC, endUTC, todayUTC, now time.Time) error { + todayEnd := todayUTC.Add(24 * time.Hour) + combinedStatsQuery := ` + WITH scoped AS ( + SELECT + created_at, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + total_cost, + actual_cost, + COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1) AS account_cost, + COALESCE(duration_ms, 0) AS duration_ms + FROM usage_logs + WHERE created_at >= LEAST($1::timestamptz, $3::timestamptz) + AND created_at < GREATEST($2::timestamptz, $4::timestamptz) + ) + SELECT + COUNT(*) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz) AS total_requests, + COALESCE(SUM(input_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_input_tokens, + COALESCE(SUM(output_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_output_tokens, + COALESCE(SUM(cache_creation_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_cache_read_tokens, + COALESCE(SUM(total_cost) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_cost, + COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_actual_cost, + COALESCE(SUM(account_cost) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_account_cost, + COALESCE(SUM(duration_ms) FILTER (WHERE created_at >= $1::timestamptz AND created_at < $2::timestamptz), 0) AS total_duration_ms, + COUNT(*) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz) AS today_requests, + COALESCE(SUM(input_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_input_tokens, + COALESCE(SUM(output_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_output_tokens, + COALESCE(SUM(cache_creation_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_cache_read_tokens, + COALESCE(SUM(total_cost) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_cost, + COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_actual_cost, + COALESCE(SUM(account_cost) FILTER (WHERE created_at >= $3::timestamptz AND created_at < $4::timestamptz), 0) AS today_account_cost + FROM scoped + ` + var totalDurationMs int64 + if err := scanSingleRow( + ctx, + r.sql, + combinedStatsQuery, + []any{startUTC, endUTC, todayUTC, todayEnd}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.TotalAccountCost, + &totalDurationMs, + &stats.TodayRequests, + &stats.TodayInputTokens, + &stats.TodayOutputTokens, + &stats.TodayCacheCreationTokens, + &stats.TodayCacheReadTokens, + &stats.TodayCost, + &stats.TodayActualCost, + &stats.TodayAccountCost, + ); err != nil { + return err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens + if stats.TotalRequests > 0 { + stats.AverageDurationMs = float64(totalDurationMs) / float64(stats.TotalRequests) + } + + stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens + + hourStart := now.UTC().Truncate(time.Hour) + hourEnd := hourStart.Add(time.Hour) + activeUsersQuery := ` + WITH scoped AS ( + SELECT user_id, created_at + FROM usage_logs + WHERE created_at >= LEAST($1::timestamptz, $3::timestamptz) + AND created_at < GREATEST($2::timestamptz, $4::timestamptz) + ) + SELECT + COUNT(DISTINCT CASE WHEN created_at >= $1::timestamptz AND created_at < $2::timestamptz THEN user_id END) AS active_users, + COUNT(DISTINCT CASE WHEN created_at >= $3::timestamptz AND created_at < $4::timestamptz THEN user_id END) AS hourly_active_users + FROM scoped + ` + if err := scanSingleRow(ctx, r.sql, activeUsersQuery, []any{todayUTC, todayEnd, hourStart, hourEnd}, &stats.ActiveUsers, &stats.HourlyActiveUsers); err != nil { + return err + } + + return nil +} + +// UserDashboardStats 用户仪表盘统计 +type UserDashboardStats = usagestats.UserDashboardStats + +// PlatformDashboardStats 单平台用量明细 +type PlatformDashboardStats = usagestats.PlatformDashboardStats + +// GetUserDashboardStats 获取用户专属的仪表盘统计 +func (r *usageLogRepository) GetUserDashboardStats(ctx context.Context, userID int64) (*UserDashboardStats, error) { + stats := &UserDashboardStats{} + today := timezone.Today() + + // API Key 统计 + if err := scanSingleRow( + ctx, + r.sql, + "SELECT COUNT(*) FROM api_keys WHERE user_id = $1 AND deleted_at IS NULL", + []any{userID}, + &stats.TotalAPIKeys, + ); err != nil { + return nil, err + } + if err := scanSingleRow( + ctx, + r.sql, + "SELECT COUNT(*) FROM api_keys WHERE user_id = $1 AND status = $2 AND deleted_at IS NULL", + []any{userID, service.StatusActive}, + &stats.ActiveAPIKeys, + ); err != nil { + return nil, err + } + + // 累计 Token 统计 + totalStatsQuery := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(duration_ms), 0) as avg_duration_ms + FROM usage_logs + WHERE user_id = $1 + ` + if err := scanSingleRow( + ctx, + r.sql, + totalStatsQuery, + []any{userID}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens + + // 今日 Token 统计 + todayStatsQuery := ` + SELECT + COUNT(*) as today_requests, + COALESCE(SUM(input_tokens), 0) as today_input_tokens, + COALESCE(SUM(output_tokens), 0) as today_output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as today_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as today_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as today_cost, + COALESCE(SUM(actual_cost), 0) as today_actual_cost + FROM usage_logs + WHERE user_id = $1 AND created_at >= $2 + ` + if err := scanSingleRow( + ctx, + r.sql, + todayStatsQuery, + []any{userID, today}, + &stats.TodayRequests, + &stats.TodayInputTokens, + &stats.TodayOutputTokens, + &stats.TodayCacheCreationTokens, + &stats.TodayCacheReadTokens, + &stats.TodayCost, + &stats.TodayActualCost, + ); err != nil { + return nil, err + } + stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens + + // 性能指标:RPM 和 TPM(最近1分钟,仅统计该用户的请求) + rpm, tpm, err := r.getPerformanceStats(ctx, userID) + if err != nil { + return nil, err + } + stats.Rpm = rpm + stats.Tpm = tpm + + // 按"有效平台"维度拆分(group.platform 优先,否则 account.platform)。 + // 与 ops 路径口径一致;HAVING 过滤掉无法确定平台的行(避免出现空字符串平台)。 + // 与上面 totalStatsQuery/todayStatsQuery 的总值可能略微差异,原因有二: + // 1) 无平台归属的极少数行(group/account 都没 platform)会被 HAVING 排除; + // 2) usageLogSuccessFilterUL 会把 actual_cost = 0 的失败 placeholder 行排除, + // 而 totalStatsQuery/todayStatsQuery 没有这层过滤、会把这些行的 request 计数算进去。 + platformQuery := ` + SELECT + ` + usageLogEffectivePlatformExpr + ` as platform, + COUNT(*) as total_requests, + COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(ul.actual_cost), 0) as total_actual_cost, + COUNT(*) FILTER (WHERE ul.created_at >= $2) as today_requests, + COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens) FILTER (WHERE ul.created_at >= $2), 0) as today_tokens, + COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $2), 0) as today_actual_cost + FROM usage_logs ul + LEFT JOIN groups g ON g.id = ul.group_id + LEFT JOIN accounts a ON a.id = ul.account_id + WHERE ul.user_id = $1 + AND ` + usageLogSuccessFilterUL + ` + GROUP BY ` + usageLogEffectivePlatformExpr + ` + HAVING ` + usageLogEffectivePlatformExpr + ` IS NOT NULL AND ` + usageLogEffectivePlatformExpr + ` <> '' + ORDER BY total_actual_cost DESC + ` + rows, err := r.sql.QueryContext(ctx, platformQuery, userID, today) + if err != nil { + return nil, err + } + for rows.Next() { + var p PlatformDashboardStats + if err := rows.Scan( + &p.Platform, + &p.TotalRequests, + &p.TotalTokens, + &p.TotalActualCost, + &p.TodayRequests, + &p.TodayTokens, + &p.TodayActualCost, + ); err != nil { + _ = rows.Close() + return nil, err + } + stats.ByPlatform = append(stats.ByPlatform, p) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + + return stats, nil +} + +// getPerformanceStatsByAPIKey 获取指定 API Key 的 RPM 和 TPM(近5分钟平均值) +func (r *usageLogRepository) getPerformanceStatsByAPIKey(ctx context.Context, apiKeyID int64) (rpm, tpm int64, err error) { + fiveMinutesAgo := time.Now().Add(-5 * time.Minute) + query := ` + SELECT + COUNT(*) as request_count, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as token_count + FROM usage_logs + WHERE created_at >= $1 AND api_key_id = $2` + args := []any{fiveMinutesAgo, apiKeyID} + + var requestCount int64 + var tokenCount int64 + if err := scanSingleRow(ctx, r.sql, query, args, &requestCount, &tokenCount); err != nil { + return 0, 0, err + } + return requestCount / 5, tokenCount / 5, nil +} + +// GetAPIKeyDashboardStats 获取指定 API Key 的仪表盘统计(按 api_key_id 过滤) +func (r *usageLogRepository) GetAPIKeyDashboardStats(ctx context.Context, apiKeyID int64) (*UserDashboardStats, error) { + stats := &UserDashboardStats{} + today := timezone.Today() + + // API Key 维度不需要统计 key 数量,设为 1 + stats.TotalAPIKeys = 1 + stats.ActiveAPIKeys = 1 + + // 累计 Token 统计 + totalStatsQuery := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(duration_ms), 0) as avg_duration_ms + FROM usage_logs + WHERE api_key_id = $1 + ` + if err := scanSingleRow( + ctx, + r.sql, + totalStatsQuery, + []any{apiKeyID}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheCreationTokens + stats.TotalCacheReadTokens + + // 今日 Token 统计 + todayStatsQuery := ` + SELECT + COUNT(*) as today_requests, + COALESCE(SUM(input_tokens), 0) as today_input_tokens, + COALESCE(SUM(output_tokens), 0) as today_output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as today_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as today_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as today_cost, + COALESCE(SUM(actual_cost), 0) as today_actual_cost + FROM usage_logs + WHERE api_key_id = $1 AND created_at >= $2 + ` + if err := scanSingleRow( + ctx, + r.sql, + todayStatsQuery, + []any{apiKeyID, today}, + &stats.TodayRequests, + &stats.TodayInputTokens, + &stats.TodayOutputTokens, + &stats.TodayCacheCreationTokens, + &stats.TodayCacheReadTokens, + &stats.TodayCost, + &stats.TodayActualCost, + ); err != nil { + return nil, err + } + stats.TodayTokens = stats.TodayInputTokens + stats.TodayOutputTokens + stats.TodayCacheCreationTokens + stats.TodayCacheReadTokens + + // 性能指标:RPM 和 TPM(最近5分钟,按 API Key 过滤) + rpm, tpm, err := r.getPerformanceStatsByAPIKey(ctx, apiKeyID) + if err != nil { + return nil, err + } + stats.Rpm = rpm + stats.Tpm = tpm + + return stats, nil +} diff --git a/backend/internal/repository/usage_log_repo_insert.go b/backend/internal/repository/usage_log_repo_insert.go new file mode 100644 index 0000000000..dfd8969512 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_insert.go @@ -0,0 +1,1297 @@ +package repository + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + "sync/atomic" + "time" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/service" +) + +// usageLogInsertArgTypes must stay in the same order as: +// 1. prepareUsageLogInsert().args +// 2. every INSERT/CTE VALUES column list in this file +// 3. execUsageLogInsertNoResult placeholder positions +// 4. scanUsageLog selected column order (via usageLogSelectColumns) +// +// When adding a usage_logs column, update all of those call sites together. +var usageLogInsertArgTypes = [...]string{ + "bigint", // user_id + "bigint", // api_key_id + "bigint", // account_id + "text", // request_id + "text", // model + "text", // requested_model + "text", // upstream_model + "bigint", // group_id + "bigint", // subscription_id + "integer", // input_tokens + "integer", // output_tokens + "integer", // cache_creation_tokens + "integer", // cache_read_tokens + "integer", // cache_creation_5m_tokens + "integer", // cache_creation_1h_tokens + "integer", // image_output_tokens + "numeric", // image_output_cost + "numeric", // input_cost + "numeric", // output_cost + "numeric", // cache_creation_cost + "numeric", // cache_read_cost + "numeric", // total_cost + "numeric", // actual_cost + "numeric", // rate_multiplier + "numeric", // account_rate_multiplier + "smallint", // billing_type + "smallint", // request_type + "boolean", // stream + "boolean", // openai_ws_mode + "integer", // duration_ms + "integer", // first_token_ms + "text", // user_agent + "text", // ip_address + "integer", // image_count + "text", // image_size + "text", // image_input_size + "text", // image_output_size + "text", // image_size_source + "jsonb", // image_size_breakdown + "integer", // video_count + "text", // video_resolution + "integer", // video_duration_seconds + "text", // service_tier + "text", // reasoning_effort + "text", // inbound_endpoint + "text", // upstream_endpoint + "boolean", // cache_ttl_overridden + "bigint", // channel_id + "text", // model_mapping_chain + "text", // billing_tier + "text", // billing_mode + "numeric", // account_stats_cost + "timestamptz", // created_at +} + +const ( + usageLogCreateBatchMaxSize = 64 + usageLogCreateBatchWindow = 3 * time.Millisecond + usageLogCreateBatchQueueCap = 4096 + usageLogCreateCancelWait = 2 * time.Second + + usageLogBestEffortBatchMaxSize = 256 + usageLogBestEffortBatchWindow = 20 * time.Millisecond + usageLogBestEffortBatchQueueCap = 32768 + usageLogBestEffortRecentTTL = 30 * time.Second +) + +type usageLogCreateRequest struct { + log *service.UsageLog + prepared usageLogInsertPrepared + shared *usageLogCreateShared + resultCh chan usageLogCreateResult +} + +type usageLogCreateResult struct { + inserted bool + err error +} + +type usageLogBestEffortRequest struct { + prepared usageLogInsertPrepared + apiKeyID int64 + resultCh chan error +} + +type usageLogInsertPrepared struct { + createdAt time.Time + requestID string + rateMultiplier float64 + requestType int16 + args []any +} + +type usageLogBatchState struct { + ID int64 + CreatedAt time.Time +} + +type usageLogBatchRow struct { + RequestID string `json:"request_id"` + APIKeyID int64 `json:"api_key_id"` + ID int64 `json:"id"` + CreatedAt time.Time `json:"created_at"` + Inserted bool `json:"inserted"` +} + +type usageLogCreateShared struct { + state atomic.Int32 +} + +const ( + usageLogCreateStateQueued int32 = iota + usageLogCreateStateProcessing + usageLogCreateStateCompleted + usageLogCreateStateCanceled +) + +func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog) (bool, error) { + if log == nil { + return false, nil + } + + if tx := dbent.TxFromContext(ctx); tx != nil { + return r.createSingle(ctx, tx.Client(), log) + } + requestID := strings.TrimSpace(log.RequestID) + if requestID == "" { + return r.createSingle(ctx, r.sql, log) + } + log.RequestID = requestID + return r.createBatched(ctx, log) +} + +func (r *usageLogRepository) CreateBestEffort(ctx context.Context, log *service.UsageLog) error { + if log == nil { + return nil + } + + if tx := dbent.TxFromContext(ctx); tx != nil { + _, err := r.createSingle(ctx, tx.Client(), log) + return err + } + if r.db == nil { + _, err := r.createSingle(ctx, r.sql, log) + return err + } + + r.ensureBestEffortBatcher() + if r.bestEffortBatchCh == nil { + _, err := r.createSingle(ctx, r.sql, log) + return err + } + + req := usageLogBestEffortRequest{ + prepared: prepareUsageLogInsert(log), + apiKeyID: log.APIKeyID, + resultCh: make(chan error, 1), + } + if key, ok := r.bestEffortRecentKey(req.prepared.requestID, req.apiKeyID); ok { + if _, exists := r.bestEffortRecent.Get(key); exists { + return nil + } + } + + // 队列满时阻塞等待而非立即丢弃:批处理器持续排空队列,短暂等待即可入队。 + // 立即丢弃会造成“已扣费但无 usage_log”的永久数据缺口(issue #3656); + // 阻塞上限由调用方 ctx 期限约束,超时后由上层同步兜底。 + select { + case r.bestEffortBatchCh <- req: + case <-ctx.Done(): + return service.MarkUsageLogCreateDropped(ctx.Err()) + } + + select { + case err := <-req.resultCh: + return err + case <-ctx.Done(): + return service.MarkUsageLogCreateDropped(ctx.Err()) + } +} + +func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, log *service.UsageLog) (bool, error) { + prepared := prepareUsageLogInsert(log) + if sqlq == nil { + sqlq = r.sql + } + if ctx != nil && ctx.Err() != nil { + return false, service.MarkUsageLogCreateNotPersisted(ctx.Err()) + } + + query := ` + INSERT INTO usage_logs ( + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + ) VALUES ( + $1, $2, $3, $4, $5, $6, $7, + $8, $9, + $10, $11, $12, $13, + $14, $15, $16, $17, + $18, $19, $20, $21, $22, $23, + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53 + ) + ON CONFLICT (request_id, api_key_id) DO NOTHING + RETURNING id, created_at + ` + + if err := scanSingleRow(ctx, sqlq, query, prepared.args, &log.ID, &log.CreatedAt); err != nil { + if errors.Is(err, sql.ErrNoRows) && prepared.requestID != "" { + selectQuery := "SELECT id, created_at FROM usage_logs WHERE request_id = $1 AND api_key_id = $2" + if err := scanSingleRow(ctx, sqlq, selectQuery, []any{prepared.requestID, log.APIKeyID}, &log.ID, &log.CreatedAt); err != nil { + return false, err + } + log.RateMultiplier = prepared.rateMultiplier + return false, nil + } else { + return false, err + } + } + log.RateMultiplier = prepared.rateMultiplier + return true, nil +} + +func (r *usageLogRepository) createBatched(ctx context.Context, log *service.UsageLog) (bool, error) { + if r.db == nil { + return r.createSingle(ctx, r.sql, log) + } + r.ensureCreateBatcher() + if r.createBatchCh == nil { + return r.createSingle(ctx, r.sql, log) + } + + req := usageLogCreateRequest{ + log: log, + prepared: prepareUsageLogInsert(log), + shared: &usageLogCreateShared{}, + resultCh: make(chan usageLogCreateResult, 1), + } + + // 队列满时阻塞等待而非立即报错:本路径是 best-effort 丢弃后的最后兜底, + // 立即失败会让日志永久丢失;阻塞上限由调用方 ctx 期限约束。 + select { + case r.createBatchCh <- req: + case <-ctx.Done(): + return false, service.MarkUsageLogCreateNotPersisted(ctx.Err()) + } + + select { + case res := <-req.resultCh: + return res.inserted, res.err + case <-ctx.Done(): + if req.shared != nil && req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateCanceled) { + return false, service.MarkUsageLogCreateNotPersisted(ctx.Err()) + } + timer := time.NewTimer(usageLogCreateCancelWait) + defer timer.Stop() + select { + case res := <-req.resultCh: + return res.inserted, res.err + case <-timer.C: + return false, ctx.Err() + } + } +} + +func (r *usageLogRepository) ensureCreateBatcher() { + if r == nil || r.db == nil { + return + } + // nil 检查必须在 Once 内部:在外层做无同步快路径读会与 Once 内的写构成数据竞争。 + r.createBatchOnce.Do(func() { + if r.createBatchCh == nil { + r.createBatchCh = make(chan usageLogCreateRequest, usageLogCreateBatchQueueCap) + go r.runCreateBatcher(r.db) + } + }) +} + +func (r *usageLogRepository) ensureBestEffortBatcher() { + if r == nil || r.db == nil { + return + } + // 同 ensureCreateBatcher:nil 检查放在 Once 内部以避免数据竞争。 + r.bestEffortBatchOnce.Do(func() { + if r.bestEffortBatchCh == nil { + r.bestEffortBatchCh = make(chan usageLogBestEffortRequest, usageLogBestEffortBatchQueueCap) + go r.runBestEffortBatcher(r.db) + } + }) +} + +func (r *usageLogRepository) runCreateBatcher(db *sql.DB) { + for { + first, ok := <-r.createBatchCh + if !ok { + return + } + + batch := make([]usageLogCreateRequest, 0, usageLogCreateBatchMaxSize) + batch = append(batch, first) + + timer := time.NewTimer(usageLogCreateBatchWindow) + batchLoop: + for len(batch) < usageLogCreateBatchMaxSize { + select { + case req, ok := <-r.createBatchCh: + if !ok { + break batchLoop + } + batch = append(batch, req) + case <-timer.C: + break batchLoop + } + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + + r.flushCreateBatch(db, batch) + } +} + +func (r *usageLogRepository) runBestEffortBatcher(db *sql.DB) { + for { + first, ok := <-r.bestEffortBatchCh + if !ok { + return + } + + batch := make([]usageLogBestEffortRequest, 0, usageLogBestEffortBatchMaxSize) + batch = append(batch, first) + + timer := time.NewTimer(usageLogBestEffortBatchWindow) + bestEffortLoop: + for len(batch) < usageLogBestEffortBatchMaxSize { + select { + case req, ok := <-r.bestEffortBatchCh: + if !ok { + break bestEffortLoop + } + batch = append(batch, req) + case <-timer.C: + break bestEffortLoop + } + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + + r.flushBestEffortBatch(db, batch) + } +} + +func (r *usageLogRepository) flushCreateBatch(db *sql.DB, batch []usageLogCreateRequest) { + if len(batch) == 0 { + return + } + + uniqueOrder := make([]string, 0, len(batch)) + preparedByKey := make(map[string]usageLogInsertPrepared, len(batch)) + requestsByKey := make(map[string][]usageLogCreateRequest, len(batch)) + fallback := make([]usageLogCreateRequest, 0) + + for _, req := range batch { + if req.log == nil { + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) + continue + } + if req.shared != nil && !req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateProcessing) { + if req.shared.state.Load() == usageLogCreateStateCanceled { + completeUsageLogCreateRequest(req, usageLogCreateResult{ + inserted: false, + err: service.MarkUsageLogCreateNotPersisted(context.Canceled), + }) + continue + } + } + prepared := req.prepared + if prepared.requestID == "" { + fallback = append(fallback, req) + continue + } + key := usageLogBatchKey(prepared.requestID, req.log.APIKeyID) + if _, exists := requestsByKey[key]; !exists { + uniqueOrder = append(uniqueOrder, key) + preparedByKey[key] = prepared + } + requestsByKey[key] = append(requestsByKey[key], req) + } + + if len(uniqueOrder) > 0 { + insertedMap, stateMap, safeFallback, err := r.batchInsertUsageLogs(db, uniqueOrder, preparedByKey) + if err != nil { + if safeFallback { + for _, key := range uniqueOrder { + fallback = append(fallback, requestsByKey[key]...) + } + } else { + for _, key := range uniqueOrder { + reqs := requestsByKey[key] + state, hasState := stateMap[key] + inserted := insertedMap[key] + for idx, req := range reqs { + req.log.RateMultiplier = preparedByKey[key].rateMultiplier + if hasState { + req.log.ID = state.ID + req.log.CreatedAt = state.CreatedAt + } + switch { + case inserted && idx == 0: + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: true, err: nil}) + case inserted: + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) + case hasState: + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) + case idx == 0: + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: err}) + default: + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil}) + } + } + } + } + } else { + for _, key := range uniqueOrder { + reqs := requestsByKey[key] + state, ok := stateMap[key] + if !ok { + for _, req := range reqs { + completeUsageLogCreateRequest(req, usageLogCreateResult{ + inserted: false, + err: fmt.Errorf("usage log batch state missing for key=%s", key), + }) + } + continue + } + for idx, req := range reqs { + req.log.ID = state.ID + req.log.CreatedAt = state.CreatedAt + req.log.RateMultiplier = preparedByKey[key].rateMultiplier + completeUsageLogCreateRequest(req, usageLogCreateResult{ + inserted: idx == 0 && insertedMap[key], + err: nil, + }) + } + } + } + } + + if len(fallback) == 0 { + return + } + + for _, req := range fallback { + fallbackCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + inserted, err := r.createSingle(fallbackCtx, db, req.log) + cancel() + completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: inserted, err: err}) + } +} + +func (r *usageLogRepository) flushBestEffortBatch(db *sql.DB, batch []usageLogBestEffortRequest) { + if len(batch) == 0 { + return + } + + type bestEffortGroup struct { + prepared usageLogInsertPrepared + apiKeyID int64 + key string + reqs []usageLogBestEffortRequest + } + + groupsByKey := make(map[string]*bestEffortGroup, len(batch)) + groupOrder := make([]*bestEffortGroup, 0, len(batch)) + preparedList := make([]usageLogInsertPrepared, 0, len(batch)) + + for idx, req := range batch { + prepared := req.prepared + key := fmt.Sprintf("__best_effort_%d", idx) + if prepared.requestID != "" { + key = usageLogBatchKey(prepared.requestID, req.apiKeyID) + } + group, exists := groupsByKey[key] + if !exists { + group = &bestEffortGroup{ + prepared: prepared, + apiKeyID: req.apiKeyID, + key: key, + } + groupsByKey[key] = group + groupOrder = append(groupOrder, group) + preparedList = append(preparedList, prepared) + } + group.reqs = append(group.reqs, req) + } + + if len(preparedList) == 0 { + for _, req := range batch { + sendUsageLogBestEffortResult(req.resultCh, nil) + } + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + query, args := buildUsageLogBestEffortInsertQuery(preparedList) + if _, err := db.ExecContext(ctx, query, args...); err != nil { + logger.LegacyPrintf("repository.usage_log", "best-effort batch insert failed: %v", err) + for _, group := range groupOrder { + singleErr := execUsageLogInsertNoResult(ctx, db, group.prepared) + if singleErr != nil { + logger.LegacyPrintf("repository.usage_log", "best-effort single fallback insert failed: %v", singleErr) + } else if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil { + r.bestEffortRecent.SetDefault(group.key, struct{}{}) + } + for _, req := range group.reqs { + sendUsageLogBestEffortResult(req.resultCh, singleErr) + } + } + return + } + for _, group := range groupOrder { + if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil { + r.bestEffortRecent.SetDefault(group.key, struct{}{}) + } + for _, req := range group.reqs { + sendUsageLogBestEffortResult(req.resultCh, nil) + } + } +} + +func sendUsageLogBestEffortResult(ch chan error, err error) { + if ch == nil { + return + } + select { + case ch <- err: + default: + } +} + +func completeUsageLogCreateRequest(req usageLogCreateRequest, res usageLogCreateResult) { + if req.shared != nil { + req.shared.state.Store(usageLogCreateStateCompleted) + } + sendUsageLogCreateResult(req.resultCh, res) +} + +func (r *usageLogRepository) batchInsertUsageLogs(db *sql.DB, keys []string, preparedByKey map[string]usageLogInsertPrepared) (map[string]bool, map[string]usageLogBatchState, bool, error) { + if len(keys) == 0 { + return map[string]bool{}, map[string]usageLogBatchState{}, false, nil + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + query, args := buildUsageLogBatchInsertQuery(keys, preparedByKey) + var payload []byte + if err := db.QueryRowContext(ctx, query, args...).Scan(&payload); err != nil { + return nil, nil, true, err + } + var rows []usageLogBatchRow + if err := json.Unmarshal(payload, &rows); err != nil { + return nil, nil, false, err + } + insertedMap := make(map[string]bool, len(keys)) + stateMap := make(map[string]usageLogBatchState, len(keys)) + for _, row := range rows { + key := usageLogBatchKey(row.RequestID, row.APIKeyID) + insertedMap[key] = row.Inserted + stateMap[key] = usageLogBatchState{ + ID: row.ID, + CreatedAt: row.CreatedAt, + } + } + if len(stateMap) != len(keys) { + return insertedMap, stateMap, false, fmt.Errorf("usage log batch state count mismatch: got=%d want=%d", len(stateMap), len(keys)) + } + return insertedMap, stateMap, false, nil +} + +func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usageLogInsertPrepared) (string, []any) { + var query strings.Builder + _, _ = query.WriteString(` + WITH input ( + input_idx, + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + ) AS (VALUES `) + + args := make([]any, 0, len(keys)*53) + argPos := 1 + for idx, key := range keys { + if idx > 0 { + _, _ = query.WriteString(",") + } + _, _ = query.WriteString("(") + _, _ = query.WriteString("$") + _, _ = query.WriteString(strconv.Itoa(argPos)) + args = append(args, idx) + argPos++ + prepared := preparedByKey[key] + for i := 0; i < len(prepared.args); i++ { + _, _ = query.WriteString(",") + _, _ = query.WriteString("$") + _, _ = query.WriteString(strconv.Itoa(argPos)) + if i < len(usageLogInsertArgTypes) { + _, _ = query.WriteString("::") + _, _ = query.WriteString(usageLogInsertArgTypes[i]) + } + argPos++ + } + _, _ = query.WriteString(")") + args = append(args, prepared.args...) + } + _, _ = query.WriteString(` + ), + inserted AS ( + INSERT INTO usage_logs ( + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + ) + SELECT + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + FROM input + ON CONFLICT (request_id, api_key_id) DO NOTHING + RETURNING request_id, api_key_id, id, created_at + ), + resolved AS ( + SELECT + input.input_idx, + input.request_id, + input.api_key_id, + COALESCE(inserted.id, existing.id) AS id, + COALESCE(inserted.created_at, existing.created_at) AS created_at, + (inserted.id IS NOT NULL) AS inserted + FROM input + LEFT JOIN inserted + ON inserted.request_id = input.request_id + AND inserted.api_key_id = input.api_key_id + LEFT JOIN usage_logs existing + ON existing.request_id = input.request_id + AND existing.api_key_id = input.api_key_id + ) + SELECT COALESCE( + json_agg( + json_build_object( + 'request_id', resolved.request_id, + 'api_key_id', resolved.api_key_id, + 'id', resolved.id, + 'created_at', resolved.created_at, + 'inserted', resolved.inserted + ) + ORDER BY resolved.input_idx + ), + '[]'::json + ) + FROM resolved + `) + return query.String(), args +} + +func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (string, []any) { + var query strings.Builder + _, _ = query.WriteString(` + WITH input ( + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + ) AS (VALUES `) + + args := make([]any, 0, len(preparedList)*53) + argPos := 1 + for idx, prepared := range preparedList { + if idx > 0 { + _, _ = query.WriteString(",") + } + _, _ = query.WriteString("(") + for i := 0; i < len(prepared.args); i++ { + if i > 0 { + _, _ = query.WriteString(",") + } + _, _ = query.WriteString("$") + _, _ = query.WriteString(strconv.Itoa(argPos)) + if i < len(usageLogInsertArgTypes) { + _, _ = query.WriteString("::") + _, _ = query.WriteString(usageLogInsertArgTypes[i]) + } + argPos++ + } + _, _ = query.WriteString(")") + args = append(args, prepared.args...) + } + + _, _ = query.WriteString(` + ) + INSERT INTO usage_logs ( + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + ) + SELECT + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + FROM input + ON CONFLICT (request_id, api_key_id) DO NOTHING + `) + + return query.String(), args +} + +func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared usageLogInsertPrepared) error { + _, err := sqlq.ExecContext(ctx, ` + INSERT INTO usage_logs ( + user_id, + api_key_id, + account_id, + request_id, + model, + requested_model, + upstream_model, + group_id, + subscription_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + cache_creation_5m_tokens, + cache_creation_1h_tokens, + image_output_tokens, + image_output_cost, + input_cost, + output_cost, + cache_creation_cost, + cache_read_cost, + total_cost, + actual_cost, + rate_multiplier, + account_rate_multiplier, + billing_type, + request_type, + stream, + openai_ws_mode, + duration_ms, + first_token_ms, + user_agent, + ip_address, + image_count, + image_size, + image_input_size, + image_output_size, + image_size_source, + image_size_breakdown, + video_count, + video_resolution, + video_duration_seconds, + service_tier, + reasoning_effort, + inbound_endpoint, + upstream_endpoint, + cache_ttl_overridden, + channel_id, + model_mapping_chain, + billing_tier, + billing_mode, + account_stats_cost, + created_at + ) VALUES ( + $1, $2, $3, $4, $5, $6, $7, + $8, $9, + $10, $11, $12, $13, + $14, $15, $16, $17, + $18, $19, $20, $21, $22, $23, + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53 + ) + ON CONFLICT (request_id, api_key_id) DO NOTHING + `, prepared.args...) + return err +} + +func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { + createdAt := log.CreatedAt + if createdAt.IsZero() { + createdAt = time.Now() + } + + requestID := strings.TrimSpace(log.RequestID) + log.RequestID = requestID + + rateMultiplier := log.RateMultiplier + log.SyncRequestTypeAndLegacyFields() + requestType := int16(log.RequestType) + + groupID := nullInt64(log.GroupID) + subscriptionID := nullInt64(log.SubscriptionID) + duration := nullInt(log.DurationMs) + firstToken := nullInt(log.FirstTokenMs) + userAgent := nullString(log.UserAgent) + ipAddress := nullString(log.IPAddress) + imageSize := nullString(log.ImageSize) + imageInputSize := nullString(log.ImageInputSize) + imageOutputSize := nullString(log.ImageOutputSize) + imageSizeSource := nullString(log.ImageSizeSource) + imageSizeBreakdown := nullStringIntMapJSON(log.ImageSizeBreakdown) + videoResolution := nullString(log.VideoResolution) + videoDurationSeconds := nullInt(log.VideoDurationSeconds) + serviceTier := nullString(log.ServiceTier) + reasoningEffort := nullString(log.ReasoningEffort) + inboundEndpoint := nullString(log.InboundEndpoint) + upstreamEndpoint := nullString(log.UpstreamEndpoint) + channelID := nullInt64(log.ChannelID) + modelMappingChain := nullString(log.ModelMappingChain) + billingTier := nullString(log.BillingTier) + billingMode := nullString(log.BillingMode) + requestedModel := strings.TrimSpace(log.RequestedModel) + if requestedModel == "" { + requestedModel = strings.TrimSpace(log.Model) + } + upstreamModel := nullString(log.UpstreamModel) + + var requestIDArg any + if requestID != "" { + requestIDArg = requestID + } + + return usageLogInsertPrepared{ + createdAt: createdAt, + requestID: requestID, + rateMultiplier: rateMultiplier, + requestType: requestType, + args: []any{ + log.UserID, + log.APIKeyID, + log.AccountID, + requestIDArg, + log.Model, + nullString(&requestedModel), + upstreamModel, + groupID, + subscriptionID, + log.InputTokens, + log.OutputTokens, + log.CacheCreationTokens, + log.CacheReadTokens, + log.CacheCreation5mTokens, + log.CacheCreation1hTokens, + log.ImageOutputTokens, + log.ImageOutputCost, + log.InputCost, + log.OutputCost, + log.CacheCreationCost, + log.CacheReadCost, + log.TotalCost, + log.ActualCost, + rateMultiplier, + log.AccountRateMultiplier, + log.BillingType, + requestType, + log.Stream, + log.OpenAIWSMode, + duration, + firstToken, + userAgent, + ipAddress, + log.ImageCount, + imageSize, + imageInputSize, + imageOutputSize, + imageSizeSource, + imageSizeBreakdown, + log.VideoCount, + videoResolution, + videoDurationSeconds, + serviceTier, + reasoningEffort, + inboundEndpoint, + upstreamEndpoint, + log.CacheTTLOverridden, + channelID, + modelMappingChain, + billingTier, + billingMode, + log.AccountStatsCost, // account_stats_cost + createdAt, + }, + } +} + +func usageLogBatchKey(requestID string, apiKeyID int64) string { + return requestID + "\x1f" + strconv.FormatInt(apiKeyID, 10) +} + +func sendUsageLogCreateResult(ch chan usageLogCreateResult, res usageLogCreateResult) { + if ch == nil { + return + } + select { + case ch <- res: + default: + } +} + +func (r *usageLogRepository) bestEffortRecentKey(requestID string, apiKeyID int64) (string, bool) { + requestID = strings.TrimSpace(requestID) + if requestID == "" || r == nil || r.bestEffortRecent == nil { + return "", false + } + return usageLogBatchKey(requestID, apiKeyID), true +} diff --git a/backend/internal/repository/usage_log_repo_query.go b/backend/internal/repository/usage_log_repo_query.go new file mode 100644 index 0000000000..c178429bab --- /dev/null +++ b/backend/internal/repository/usage_log_repo_query.go @@ -0,0 +1,726 @@ +package repository + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + "time" + + dbaccount "github.com/Wei-Shaw/sub2api/ent/account" + dbapikey "github.com/Wei-Shaw/sub2api/ent/apikey" + dbgroup "github.com/Wei-Shaw/sub2api/ent/group" + "github.com/Wei-Shaw/sub2api/ent/schema/mixins" + dbuser "github.com/Wei-Shaw/sub2api/ent/user" + dbusersub "github.com/Wei-Shaw/sub2api/ent/usersubscription" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/Wei-Shaw/sub2api/internal/service" +) + +const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" + +func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) { + query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1" + rows, err := r.sql.QueryContext(ctx, query, id) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + log = nil + } + }() + if !rows.Next() { + if err = rows.Err(); err != nil { + return nil, err + } + return nil, service.ErrUsageLogNotFound + } + log, err = scanUsageLog(rows) + if err != nil { + return nil, err + } + if err = rows.Err(); err != nil { + return nil, err + } + return log, nil +} + +func (r *usageLogRepository) ListByUser(ctx context.Context, userID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { + return r.listUsageLogsWithPagination(ctx, "WHERE user_id = $1", []any{userID}, params) +} + +func (r *usageLogRepository) ListByAPIKey(ctx context.Context, apiKeyID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { + return r.listUsageLogsWithPagination(ctx, "WHERE api_key_id = $1", []any{apiKeyID}, params) +} + +func (r *usageLogRepository) ListByAccount(ctx context.Context, accountID int64, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { + return r.listUsageLogsWithPagination(ctx, "WHERE account_id = $1", []any{accountID}, params) +} + +func (r *usageLogRepository) ListByUserAndTimeRange(ctx context.Context, userID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { + query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000" + logs, err := r.queryUsageLogs(ctx, query, userID, startTime, endTime) + return logs, nil, err +} + +func (r *usageLogRepository) ListByAPIKeyAndTimeRange(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { + query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE api_key_id = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000" + logs, err := r.queryUsageLogs(ctx, query, apiKeyID, startTime, endTime) + return logs, nil, err +} + +func (r *usageLogRepository) ListByAccountAndTimeRange(ctx context.Context, accountID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { + query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE account_id = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000" + logs, err := r.queryUsageLogs(ctx, query, accountID, startTime, endTime) + return logs, nil, err +} + +func (r *usageLogRepository) ListByModelAndTimeRange(ctx context.Context, modelName string, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) { + query := fmt.Sprintf("SELECT %s FROM usage_logs WHERE %s = $1 AND created_at >= $2 AND created_at < $3 ORDER BY id DESC LIMIT 10000", usageLogSelectColumns, rawUsageLogModelColumn) + logs, err := r.queryUsageLogs(ctx, query, modelName, startTime, endTime) + return logs, nil, err +} + +func (r *usageLogRepository) Delete(ctx context.Context, id int64) error { + _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_logs WHERE id = $1", id) + return err +} + +// UsageLogFilters represents filters for usage log queries +type UsageLogFilters = usagestats.UsageLogFilters + +// ListWithFilters lists usage logs with optional filters (for admin) +func (r *usageLogRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters UsageLogFilters) ([]service.UsageLog, *pagination.PaginationResult, error) { + conditions := make([]string, 0, 9) + args := make([]any, 0, 9) + + if filters.UserID > 0 { + conditions = append(conditions, fmt.Sprintf("user_id = $%d", len(args)+1)) + args = append(args, filters.UserID) + } + if filters.APIKeyID > 0 { + conditions = append(conditions, fmt.Sprintf("api_key_id = $%d", len(args)+1)) + args = append(args, filters.APIKeyID) + } + if filters.AccountID > 0 { + conditions = append(conditions, fmt.Sprintf("account_id = $%d", len(args)+1)) + args = append(args, filters.AccountID) + } + if filters.GroupID > 0 { + conditions = append(conditions, fmt.Sprintf("group_id = $%d", len(args)+1)) + args = append(args, filters.GroupID) + } + conditions, args = appendUsageLogModelWhereCondition(conditions, args, filters.Model, filters.ModelFilterSource) + conditions, args = appendRequestTypeOrStreamWhereCondition(conditions, args, filters.RequestType, filters.Stream) + if filters.BillingType != nil { + conditions = append(conditions, fmt.Sprintf("billing_type = $%d", len(args)+1)) + args = append(args, int16(*filters.BillingType)) + } + conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode) + if filters.StartTime != nil { + conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1)) + args = append(args, *filters.StartTime) + } + if filters.EndTime != nil { + conditions = append(conditions, fmt.Sprintf("created_at < $%d", len(args)+1)) + args = append(args, *filters.EndTime) + } + + whereClause := buildWhere(conditions) + var ( + logs []service.UsageLog + page *pagination.PaginationResult + err error + ) + if shouldUseFastUsageLogTotal(filters) { + logs, page, err = r.listUsageLogsWithFastPagination(ctx, whereClause, args, params) + } else { + logs, page, err = r.listUsageLogsWithPagination(ctx, whereClause, args, params) + } + if err != nil { + return nil, nil, err + } + + if err := r.hydrateUsageLogAssociations(ctx, logs); err != nil { + return nil, nil, err + } + return logs, page, nil +} + +func shouldUseFastUsageLogTotal(filters UsageLogFilters) bool { + if filters.ExactTotal { + return false + } + // 强选择过滤下记录集通常较小,保留精确总数。 + return filters.UserID == 0 && filters.APIKeyID == 0 && filters.AccountID == 0 +} + +func (r *usageLogRepository) listUsageLogsWithPagination(ctx context.Context, whereClause string, args []any, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { + countQuery := "SELECT COUNT(*) FROM usage_logs " + whereClause + var total int64 + if err := scanSingleRow(ctx, r.sql, countQuery, args, &total); err != nil { + return nil, nil, err + } + + limitPos := len(args) + 1 + offsetPos := len(args) + 2 + listArgs := append(append([]any{}, args...), params.Limit(), params.Offset()) + query := fmt.Sprintf("SELECT %s FROM usage_logs %s ORDER BY %s LIMIT $%d OFFSET $%d", usageLogSelectColumns, whereClause, usageLogOrderBy(params), limitPos, offsetPos) + logs, err := r.queryUsageLogs(ctx, query, listArgs...) + if err != nil { + return nil, nil, err + } + return logs, paginationResultFromTotal(total, params), nil +} + +func (r *usageLogRepository) listUsageLogsWithFastPagination(ctx context.Context, whereClause string, args []any, params pagination.PaginationParams) ([]service.UsageLog, *pagination.PaginationResult, error) { + limit := params.Limit() + offset := params.Offset() + + limitPos := len(args) + 1 + offsetPos := len(args) + 2 + listArgs := append(append([]any{}, args...), limit+1, offset) + query := fmt.Sprintf("SELECT %s FROM usage_logs %s ORDER BY %s LIMIT $%d OFFSET $%d", usageLogSelectColumns, whereClause, usageLogOrderBy(params), limitPos, offsetPos) + + logs, err := r.queryUsageLogs(ctx, query, listArgs...) + if err != nil { + return nil, nil, err + } + + hasMore := false + if len(logs) > limit { + hasMore = true + logs = logs[:limit] + } + + total := int64(offset) + int64(len(logs)) + if hasMore { + // 只保证“还有下一页”,避免对超大表做全量 COUNT(*)。 + total = int64(offset) + int64(limit) + 1 + } + + return logs, paginationResultFromTotal(total, params), nil +} + +func usageLogOrderBy(params pagination.PaginationParams) string { + sortBy := strings.ToLower(strings.TrimSpace(params.SortBy)) + sortOrder := strings.ToUpper(params.NormalizedSortOrder(pagination.SortOrderDesc)) + + var column string + switch sortBy { + case "model": + column = "COALESCE(NULLIF(TRIM(requested_model), ''), model)" + case "created_at": + column = "created_at" + default: + column = "id" + } + + if column == "id" { + return fmt.Sprintf("id %s", sortOrder) + } + return fmt.Sprintf("%s %s, id %s", column, sortOrder, sortOrder) +} + +func (r *usageLogRepository) queryUsageLogs(ctx context.Context, query string, args ...any) (logs []service.UsageLog, err error) { + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + logs = nil + } + }() + + logs = make([]service.UsageLog, 0) + for rows.Next() { + var log *service.UsageLog + log, err = scanUsageLog(rows) + if err != nil { + return nil, err + } + logs = append(logs, *log) + } + if err = rows.Err(); err != nil { + return nil, err + } + return logs, nil +} + +func (r *usageLogRepository) hydrateUsageLogAssociations(ctx context.Context, logs []service.UsageLog) error { + // 关联数据使用 Ent 批量加载,避免把复杂 SQL 继续膨胀。 + if len(logs) == 0 { + return nil + } + + ids := collectUsageLogIDs(logs) + users, err := r.loadUsers(ctx, ids.userIDs) + if err != nil { + return err + } + apiKeys, err := r.loadAPIKeys(ctx, ids.apiKeyIDs) + if err != nil { + return err + } + accounts, err := r.loadAccounts(ctx, ids.accountIDs) + if err != nil { + return err + } + groups, err := r.loadGroups(ctx, ids.groupIDs) + if err != nil { + return err + } + subs, err := r.loadSubscriptions(ctx, ids.subscriptionIDs) + if err != nil { + return err + } + + for i := range logs { + if user, ok := users[logs[i].UserID]; ok { + logs[i].User = user + } + if key, ok := apiKeys[logs[i].APIKeyID]; ok { + logs[i].APIKey = key + } + if acc, ok := accounts[logs[i].AccountID]; ok { + logs[i].Account = acc + } + if logs[i].GroupID != nil { + if group, ok := groups[*logs[i].GroupID]; ok { + logs[i].Group = group + } + } + if logs[i].SubscriptionID != nil { + if sub, ok := subs[*logs[i].SubscriptionID]; ok { + logs[i].Subscription = sub + } + } + } + return nil +} + +type usageLogIDs struct { + userIDs []int64 + apiKeyIDs []int64 + accountIDs []int64 + groupIDs []int64 + subscriptionIDs []int64 +} + +func collectUsageLogIDs(logs []service.UsageLog) usageLogIDs { + idSet := func() map[int64]struct{} { return make(map[int64]struct{}) } + + userIDs := idSet() + apiKeyIDs := idSet() + accountIDs := idSet() + groupIDs := idSet() + subscriptionIDs := idSet() + + for i := range logs { + userIDs[logs[i].UserID] = struct{}{} + apiKeyIDs[logs[i].APIKeyID] = struct{}{} + accountIDs[logs[i].AccountID] = struct{}{} + if logs[i].GroupID != nil { + groupIDs[*logs[i].GroupID] = struct{}{} + } + if logs[i].SubscriptionID != nil { + subscriptionIDs[*logs[i].SubscriptionID] = struct{}{} + } + } + + return usageLogIDs{ + userIDs: setToSlice(userIDs), + apiKeyIDs: setToSlice(apiKeyIDs), + accountIDs: setToSlice(accountIDs), + groupIDs: setToSlice(groupIDs), + subscriptionIDs: setToSlice(subscriptionIDs), + } +} + +func (r *usageLogRepository) loadUsers(ctx context.Context, ids []int64) (map[int64]*service.User, error) { + out := make(map[int64]*service.User) + if len(ids) == 0 { + return out, nil + } + // 无条件穿透软删除:ids 来自调用方已按 user_id 筛选的日志行;普通用户路径强制 UserID=本人(本人必为活跃用户),不会借此解析他人已删身份;仅 admin 路径可借此显示已删用户。 + models, err := r.client.User.Query().Where(dbuser.IDIn(ids...)).All(mixins.SkipSoftDelete(ctx)) + if err != nil { + return nil, err + } + for _, m := range models { + out[m.ID] = userEntityToService(m) + } + return out, nil +} + +func (r *usageLogRepository) loadAPIKeys(ctx context.Context, ids []int64) (map[int64]*service.APIKey, error) { + out := make(map[int64]*service.APIKey) + if len(ids) == 0 { + return out, nil + } + models, err := r.client.APIKey.Query().Where(dbapikey.IDIn(ids...)).All(ctx) + if err != nil { + return nil, err + } + for _, m := range models { + out[m.ID] = apiKeyEntityToService(m) + } + return out, nil +} + +func (r *usageLogRepository) loadAccounts(ctx context.Context, ids []int64) (map[int64]*service.Account, error) { + out := make(map[int64]*service.Account) + if len(ids) == 0 { + return out, nil + } + models, err := r.client.Account.Query().Where(dbaccount.IDIn(ids...)).All(ctx) + if err != nil { + return nil, err + } + for _, m := range models { + out[m.ID] = accountEntityToService(m) + } + return out, nil +} + +func (r *usageLogRepository) loadGroups(ctx context.Context, ids []int64) (map[int64]*service.Group, error) { + out := make(map[int64]*service.Group) + if len(ids) == 0 { + return out, nil + } + models, err := r.client.Group.Query().Where(dbgroup.IDIn(ids...)).All(ctx) + if err != nil { + return nil, err + } + for _, m := range models { + out[m.ID] = groupEntityToService(m) + } + return out, nil +} + +func (r *usageLogRepository) loadSubscriptions(ctx context.Context, ids []int64) (map[int64]*service.UserSubscription, error) { + out := make(map[int64]*service.UserSubscription) + if len(ids) == 0 { + return out, nil + } + models, err := r.client.UserSubscription.Query().Where(dbusersub.IDIn(ids...)).All(ctx) + if err != nil { + return nil, err + } + for _, m := range models { + out[m.ID] = userSubscriptionEntityToService(m) + } + return out, nil +} + +func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, error) { + var ( + id int64 + userID int64 + apiKeyID int64 + accountID int64 + requestID sql.NullString + model string + requestedModel sql.NullString + upstreamModel sql.NullString + groupID sql.NullInt64 + subscriptionID sql.NullInt64 + inputTokens int + outputTokens int + cacheCreationTokens int + cacheReadTokens int + cacheCreation5m int + cacheCreation1h int + imageOutputTokens int + imageOutputCost float64 + inputCost float64 + outputCost float64 + cacheCreationCost float64 + cacheReadCost float64 + totalCost float64 + actualCost float64 + rateMultiplier float64 + accountRateMultiplier sql.NullFloat64 + billingType int16 + requestTypeRaw int16 + stream bool + openaiWSMode bool + durationMs sql.NullInt64 + firstTokenMs sql.NullInt64 + userAgent sql.NullString + ipAddress sql.NullString + imageCount int + imageSize sql.NullString + imageInputSize sql.NullString + imageOutputSize sql.NullString + imageSizeSource sql.NullString + imageSizeBreakdown sql.NullString + videoCount int + videoResolution sql.NullString + videoDurationSeconds sql.NullInt64 + serviceTier sql.NullString + reasoningEffort sql.NullString + inboundEndpoint sql.NullString + upstreamEndpoint sql.NullString + cacheTTLOverridden bool + channelID sql.NullInt64 + modelMappingChain sql.NullString + billingTier sql.NullString + billingMode sql.NullString + accountStatsCost sql.NullFloat64 + createdAt time.Time + ) + + if err := scanner.Scan( + &id, + &userID, + &apiKeyID, + &accountID, + &requestID, + &model, + &requestedModel, + &upstreamModel, + &groupID, + &subscriptionID, + &inputTokens, + &outputTokens, + &cacheCreationTokens, + &cacheReadTokens, + &cacheCreation5m, + &cacheCreation1h, + &imageOutputTokens, + &imageOutputCost, + &inputCost, + &outputCost, + &cacheCreationCost, + &cacheReadCost, + &totalCost, + &actualCost, + &rateMultiplier, + &accountRateMultiplier, + &billingType, + &requestTypeRaw, + &stream, + &openaiWSMode, + &durationMs, + &firstTokenMs, + &userAgent, + &ipAddress, + &imageCount, + &imageSize, + &imageInputSize, + &imageOutputSize, + &imageSizeSource, + &imageSizeBreakdown, + &videoCount, + &videoResolution, + &videoDurationSeconds, + &serviceTier, + &reasoningEffort, + &inboundEndpoint, + &upstreamEndpoint, + &cacheTTLOverridden, + &channelID, + &modelMappingChain, + &billingTier, + &billingMode, + &accountStatsCost, + &createdAt, + ); err != nil { + return nil, err + } + + log := &service.UsageLog{ + ID: id, + UserID: userID, + APIKeyID: apiKeyID, + AccountID: accountID, + Model: model, + RequestedModel: coalesceTrimmedString(requestedModel, model), + InputTokens: inputTokens, + OutputTokens: outputTokens, + CacheCreationTokens: cacheCreationTokens, + CacheReadTokens: cacheReadTokens, + CacheCreation5mTokens: cacheCreation5m, + CacheCreation1hTokens: cacheCreation1h, + ImageOutputTokens: imageOutputTokens, + ImageOutputCost: imageOutputCost, + InputCost: inputCost, + OutputCost: outputCost, + CacheCreationCost: cacheCreationCost, + CacheReadCost: cacheReadCost, + TotalCost: totalCost, + ActualCost: actualCost, + RateMultiplier: rateMultiplier, + AccountRateMultiplier: nullFloat64Ptr(accountRateMultiplier), + BillingType: int8(billingType), + RequestType: service.RequestTypeFromInt16(requestTypeRaw), + ImageCount: imageCount, + VideoCount: videoCount, + CacheTTLOverridden: cacheTTLOverridden, + CreatedAt: createdAt, + } + // 先回填 legacy 字段,再基于 legacy + request_type 计算最终请求类型,保证历史数据兼容。 + log.Stream = stream + log.OpenAIWSMode = openaiWSMode + log.RequestType = log.EffectiveRequestType() + log.Stream, log.OpenAIWSMode = service.ApplyLegacyRequestFields(log.RequestType, stream, openaiWSMode) + + if requestID.Valid { + log.RequestID = requestID.String + } + if groupID.Valid { + value := groupID.Int64 + log.GroupID = &value + } + if subscriptionID.Valid { + value := subscriptionID.Int64 + log.SubscriptionID = &value + } + if durationMs.Valid { + value := int(durationMs.Int64) + log.DurationMs = &value + } + if firstTokenMs.Valid { + value := int(firstTokenMs.Int64) + log.FirstTokenMs = &value + } + if userAgent.Valid { + log.UserAgent = &userAgent.String + } + if ipAddress.Valid { + log.IPAddress = &ipAddress.String + } + if imageSize.Valid { + log.ImageSize = &imageSize.String + } + if imageInputSize.Valid { + log.ImageInputSize = &imageInputSize.String + } + if imageOutputSize.Valid { + log.ImageOutputSize = &imageOutputSize.String + } + if imageSizeSource.Valid { + log.ImageSizeSource = &imageSizeSource.String + } + log.ImageSizeBreakdown = stringIntMapFromNullJSON(imageSizeBreakdown) + if videoResolution.Valid { + log.VideoResolution = &videoResolution.String + } + if videoDurationSeconds.Valid { + value := int(videoDurationSeconds.Int64) + log.VideoDurationSeconds = &value + } + if serviceTier.Valid { + log.ServiceTier = &serviceTier.String + } + if reasoningEffort.Valid { + log.ReasoningEffort = &reasoningEffort.String + } + if inboundEndpoint.Valid { + log.InboundEndpoint = &inboundEndpoint.String + } + if upstreamEndpoint.Valid { + log.UpstreamEndpoint = &upstreamEndpoint.String + } + if upstreamModel.Valid { + log.UpstreamModel = &upstreamModel.String + } + if channelID.Valid { + value := channelID.Int64 + log.ChannelID = &value + } + if modelMappingChain.Valid { + log.ModelMappingChain = &modelMappingChain.String + } + if billingTier.Valid { + log.BillingTier = &billingTier.String + } + if billingMode.Valid { + log.BillingMode = &billingMode.String + } + if accountStatsCost.Valid { + log.AccountStatsCost = &accountStatsCost.Float64 + } + + return log, nil +} + +func nullInt64(v *int64) sql.NullInt64 { + if v == nil { + return sql.NullInt64{} + } + return sql.NullInt64{Int64: *v, Valid: true} +} + +func nullInt(v *int) sql.NullInt64 { + if v == nil { + return sql.NullInt64{} + } + return sql.NullInt64{Int64: int64(*v), Valid: true} +} + +func nullFloat64Ptr(v sql.NullFloat64) *float64 { + if !v.Valid { + return nil + } + out := v.Float64 + return &out +} + +func nullString(v *string) sql.NullString { + if v == nil || *v == "" { + return sql.NullString{} + } + return sql.NullString{String: *v, Valid: true} +} + +func nullStringIntMapJSON(v map[string]int) any { + if len(v) == 0 { + return nil + } + payload, err := json.Marshal(v) + if err != nil { + return nil + } + return string(payload) +} + +func stringIntMapFromNullJSON(v sql.NullString) map[string]int { + if !v.Valid || strings.TrimSpace(v.String) == "" { + return nil + } + var out map[string]int + if err := json.Unmarshal([]byte(v.String), &out); err != nil { + return nil + } + if len(out) == 0 { + return nil + } + return out +} + +func coalesceTrimmedString(v sql.NullString, fallback string) string { + if v.Valid && strings.TrimSpace(v.String) != "" { + return v.String + } + return fallback +} + +func setToSlice(set map[int64]struct{}) []int64 { + out := make([]int64, 0, len(set)) + for id := range set { + out = append(out, id) + } + return out +} diff --git a/backend/internal/repository/usage_log_repo_request_type_test.go b/backend/internal/repository/usage_log_repo_request_type_test.go index e4d4e9a4fa..c32ad2b63f 100644 --- a/backend/internal/repository/usage_log_repo_request_type_test.go +++ b/backend/internal/repository/usage_log_repo_request_type_test.go @@ -80,6 +80,9 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) { sqlmock.AnyArg(), // image_output_size sqlmock.AnyArg(), // image_size_source sqlmock.AnyArg(), // image_size_breakdown + sqlmock.AnyArg(), // video_count + sqlmock.AnyArg(), // video_resolution + sqlmock.AnyArg(), // video_duration_seconds sqlmock.AnyArg(), // service_tier sqlmock.AnyArg(), // reasoning_effort sqlmock.AnyArg(), // inbound_endpoint @@ -163,6 +166,9 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) { sqlmock.AnyArg(), // image_output_size sqlmock.AnyArg(), // image_size_source sqlmock.AnyArg(), // image_size_breakdown + sqlmock.AnyArg(), // video_count + sqlmock.AnyArg(), // video_resolution + sqlmock.AnyArg(), // video_duration_seconds serviceTier, sqlmock.AnyArg(), sqlmock.AnyArg(), @@ -281,9 +287,14 @@ func TestAppendUsageLogBillingModeWhereCondition(t *testing.T) { wantCondition string }{ { - name: "image includes legacy image rows", + name: "image includes explicit image and legacy image rows", billingMode: string(service.BillingModeImage), - wantCondition: "(billing_mode = $1 OR COALESCE(image_count, 0) > 0)", + wantCondition: "(billing_mode = $1 OR ((billing_mode IS NULL OR billing_mode = '') AND COALESCE(image_count, 0) > 0))", + }, + { + name: "video remains exact", + billingMode: string(service.BillingModeVideo), + wantCondition: "billing_mode = $1", }, { name: "token includes legacy non-image rows", @@ -309,7 +320,7 @@ func TestAppendUsageLogBillingModeWhereCondition(t *testing.T) { func TestAppendUsageLogBillingModeWhereConditionWithAlias(t *testing.T) { conditions, args := appendUsageLogBillingModeWhereConditionWithAlias(nil, nil, string(service.BillingModeImage), "ul") - require.Equal(t, []string{"(ul.billing_mode = $1 OR COALESCE(ul.image_count, 0) > 0)"}, conditions) + require.Equal(t, []string{"(ul.billing_mode = $1 OR ((ul.billing_mode IS NULL OR ul.billing_mode = '') AND COALESCE(ul.image_count, 0) > 0))"}, conditions) require.Equal(t, []any{string(service.BillingModeImage)}, args) } @@ -794,6 +805,9 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{Valid: true, String: "3840x2160"}, sql.NullString{Valid: true, String: "output"}, sql.NullString{Valid: true, String: `{"4K":2}`}, + 0, // video_count + sql.NullString{}, // video_resolution + sql.NullInt64{}, // video_duration_seconds sql.NullString{}, sql.NullString{}, sql.NullString{}, @@ -862,6 +876,9 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, // image_output_size sql.NullString{}, // image_size_source sql.NullString{}, // image_size_breakdown + 0, // video_count + sql.NullString{}, // video_resolution + sql.NullInt64{}, // video_duration_seconds sql.NullString{Valid: true, String: "priority"}, sql.NullString{}, sql.NullString{}, @@ -914,6 +931,9 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, // image_output_size sql.NullString{}, // image_size_source sql.NullString{}, // image_size_breakdown + 0, // video_count + sql.NullString{}, // video_resolution + sql.NullInt64{}, // video_duration_seconds sql.NullString{Valid: true, String: "flex"}, sql.NullString{}, sql.NullString{}, @@ -966,6 +986,9 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, // image_output_size sql.NullString{}, // image_size_source sql.NullString{}, // image_size_breakdown + 0, // video_count + sql.NullString{}, // video_resolution + sql.NullInt64{}, // video_duration_seconds sql.NullString{Valid: true, String: "priority"}, sql.NullString{}, sql.NullString{}, diff --git a/backend/internal/repository/usage_log_repo_stats.go b/backend/internal/repository/usage_log_repo_stats.go new file mode 100644 index 0000000000..aeccd49b30 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_stats.go @@ -0,0 +1,1145 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + "fmt" + "os" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/lib/pq" + "golang.org/x/sync/errgroup" +) + +// GetUserStatsAggregated returns aggregated usage statistics for a user using database-level aggregation +func (r *usageLogRepository) GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) { + query := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms + FROM usage_logs + WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 + ` + + var stats usagestats.UsageStats + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{userID, startTime, endTime}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens + return &stats, nil +} + +// GetAPIKeyStatsAggregated returns aggregated usage statistics for an API key using database-level aggregation +func (r *usageLogRepository) GetAPIKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) { + query := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms + FROM usage_logs + WHERE api_key_id = $1 AND created_at >= $2 AND created_at < $3 + ` + + var stats usagestats.UsageStats + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{apiKeyID, startTime, endTime}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens + return &stats, nil +} + +// GetAccountStatsAggregated 使用 SQL 聚合统计账号使用数据 +// +// 性能优化说明: +// 原实现先查询所有日志记录,再在应用层循环计算统计值: +// 1. 需要传输大量数据到应用层 +// 2. 应用层循环计算增加 CPU 和内存开销 +// +// 新实现使用 SQL 聚合函数: +// 1. 在数据库层完成 COUNT/SUM/AVG 计算 +// 2. 只返回单行聚合结果,大幅减少数据传输量 +// 3. 利用数据库索引优化聚合查询性能 +func (r *usageLogRepository) GetAccountStatsAggregated(ctx context.Context, accountID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) { + query := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms + FROM usage_logs + WHERE account_id = $1 AND created_at >= $2 AND created_at < $3 + ` + + var stats usagestats.UsageStats + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{accountID, startTime, endTime}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens + return &stats, nil +} + +// GetModelStatsAggregated 使用 SQL 聚合统计模型使用数据 +// 性能优化:数据库层聚合计算,避免应用层循环统计 +func (r *usageLogRepository) GetModelStatsAggregated(ctx context.Context, modelName string, startTime, endTime time.Time) (*usagestats.UsageStats, error) { + query := fmt.Sprintf(` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms + FROM usage_logs + WHERE %s = $1 AND created_at >= $2 AND created_at < $3 + `, rawUsageLogModelColumn) + + var stats usagestats.UsageStats + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{modelName, startTime, endTime}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens + return &stats, nil +} + +// GetDailyStatsAggregated 使用 SQL 聚合统计用户的每日使用数据 +// 性能优化:使用 GROUP BY 在数据库层按日期分组聚合,避免应用层循环分组统计 +func (r *usageLogRepository) GetDailyStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (result []map[string]any, err error) { + tzName := resolveUsageStatsTimezone() + query := ` + SELECT + -- 使用应用时区分组,避免数据库会话时区导致日边界偏移。 + TO_CHAR(created_at AT TIME ZONE $4, 'YYYY-MM-DD') as date, + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms + FROM usage_logs + WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 + GROUP BY 1 + ORDER BY 1 + ` + + rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime, tzName) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + result = nil + } + }() + + result = make([]map[string]any, 0) + for rows.Next() { + var ( + date string + totalRequests int64 + totalInputTokens int64 + totalOutputTokens int64 + totalCacheTokens int64 + totalCost float64 + totalActualCost float64 + avgDurationMs float64 + ) + if err = rows.Scan( + &date, + &totalRequests, + &totalInputTokens, + &totalOutputTokens, + &totalCacheTokens, + &totalCost, + &totalActualCost, + &avgDurationMs, + ); err != nil { + return nil, err + } + result = append(result, map[string]any{ + "date": date, + "total_requests": totalRequests, + "total_input_tokens": totalInputTokens, + "total_output_tokens": totalOutputTokens, + "total_cache_tokens": totalCacheTokens, + "total_tokens": totalInputTokens + totalOutputTokens + totalCacheTokens, + "total_cost": totalCost, + "total_actual_cost": totalActualCost, + "average_duration_ms": avgDurationMs, + }) + } + + if err = rows.Err(); err != nil { + return nil, err + } + + return result, nil +} + +// resolveUsageStatsTimezone 获取用于 SQL 分组的时区名称。 +// 优先使用应用初始化的时区,其次尝试读取 TZ 环境变量,最后回落为 UTC。 +func resolveUsageStatsTimezone() string { + tzName := timezone.Name() + if tzName != "" && tzName != "Local" { + return tzName + } + if envTZ := strings.TrimSpace(os.Getenv("TZ")); envTZ != "" { + return envTZ + } + return "UTC" +} + +// GetAccountTodayStats 获取账号今日统计 +func (r *usageLogRepository) GetAccountTodayStats(ctx context.Context, accountID int64) (*usagestats.AccountStats, error) { + today := timezone.Today() + + query := ` + SELECT + COUNT(*) as requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, + COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost, + COALESCE(SUM(total_cost), 0) as standard_cost, + COALESCE(SUM(actual_cost), 0) as user_cost + FROM usage_logs + WHERE account_id = $1 AND created_at >= $2 + ` + + stats := &usagestats.AccountStats{} + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{accountID, today}, + &stats.Requests, + &stats.Tokens, + &stats.Cost, + &stats.StandardCost, + &stats.UserCost, + ); err != nil { + return nil, err + } + return stats, nil +} + +// GetAccountWindowStats 获取账号时间窗口内的统计 +func (r *usageLogRepository) GetAccountWindowStats(ctx context.Context, accountID int64, startTime time.Time) (*usagestats.AccountStats, error) { + query := ` + SELECT + COUNT(*) as requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, + COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost, + COALESCE(SUM(total_cost), 0) as standard_cost, + COALESCE(SUM(actual_cost), 0) as user_cost + FROM usage_logs + WHERE account_id = $1 AND created_at >= $2 + ` + + stats := &usagestats.AccountStats{} + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{accountID, startTime}, + &stats.Requests, + &stats.Tokens, + &stats.Cost, + &stats.StandardCost, + &stats.UserCost, + ); err != nil { + return nil, err + } + return stats, nil +} + +// GetAccountWindowStatsBatch 批量获取同一窗口起点下多个账号的统计数据。 +// 返回 map[accountID]*AccountStats,未命中的账号会返回零值统计,便于上层直接复用。 +func (r *usageLogRepository) GetAccountWindowStatsBatch(ctx context.Context, accountIDs []int64, startTime time.Time) (map[int64]*usagestats.AccountStats, error) { + result := make(map[int64]*usagestats.AccountStats, len(accountIDs)) + if len(accountIDs) == 0 { + return result, nil + } + + query := ` + SELECT + account_id, + COUNT(*) as requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, + COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost, + COALESCE(SUM(total_cost), 0) as standard_cost, + COALESCE(SUM(actual_cost), 0) as user_cost + FROM usage_logs + WHERE account_id = ANY($1) AND created_at >= $2 + GROUP BY account_id + ` + rows, err := r.sql.QueryContext(ctx, query, pq.Array(accountIDs), startTime) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + for rows.Next() { + var accountID int64 + stats := &usagestats.AccountStats{} + if err := rows.Scan( + &accountID, + &stats.Requests, + &stats.Tokens, + &stats.Cost, + &stats.StandardCost, + &stats.UserCost, + ); err != nil { + return nil, err + } + result[accountID] = stats + } + if err := rows.Err(); err != nil { + return nil, err + } + + for _, accountID := range accountIDs { + if _, ok := result[accountID]; !ok { + result[accountID] = &usagestats.AccountStats{} + } + } + return result, nil +} + +// GetGeminiUsageTotalsBatch 批量聚合 Gemini 账号在窗口内的 Pro/Flash 请求与用量。 +// 模型分类规则与 service.geminiModelClassFromName 一致:model 包含 flash/lite 视为 flash,其余视为 pro。 +func (r *usageLogRepository) GetGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, startTime, endTime time.Time) (map[int64]service.GeminiUsageTotals, error) { + result := make(map[int64]service.GeminiUsageTotals, len(accountIDs)) + if len(accountIDs) == 0 { + return result, nil + } + + query := ` + SELECT + account_id, + COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 1 ELSE 0 END), 0) AS flash_requests, + COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE 1 END), 0) AS pro_requests, + COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) ELSE 0 END), 0) AS flash_tokens, + COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) END), 0) AS pro_tokens, + COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN actual_cost ELSE 0 END), 0) AS flash_cost, + COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE actual_cost END), 0) AS pro_cost + FROM usage_logs + WHERE account_id = ANY($1) AND created_at >= $2 AND created_at < $3 + GROUP BY account_id + ` + rows, err := r.sql.QueryContext(ctx, query, pq.Array(accountIDs), startTime, endTime) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + for rows.Next() { + var accountID int64 + var totals service.GeminiUsageTotals + if err := rows.Scan( + &accountID, + &totals.FlashRequests, + &totals.ProRequests, + &totals.FlashTokens, + &totals.ProTokens, + &totals.FlashCost, + &totals.ProCost, + ); err != nil { + return nil, err + } + result[accountID] = totals + } + if err := rows.Err(); err != nil { + return nil, err + } + + for _, accountID := range accountIDs { + if _, ok := result[accountID]; !ok { + result[accountID] = service.GeminiUsageTotals{} + } + } + return result, nil +} + +// UsageStats represents usage statistics +type UsageStats = usagestats.UsageStats + +// BatchUserUsageStats represents usage stats for a single user +type BatchUserUsageStats = usagestats.BatchUserUsageStats + +// PlatformUsage represents per-platform usage breakdown +type PlatformUsage = usagestats.PlatformUsage + +func normalizePositiveInt64IDs(ids []int64) []int64 { + if len(ids) == 0 { + return nil + } + seen := make(map[int64]struct{}, len(ids)) + out := make([]int64, 0, len(ids)) + for _, id := range ids { + if id <= 0 { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + out = append(out, id) + } + return out +} + +// GetBatchUserUsageStats gets today and total actual_cost for multiple users within a time range. +// If startTime is zero, defaults to 30 days ago. +func (r *usageLogRepository) GetBatchUserUsageStats(ctx context.Context, userIDs []int64, startTime, endTime time.Time) (map[int64]*BatchUserUsageStats, error) { + result := make(map[int64]*BatchUserUsageStats) + normalizedUserIDs := normalizePositiveInt64IDs(userIDs) + if len(normalizedUserIDs) == 0 { + return result, nil + } + + // 默认最近 30 天 + if startTime.IsZero() { + startTime = time.Now().AddDate(0, 0, -30) + } + if endTime.IsZero() { + endTime = time.Now() + } + + for _, id := range normalizedUserIDs { + result[id] = &BatchUserUsageStats{UserID: id} + } + + // GROUP BY (user_id, effective_platform) 一次查询同时得到总值与按平台拆分。 + // 应用层把同一 user_id 的多行累加为总值,并把非空 platform 行收集到 ByPlatform。 + query := ` + SELECT + ul.user_id, + ` + usageLogEffectivePlatformExpr + ` as platform, + COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $2 AND ul.created_at < $3), 0) as total_cost, + COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $4), 0) as today_cost + FROM usage_logs ul + LEFT JOIN groups g ON g.id = ul.group_id + LEFT JOIN accounts a ON a.id = ul.account_id + WHERE ul.user_id = ANY($1) + AND ul.created_at >= LEAST($2, $4) + AND ` + usageLogSuccessFilterUL + ` + GROUP BY ul.user_id, ` + usageLogEffectivePlatformExpr + ` + ` + today := timezone.Today() + rows, err := r.sql.QueryContext(ctx, query, pq.Array(normalizedUserIDs), startTime, endTime, today) + if err != nil { + return nil, err + } + for rows.Next() { + var userID int64 + var platform sql.NullString + var total float64 + var todayTotal float64 + if err := rows.Scan(&userID, &platform, &total, &todayTotal); err != nil { + _ = rows.Close() + return nil, err + } + stats, ok := result[userID] + if !ok { + continue + } + stats.TotalActualCost += total + stats.TodayActualCost += todayTotal + if platform.Valid && platform.String != "" { + stats.ByPlatform = append(stats.ByPlatform, PlatformUsage{ + Platform: platform.String, + TotalActualCost: total, + TodayActualCost: todayTotal, + }) + } + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + + return result, nil +} + +// BatchAPIKeyUsageStats represents usage stats for a single API key +type BatchAPIKeyUsageStats = usagestats.BatchAPIKeyUsageStats + +// GetBatchAPIKeyUsageStats gets today and total actual_cost for multiple API keys within a time range. +// If startTime is zero, defaults to 30 days ago. +func (r *usageLogRepository) GetBatchAPIKeyUsageStats(ctx context.Context, apiKeyIDs []int64, startTime, endTime time.Time) (map[int64]*BatchAPIKeyUsageStats, error) { + result := make(map[int64]*BatchAPIKeyUsageStats) + normalizedAPIKeyIDs := normalizePositiveInt64IDs(apiKeyIDs) + if len(normalizedAPIKeyIDs) == 0 { + return result, nil + } + + // 默认最近 30 天 + if startTime.IsZero() { + startTime = time.Now().AddDate(0, 0, -30) + } + if endTime.IsZero() { + endTime = time.Now() + } + + for _, id := range normalizedAPIKeyIDs { + result[id] = &BatchAPIKeyUsageStats{APIKeyID: id} + } + + query := ` + SELECT + api_key_id, + COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $2 AND created_at < $3), 0) as total_cost, + COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $4), 0) as today_cost + FROM usage_logs + WHERE api_key_id = ANY($1) + AND created_at >= LEAST($2, $4) + GROUP BY api_key_id + ` + today := timezone.Today() + rows, err := r.sql.QueryContext(ctx, query, pq.Array(normalizedAPIKeyIDs), startTime, endTime, today) + if err != nil { + return nil, err + } + for rows.Next() { + var apiKeyID int64 + var total float64 + var todayTotal float64 + if err := rows.Scan(&apiKeyID, &total, &todayTotal); err != nil { + _ = rows.Close() + return nil, err + } + if stats, ok := result[apiKeyID]; ok { + stats.TotalActualCost = total + stats.TodayActualCost = todayTotal + } + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + + return result, nil +} + +// resolveEndpointColumn maps endpoint type to the corresponding DB column name. +func resolveEndpointColumn(endpointType string) string { + switch endpointType { + case "upstream": + return "ul.upstream_endpoint" + case "path": + return "ul.inbound_endpoint || ' -> ' || ul.upstream_endpoint" + default: + return "ul.inbound_endpoint" + } +} + +// GetGlobalStats gets usage statistics for all users within a time range +func (r *usageLogRepository) GetGlobalStats(ctx context.Context, startTime, endTime time.Time) (*UsageStats, error) { + query := ` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(AVG(duration_ms), 0) as avg_duration_ms + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + ` + + stats := &UsageStats{} + if err := scanSingleRow( + ctx, + r.sql, + query, + []any{startTime, endTime}, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &stats.AverageDurationMs, + ); err != nil { + return nil, err + } + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens + return stats, nil +} + +// GetStatsWithFilters gets usage statistics with optional filters +func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters UsageLogFilters) (*UsageStats, error) { + conditions := make([]string, 0, 9) + args := make([]any, 0, 9) + + if filters.UserID > 0 { + conditions = append(conditions, fmt.Sprintf("user_id = $%d", len(args)+1)) + args = append(args, filters.UserID) + } + if filters.APIKeyID > 0 { + conditions = append(conditions, fmt.Sprintf("api_key_id = $%d", len(args)+1)) + args = append(args, filters.APIKeyID) + } + if filters.AccountID > 0 { + conditions = append(conditions, fmt.Sprintf("account_id = $%d", len(args)+1)) + args = append(args, filters.AccountID) + } + if filters.GroupID > 0 { + conditions = append(conditions, fmt.Sprintf("group_id = $%d", len(args)+1)) + args = append(args, filters.GroupID) + } + conditions, args = appendUsageLogModelWhereCondition(conditions, args, filters.Model, filters.ModelFilterSource) + conditions, args = appendRequestTypeOrStreamWhereCondition(conditions, args, filters.RequestType, filters.Stream) + if filters.BillingType != nil { + conditions = append(conditions, fmt.Sprintf("billing_type = $%d", len(args)+1)) + args = append(args, int16(*filters.BillingType)) + } + conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode) + if filters.StartTime != nil { + conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1)) + args = append(args, *filters.StartTime) + } + if filters.EndTime != nil { + conditions = append(conditions, fmt.Sprintf("created_at < $%d", len(args)+1)) + args = append(args, *filters.EndTime) + } + + query := fmt.Sprintf(` + SELECT + COUNT(*) as total_requests, + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens, + COALESCE(SUM(total_cost), 0) as total_cost, + COALESCE(SUM(actual_cost), 0) as total_actual_cost, + COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as total_account_cost, + COALESCE(AVG(duration_ms), 0) as avg_duration_ms + FROM usage_logs + %s + `, buildWhere(conditions)) + + stats := &UsageStats{} + var totalAccountCost float64 + + start := time.Unix(0, 0).UTC() + if filters.StartTime != nil { + start = *filters.StartTime + } + end := time.Now().UTC() + if filters.EndTime != nil { + end = *filters.EndTime + } + + var endpoints, upstreamEndpoints, endpointPaths []EndpointStat + + // 汇总查询:失败即致命。 + runSummary := func(c context.Context) error { + return scanSingleRow( + c, r.sql, query, args, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalCacheCreationTokens, + &stats.TotalCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &totalAccountCost, + &stats.AverageDurationMs, + ) + } + // endpoint 明细:best-effort(失败 log + 返空),不致命。 + runEndpoints := func(c context.Context) { + res, err := r.getEndpointStatsByColumnWithFilters(c, "inbound_endpoint", start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + endpoints = res + } + runUpstream := func(c context.Context) { + res, err := r.getEndpointStatsByColumnWithFilters(c, "upstream_endpoint", start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + upstreamEndpoints = res + } + runPaths := func(c context.Context) { + res, err := r.getEndpointPathStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "getEndpointPathStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + endpointPaths = res + } + + if r.db != nil { + // 生产路径:r.sql 是 *sql.DB 连接池,可并发。4 条查询并行,延迟取最大值。 + g, gctx := errgroup.WithContext(ctx) + g.Go(func() error { return runSummary(gctx) }) + g.Go(func() error { runEndpoints(gctx); return nil }) + g.Go(func() error { runUpstream(gctx); return nil }) + g.Go(func() error { runPaths(gctx); return nil }) + if err := g.Wait(); err != nil { + return nil, err + } + } else { + // 事务路径(ent.Tx 不能并发查询):顺序执行,行为与重构前一致。 + if err := runSummary(ctx); err != nil { + return nil, err + } + runEndpoints(ctx) + runUpstream(ctx) + runPaths(ctx) + } + + stats.TotalAccountCost = &totalAccountCost + stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens + stats.Endpoints = endpoints + stats.UpstreamEndpoints = upstreamEndpoints + stats.EndpointPaths = endpointPaths + + return stats, nil +} + +// AccountUsageHistory represents daily usage history for an account +type AccountUsageHistory = usagestats.AccountUsageHistory + +// AccountUsageSummary represents summary statistics for an account +type AccountUsageSummary = usagestats.AccountUsageSummary + +// AccountUsageStatsResponse represents the full usage statistics response for an account +type AccountUsageStatsResponse = usagestats.AccountUsageStatsResponse + +// EndpointStat represents endpoint usage statistics row. +type EndpointStat = usagestats.EndpointStat + +func (r *usageLogRepository) getEndpointStatsByColumnWithFilters(ctx context.Context, endpointColumn string, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []EndpointStat, err error) { + actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost" + if accountID > 0 && userID == 0 && apiKeyID == 0 { + actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost" + } + + query := fmt.Sprintf(` + SELECT + COALESCE(NULLIF(TRIM(%s), ''), 'unknown') AS endpoint, + COUNT(*) AS requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, + COALESCE(SUM(total_cost), 0) as cost, + %s + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + `, endpointColumn, actualCostExpr) + + args := []any{startTime, endTime} + if userID > 0 { + query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) + args = append(args, userID) + } + if apiKeyID > 0 { + query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) + args = append(args, apiKeyID) + } + if accountID > 0 { + query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) + args = append(args, accountID) + } + if groupID > 0 { + query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) + args = append(args, groupID) + } + query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource) + query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) + if billingType != nil { + query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) + args = append(args, int16(*billingType)) + } + query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") + query += " GROUP BY endpoint ORDER BY requests DESC" + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results = make([]EndpointStat, 0) + for rows.Next() { + var row EndpointStat + if err := rows.Scan(&row.Endpoint, &row.Requests, &row.TotalTokens, &row.Cost, &row.ActualCost); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} + +func (r *usageLogRepository) getEndpointPathStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []EndpointStat, err error) { + actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost" + if accountID > 0 && userID == 0 && apiKeyID == 0 { + actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost" + } + + query := fmt.Sprintf(` + SELECT + CONCAT( + COALESCE(NULLIF(TRIM(inbound_endpoint), ''), 'unknown'), + ' -> ', + COALESCE(NULLIF(TRIM(upstream_endpoint), ''), 'unknown') + ) AS endpoint, + COUNT(*) AS requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, + COALESCE(SUM(total_cost), 0) as cost, + %s + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + `, actualCostExpr) + + args := []any{startTime, endTime} + if userID > 0 { + query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) + args = append(args, userID) + } + if apiKeyID > 0 { + query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) + args = append(args, apiKeyID) + } + if accountID > 0 { + query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) + args = append(args, accountID) + } + if groupID > 0 { + query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) + args = append(args, groupID) + } + query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource) + query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) + if billingType != nil { + query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) + args = append(args, int16(*billingType)) + } + query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") + query += " GROUP BY endpoint ORDER BY requests DESC" + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results = make([]EndpointStat, 0) + for rows.Next() { + var row EndpointStat + if err := rows.Scan(&row.Endpoint, &row.Requests, &row.TotalTokens, &row.Cost, &row.ActualCost); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} + +// GetEndpointStatsWithFilters returns inbound endpoint statistics with optional filters. +func (r *usageLogRepository) GetEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]EndpointStat, error) { + return r.getEndpointStatsByColumnWithFilters(ctx, "inbound_endpoint", startTime, endTime, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "") +} + +// GetUpstreamEndpointStatsWithFilters returns upstream endpoint statistics with optional filters. +func (r *usageLogRepository) GetUpstreamEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]EndpointStat, error) { + return r.getEndpointStatsByColumnWithFilters(ctx, "upstream_endpoint", startTime, endTime, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "") +} + +// GetAccountUsageStats returns comprehensive usage statistics for an account over a time range +func (r *usageLogRepository) GetAccountUsageStats(ctx context.Context, accountID int64, startTime, endTime time.Time) (resp *AccountUsageStatsResponse, err error) { + daysCount := int(endTime.Sub(startTime).Hours()/24) + 1 + if daysCount <= 0 { + daysCount = 30 + } + + query := ` + SELECT + TO_CHAR(created_at, 'YYYY-MM-DD') as date, + COUNT(*) as requests, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens, + COALESCE(SUM(total_cost), 0) as cost, + COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost, + COALESCE(SUM(actual_cost), 0) as user_cost + FROM usage_logs + WHERE account_id = $1 AND created_at >= $2 AND created_at < $3 + GROUP BY date + ORDER BY date ASC + ` + + rows, err := r.sql.QueryContext(ctx, query, accountID, startTime, endTime) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + resp = nil + } + }() + + history := make([]AccountUsageHistory, 0) + for rows.Next() { + var date string + var requests int64 + var tokens int64 + var cost float64 + var actualCost float64 + var userCost float64 + if err = rows.Scan(&date, &requests, &tokens, &cost, &actualCost, &userCost); err != nil { + return nil, err + } + t, _ := time.Parse("2006-01-02", date) + history = append(history, AccountUsageHistory{ + Date: date, + Label: t.Format("01/02"), + Requests: requests, + Tokens: tokens, + Cost: cost, + ActualCost: actualCost, + UserCost: userCost, + }) + } + if err = rows.Err(); err != nil { + return nil, err + } + + var totalAccountCost, totalUserCost, totalStandardCost float64 + var totalRequests, totalTokens int64 + var highestCostDay, highestRequestDay *AccountUsageHistory + + for i := range history { + h := &history[i] + totalAccountCost += h.ActualCost + totalUserCost += h.UserCost + totalStandardCost += h.Cost + totalRequests += h.Requests + totalTokens += h.Tokens + + if highestCostDay == nil || h.ActualCost > highestCostDay.ActualCost { + highestCostDay = h + } + if highestRequestDay == nil || h.Requests > highestRequestDay.Requests { + highestRequestDay = h + } + } + + actualDaysUsed := len(history) + if actualDaysUsed == 0 { + actualDaysUsed = 1 + } + + avgQuery := "SELECT COALESCE(AVG(duration_ms), 0) as avg_duration_ms FROM usage_logs WHERE account_id = $1 AND created_at >= $2 AND created_at < $3" + var avgDuration float64 + if err := scanSingleRow(ctx, r.sql, avgQuery, []any{accountID, startTime, endTime}, &avgDuration); err != nil { + return nil, err + } + + summary := AccountUsageSummary{ + Days: daysCount, + ActualDaysUsed: actualDaysUsed, + TotalCost: totalAccountCost, + TotalUserCost: totalUserCost, + TotalStandardCost: totalStandardCost, + TotalRequests: totalRequests, + TotalTokens: totalTokens, + AvgDailyCost: totalAccountCost / float64(actualDaysUsed), + AvgDailyUserCost: totalUserCost / float64(actualDaysUsed), + AvgDailyRequests: float64(totalRequests) / float64(actualDaysUsed), + AvgDailyTokens: float64(totalTokens) / float64(actualDaysUsed), + AvgDurationMs: avgDuration, + } + + todayStr := timezone.Now().Format("2006-01-02") + for i := range history { + if history[i].Date == todayStr { + summary.Today = &struct { + Date string `json:"date"` + Cost float64 `json:"cost"` + UserCost float64 `json:"user_cost"` + Requests int64 `json:"requests"` + Tokens int64 `json:"tokens"` + }{ + Date: history[i].Date, + Cost: history[i].ActualCost, + UserCost: history[i].UserCost, + Requests: history[i].Requests, + Tokens: history[i].Tokens, + } + break + } + } + + if highestCostDay != nil { + summary.HighestCostDay = &struct { + Date string `json:"date"` + Label string `json:"label"` + Cost float64 `json:"cost"` + UserCost float64 `json:"user_cost"` + Requests int64 `json:"requests"` + }{ + Date: highestCostDay.Date, + Label: highestCostDay.Label, + Cost: highestCostDay.ActualCost, + UserCost: highestCostDay.UserCost, + Requests: highestCostDay.Requests, + } + } + + if highestRequestDay != nil { + summary.HighestRequestDay = &struct { + Date string `json:"date"` + Label string `json:"label"` + Requests int64 `json:"requests"` + Cost float64 `json:"cost"` + UserCost float64 `json:"user_cost"` + }{ + Date: highestRequestDay.Date, + Label: highestRequestDay.Label, + Requests: highestRequestDay.Requests, + Cost: highestRequestDay.ActualCost, + UserCost: highestRequestDay.UserCost, + } + } + + models, err := r.GetModelStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, nil, nil, nil) + if err != nil { + models = []ModelStat{} + } + endpoints, endpointErr := r.GetEndpointStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, "", nil, nil, nil) + if endpointErr != nil { + logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetAccountUsageStats: %v", endpointErr) + endpoints = []EndpointStat{} + } + upstreamEndpoints, upstreamEndpointErr := r.GetUpstreamEndpointStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, "", nil, nil, nil) + if upstreamEndpointErr != nil { + logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetAccountUsageStats: %v", upstreamEndpointErr) + upstreamEndpoints = []EndpointStat{} + } + + resp = &AccountUsageStatsResponse{ + History: history, + Summary: summary, + Models: models, + Endpoints: endpoints, + UpstreamEndpoints: upstreamEndpoints, + } + return resp, nil +} diff --git a/backend/internal/repository/usage_log_repo_trend.go b/backend/internal/repository/usage_log_repo_trend.go new file mode 100644 index 0000000000..7302296683 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_trend.go @@ -0,0 +1,811 @@ +package repository + +import ( + "context" + "database/sql" + "fmt" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" +) + +// TrendDataPoint represents a single point in trend data +type TrendDataPoint = usagestats.TrendDataPoint + +// ModelStat represents usage statistics for a single model +type ModelStat = usagestats.ModelStat + +// UserUsageTrendPoint represents user usage trend data point +type UserUsageTrendPoint = usagestats.UserUsageTrendPoint + +// UserSpendingRankingItem represents a user spending ranking row. +type UserSpendingRankingItem = usagestats.UserSpendingRankingItem +type UserSpendingRankingResponse = usagestats.UserSpendingRankingResponse + +// APIKeyUsageTrendPoint represents API key usage trend data point +type APIKeyUsageTrendPoint = usagestats.APIKeyUsageTrendPoint + +// GetAPIKeyUsageTrend returns usage trend data grouped by API key and date +func (r *usageLogRepository) GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) (results []APIKeyUsageTrendPoint, err error) { + dateFormat := safeDateFormat(granularity) + + query := fmt.Sprintf(` + WITH top_keys AS ( + SELECT api_key_id + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + GROUP BY api_key_id + ORDER BY SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) DESC + LIMIT $3 + ) + SELECT + TO_CHAR(u.created_at, '%s') as date, + u.api_key_id, + COALESCE(k.name, '') as key_name, + COUNT(*) as requests, + COALESCE(SUM(u.input_tokens + u.output_tokens + u.cache_creation_tokens + u.cache_read_tokens), 0) as tokens + FROM usage_logs u + LEFT JOIN api_keys k ON u.api_key_id = k.id + WHERE u.api_key_id IN (SELECT api_key_id FROM top_keys) + AND u.created_at >= $4 AND u.created_at < $5 + GROUP BY date, u.api_key_id, k.name + ORDER BY date ASC, tokens DESC + `, dateFormat) + + rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit, startTime, endTime) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results = make([]APIKeyUsageTrendPoint, 0) + for rows.Next() { + var row APIKeyUsageTrendPoint + if err = rows.Scan(&row.Date, &row.APIKeyID, &row.KeyName, &row.Requests, &row.Tokens); err != nil { + return nil, err + } + results = append(results, row) + } + if err = rows.Err(); err != nil { + return nil, err + } + + return results, nil +} + +// GetUserUsageTrend returns usage trend data grouped by user and date +func (r *usageLogRepository) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) (results []UserUsageTrendPoint, err error) { + dateFormat := safeDateFormat(granularity) + + query := fmt.Sprintf(` + WITH top_users AS ( + SELECT user_id + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + GROUP BY user_id + ORDER BY SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) DESC + LIMIT $3 + ) + SELECT + TO_CHAR(u.created_at, '%s') as date, + u.user_id, + COALESCE(us.email, '') as email, + COALESCE(us.username, '') as username, + COUNT(*) as requests, + COALESCE(SUM(u.input_tokens + u.output_tokens + u.cache_creation_tokens + u.cache_read_tokens), 0) as tokens, + COALESCE(SUM(u.total_cost), 0) as cost, + COALESCE(SUM(u.actual_cost), 0) as actual_cost + FROM usage_logs u + LEFT JOIN users us ON u.user_id = us.id + WHERE u.user_id IN (SELECT user_id FROM top_users) + AND u.created_at >= $4 AND u.created_at < $5 + GROUP BY date, u.user_id, us.email, us.username + ORDER BY date ASC, tokens DESC + `, dateFormat) + + rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit, startTime, endTime) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results = make([]UserUsageTrendPoint, 0) + for rows.Next() { + var row UserUsageTrendPoint + if err = rows.Scan(&row.Date, &row.UserID, &row.Email, &row.Username, &row.Requests, &row.Tokens, &row.Cost, &row.ActualCost); err != nil { + return nil, err + } + results = append(results, row) + } + if err = rows.Err(); err != nil { + return nil, err + } + + return results, nil +} + +// GetUserSpendingRanking returns user spending ranking aggregated within the time range. +func (r *usageLogRepository) GetUserSpendingRanking(ctx context.Context, startTime, endTime time.Time, limit int) (result *UserSpendingRankingResponse, err error) { + if limit <= 0 { + limit = 12 + } + + query := ` + WITH user_spend AS ( + SELECT + u.user_id, + COALESCE(us.email, '') as email, + COALESCE(SUM(u.actual_cost), 0) as actual_cost, + COUNT(*) as requests, + COALESCE(SUM(u.input_tokens + u.output_tokens + u.cache_creation_tokens + u.cache_read_tokens), 0) as tokens + FROM usage_logs u + LEFT JOIN users us ON u.user_id = us.id + WHERE u.created_at >= $1 AND u.created_at < $2 + GROUP BY u.user_id, us.email + ), + ranked AS ( + SELECT + user_id, + email, + actual_cost, + requests, + tokens, + COALESCE(SUM(actual_cost) OVER (), 0) as total_actual_cost, + COALESCE(SUM(requests) OVER (), 0) as total_requests, + COALESCE(SUM(tokens) OVER (), 0) as total_tokens + FROM user_spend + ORDER BY actual_cost DESC, tokens DESC, user_id ASC + LIMIT $3 + ) + SELECT + user_id, + email, + actual_cost, + requests, + tokens, + total_actual_cost, + total_requests, + total_tokens + FROM ranked + ORDER BY actual_cost DESC, tokens DESC, user_id ASC + ` + + rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + result = nil + } + }() + + ranking := make([]UserSpendingRankingItem, 0) + totalActualCost := 0.0 + totalRequests := int64(0) + totalTokens := int64(0) + for rows.Next() { + var row UserSpendingRankingItem + if err = rows.Scan(&row.UserID, &row.Email, &row.ActualCost, &row.Requests, &row.Tokens, &totalActualCost, &totalRequests, &totalTokens); err != nil { + return nil, err + } + ranking = append(ranking, row) + } + if err = rows.Err(); err != nil { + return nil, err + } + + return &UserSpendingRankingResponse{ + Ranking: ranking, + TotalActualCost: totalActualCost, + TotalRequests: totalRequests, + TotalTokens: totalTokens, + }, nil +} + +// GetUserUsageTrendByUserID 获取指定用户的使用趋势 +func (r *usageLogRepository) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (results []TrendDataPoint, err error) { + dateFormat := safeDateFormat(granularity) + + query := fmt.Sprintf(` + SELECT + TO_CHAR(created_at, '%s') as date, + COUNT(*) as requests, + COALESCE(SUM(input_tokens), 0) as input_tokens, + COALESCE(SUM(output_tokens), 0) as output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(total_cost), 0) as cost, + COALESCE(SUM(actual_cost), 0) as actual_cost + FROM usage_logs + WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 + GROUP BY date + ORDER BY date ASC + `, dateFormat) + + rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results, err = scanTrendRows(rows) + if err != nil { + return nil, err + } + return results, nil +} + +// GetUserModelStats 获取指定用户的模型统计 +func (r *usageLogRepository) GetUserModelStats(ctx context.Context, userID int64, startTime, endTime time.Time) (results []ModelStat, err error) { + return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, 0, 0, 0, "", nil, nil, nil, usagestats.ModelSourceRequested, "") +} + +// GetUsageTrendWithFilters returns usage trend data with optional filters +func (r *usageLogRepository) GetUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) (results []TrendDataPoint, err error) { + return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "") +} + +func (r *usageLogRepository) GetUsageTrendWithUsageFilters(ctx context.Context, startTime, endTime time.Time, granularity string, filters UsageLogFilters) (results []TrendDataPoint, err error) { + return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) +} + +func (r *usageLogRepository) getUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []TrendDataPoint, err error) { + if shouldUsePreaggregatedTrend(granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType, billingMode) { + aggregated, aggregatedErr := r.getUsageTrendFromAggregates(ctx, startTime, endTime, granularity) + if aggregatedErr == nil && len(aggregated) > 0 { + return aggregated, nil + } + } + + dateFormat := safeDateFormat(granularity) + + query := fmt.Sprintf(` + SELECT + TO_CHAR(created_at, '%s') as date, + COUNT(*) as requests, + COALESCE(SUM(input_tokens), 0) as input_tokens, + COALESCE(SUM(output_tokens), 0) as output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(total_cost), 0) as cost, + COALESCE(SUM(actual_cost), 0) as actual_cost + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + `, dateFormat) + + args := []any{startTime, endTime} + if userID > 0 { + query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) + args = append(args, userID) + } + if apiKeyID > 0 { + query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) + args = append(args, apiKeyID) + } + if accountID > 0 { + query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) + args = append(args, accountID) + } + if groupID > 0 { + query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) + args = append(args, groupID) + } + query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource) + query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) + if billingType != nil { + query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) + args = append(args, int16(*billingType)) + } + query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") + query += " GROUP BY date ORDER BY date ASC" + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results, err = scanTrendRows(rows) + if err != nil { + return nil, err + } + return results, nil +} + +func shouldUsePreaggregatedTrend(granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string) bool { + if granularity != "day" && granularity != "hour" { + return false + } + return userID == 0 && + apiKeyID == 0 && + accountID == 0 && + groupID == 0 && + model == "" && + requestType == nil && + stream == nil && + billingType == nil && + billingMode == "" +} + +func (r *usageLogRepository) getUsageTrendFromAggregates(ctx context.Context, startTime, endTime time.Time, granularity string) (results []TrendDataPoint, err error) { + dateFormat := safeDateFormat(granularity) + query := "" + args := []any{startTime, endTime} + + switch granularity { + case "hour": + query = fmt.Sprintf(` + SELECT + TO_CHAR(bucket_start, '%s') as date, + total_requests as requests, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) as total_tokens, + total_cost as cost, + actual_cost + FROM usage_dashboard_hourly + WHERE bucket_start >= $1 AND bucket_start < $2 + ORDER BY bucket_start ASC + `, dateFormat) + case "day": + query = fmt.Sprintf(` + SELECT + TO_CHAR(bucket_date::timestamp, '%s') as date, + total_requests as requests, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) as total_tokens, + total_cost as cost, + actual_cost + FROM usage_dashboard_daily + WHERE bucket_date >= $1::date AND bucket_date < $2::date + ORDER BY bucket_date ASC + `, dateFormat) + default: + return nil, nil + } + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results, err = scanTrendRows(rows) + if err != nil { + return nil, err + } + return results, nil +} + +// GetModelStatsWithFilters returns model statistics with optional filters +func (r *usageLogRepository) GetModelStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []ModelStat, err error) { + return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, usagestats.ModelSourceRequested, "") +} + +// GetModelStatsWithFiltersBySource returns model statistics with optional filters and model source dimension. +// source: requested | upstream | mapping. +func (r *usageLogRepository) GetModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8, source string) (results []ModelStat, err error) { + return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, source, "") +} + +func (r *usageLogRepository) GetModelStatsWithUsageFiltersBySource(ctx context.Context, startTime, endTime time.Time, filters UsageLogFilters, source string) (results []ModelStat, err error) { + return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, source, filters.BillingMode) +} + +func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, source string, billingMode string) (results []ModelStat, err error) { + actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost" + // 当仅按 account_id 聚合时,实际费用使用账号倍率(total_cost * account_rate_multiplier)。 + if accountID > 0 && userID == 0 && apiKeyID == 0 { + actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost" + } + accountCostExpr := "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as account_cost" + modelExpr := resolveModelDimensionExpression(source) + + query := fmt.Sprintf(` + SELECT + %s as model, + COUNT(*) as requests, + COALESCE(SUM(input_tokens), 0) as input_tokens, + COALESCE(SUM(output_tokens), 0) as output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(total_cost), 0) as cost, + %s, + %s + FROM usage_logs + WHERE created_at >= $1 AND created_at < $2 + `, modelExpr, actualCostExpr, accountCostExpr) + + args := []any{startTime, endTime} + if userID > 0 { + query += fmt.Sprintf(" AND user_id = $%d", len(args)+1) + args = append(args, userID) + } + if apiKeyID > 0 { + query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1) + args = append(args, apiKeyID) + } + if accountID > 0 { + query += fmt.Sprintf(" AND account_id = $%d", len(args)+1) + args = append(args, accountID) + } + if groupID > 0 { + query += fmt.Sprintf(" AND group_id = $%d", len(args)+1) + args = append(args, groupID) + } + if strings.TrimSpace(model) != "" { + query += fmt.Sprintf(" AND %s = $%d", modelExpr, len(args)+1) + args = append(args, model) + } + query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) + if billingType != nil { + query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1) + args = append(args, int16(*billingType)) + } + query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "") + query += fmt.Sprintf(" GROUP BY %s ORDER BY total_tokens DESC", modelExpr) + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + // 保持主错误优先;仅在无错误时回传 Close 失败。 + // 同时清空返回值,避免误用不完整结果。 + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results, err = scanModelStatsRows(rows) + if err != nil { + return nil, err + } + return results, nil +} + +// GetGroupStatsWithFilters returns group usage statistics with optional filters +func (r *usageLogRepository) GetGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []usagestats.GroupStat, err error) { + return r.getGroupStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, "") +} + +func (r *usageLogRepository) GetGroupStatsWithUsageFilters(ctx context.Context, startTime, endTime time.Time, filters UsageLogFilters) (results []usagestats.GroupStat, err error) { + return r.getGroupStatsWithFilters(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode) +} + +func (r *usageLogRepository) getGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []usagestats.GroupStat, err error) { + query := ` + SELECT + COALESCE(ul.group_id, 0) as group_id, + COALESCE(g.name, '') as group_name, + COUNT(*) as requests, + COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(ul.total_cost), 0) as cost, + COALESCE(SUM(ul.actual_cost), 0) as actual_cost, + COALESCE(SUM(COALESCE(ul.account_stats_cost, ul.total_cost) * COALESCE(ul.account_rate_multiplier, 1)), 0) as account_cost + FROM usage_logs ul + LEFT JOIN groups g ON g.id = ul.group_id + WHERE ul.created_at >= $1 AND ul.created_at < $2 + ` + + args := []any{startTime, endTime} + if userID > 0 { + query += fmt.Sprintf(" AND ul.user_id = $%d", len(args)+1) + args = append(args, userID) + } + if apiKeyID > 0 { + query += fmt.Sprintf(" AND ul.api_key_id = $%d", len(args)+1) + args = append(args, apiKeyID) + } + if accountID > 0 { + query += fmt.Sprintf(" AND ul.account_id = $%d", len(args)+1) + args = append(args, accountID) + } + if groupID > 0 { + query += fmt.Sprintf(" AND ul.group_id = $%d", len(args)+1) + args = append(args, groupID) + } + if strings.TrimSpace(model) != "" { + modelExpr := resolveModelDimensionExpressionWithAlias(usagestats.ModelSourceRequested, "ul") + query += fmt.Sprintf(" AND %s = $%d", modelExpr, len(args)+1) + args = append(args, model) + } + query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream) + if billingType != nil { + query += fmt.Sprintf(" AND ul.billing_type = $%d", len(args)+1) + args = append(args, int16(*billingType)) + } + query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "ul") + query += " GROUP BY ul.group_id, g.name ORDER BY total_tokens DESC" + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results = make([]usagestats.GroupStat, 0) + for rows.Next() { + var row usagestats.GroupStat + if err := rows.Scan( + &row.GroupID, + &row.GroupName, + &row.Requests, + &row.TotalTokens, + &row.Cost, + &row.ActualCost, + &row.AccountCost, + ); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} + +// GetUserBreakdownStats returns per-user usage breakdown within a specific dimension. +func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) (results []usagestats.UserBreakdownItem, err error) { + query := ` + SELECT + COALESCE(ul.user_id, 0) as user_id, + COALESCE(u.email, '') as email, + COUNT(*) as requests, + COALESCE(SUM(ul.input_tokens), 0) as input_tokens, + COALESCE(SUM(ul.output_tokens), 0) as output_tokens, + COALESCE(SUM(ul.cache_creation_tokens + ul.cache_read_tokens), 0) as cache_tokens, + COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens, + COALESCE(SUM(ul.total_cost), 0) as cost, + COALESCE(SUM(ul.actual_cost), 0) as actual_cost, + COALESCE(SUM(COALESCE(ul.account_stats_cost, ul.total_cost) * COALESCE(ul.account_rate_multiplier, 1)), 0) as account_cost + FROM usage_logs ul + LEFT JOIN users u ON u.id = ul.user_id + WHERE ul.created_at >= $1 AND ul.created_at < $2 + ` + args := []any{startTime, endTime} + + if dim.GroupID > 0 { + query += fmt.Sprintf(" AND ul.group_id = $%d", len(args)+1) + args = append(args, dim.GroupID) + } + if dim.Model != "" { + query += fmt.Sprintf(" AND %s = $%d", resolveModelDimensionExpression(dim.ModelType), len(args)+1) + args = append(args, dim.Model) + } + if dim.Endpoint != "" { + col := resolveEndpointColumn(dim.EndpointType) + query += fmt.Sprintf(" AND %s = $%d", col, len(args)+1) + args = append(args, dim.Endpoint) + } + if dim.UserID > 0 { + query += fmt.Sprintf(" AND ul.user_id = $%d", len(args)+1) + args = append(args, dim.UserID) + } + if dim.APIKeyID > 0 { + query += fmt.Sprintf(" AND ul.api_key_id = $%d", len(args)+1) + args = append(args, dim.APIKeyID) + } + if dim.AccountID > 0 { + query += fmt.Sprintf(" AND ul.account_id = $%d", len(args)+1) + args = append(args, dim.AccountID) + } + if dim.RequestType != nil { + query += fmt.Sprintf(" AND ul.request_type = $%d", len(args)+1) + args = append(args, *dim.RequestType) + } + if dim.Stream != nil { + query += fmt.Sprintf(" AND ul.stream = $%d", len(args)+1) + args = append(args, *dim.Stream) + } + if dim.BillingType != nil { + query += fmt.Sprintf(" AND ul.billing_type = $%d", len(args)+1) + args = append(args, *dim.BillingType) + } + + // ORDER BY 列来自固定 allowlist(非用户原样字符串),避免 SQL 注入。 + orderBy := "actual_cost" + switch dim.SortBy { + case "total_tokens", "input_tokens", "output_tokens", "cache_tokens", "requests", "cost", "actual_cost": + orderBy = dim.SortBy + } + query += " GROUP BY ul.user_id, u.email ORDER BY " + orderBy + " DESC" + if limit > 0 { + query += fmt.Sprintf(" LIMIT %d", limit) + } + + rows, err := r.sql.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { + if closeErr := rows.Close(); closeErr != nil && err == nil { + err = closeErr + results = nil + } + }() + + results = make([]usagestats.UserBreakdownItem, 0) + for rows.Next() { + var row usagestats.UserBreakdownItem + if err := rows.Scan( + &row.UserID, + &row.Email, + &row.Requests, + &row.InputTokens, + &row.OutputTokens, + &row.CacheTokens, + &row.TotalTokens, + &row.Cost, + &row.ActualCost, + &row.AccountCost, + ); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} + +// GetAllGroupUsageSummary returns today's and cumulative actual_cost for every group. +// todayStart is the start-of-day in the caller's timezone (UTC-based). +// TODO(perf): This query scans ALL usage_logs rows for total_cost aggregation. +// When usage_logs exceeds ~1M rows, consider adding a short-lived cache (30s) +// or a materialized view / pre-aggregation table for cumulative costs. +func (r *usageLogRepository) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) { + query := ` + SELECT + g.id AS group_id, + COALESCE(SUM(ul.actual_cost), 0) AS total_cost, + COALESCE(SUM(CASE WHEN ul.created_at >= $1 THEN ul.actual_cost ELSE 0 END), 0) AS today_cost + FROM groups g + LEFT JOIN usage_logs ul ON ul.group_id = g.id + GROUP BY g.id + ` + + rows, err := r.sql.QueryContext(ctx, query, todayStart) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var results []usagestats.GroupUsageSummary + for rows.Next() { + var row usagestats.GroupUsageSummary + if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} + +// resolveModelDimensionExpression maps model source type to a safe SQL expression. +func resolveModelDimensionExpression(modelType string) string { + return resolveModelDimensionExpressionWithAlias(modelType, "") +} + +func resolveModelDimensionExpressionWithAlias(modelType, alias string) string { + column := func(name string) string { + if alias == "" { + return name + } + return alias + "." + name + } + requestedExpr := fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("requested_model"), column("model")) + switch usagestats.NormalizeModelSource(modelType) { + case usagestats.ModelSourceUpstream: + return fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("upstream_model"), requestedExpr) + case usagestats.ModelSourceMapping: + return fmt.Sprintf("(%s || ' -> ' || COALESCE(NULLIF(TRIM(%s), ''), %s))", requestedExpr, column("upstream_model"), requestedExpr) + default: + return requestedExpr + } +} + +func scanTrendRows(rows *sql.Rows) ([]TrendDataPoint, error) { + results := make([]TrendDataPoint, 0) + for rows.Next() { + var row TrendDataPoint + if err := rows.Scan( + &row.Date, + &row.Requests, + &row.InputTokens, + &row.OutputTokens, + &row.CacheCreationTokens, + &row.CacheReadTokens, + &row.TotalTokens, + &row.Cost, + &row.ActualCost, + ); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} + +func scanModelStatsRows(rows *sql.Rows) ([]ModelStat, error) { + results := make([]ModelStat, 0) + for rows.Next() { + var row ModelStat + if err := rows.Scan( + &row.Model, + &row.Requests, + &row.InputTokens, + &row.OutputTokens, + &row.CacheCreationTokens, + &row.CacheReadTokens, + &row.TotalTokens, + &row.Cost, + &row.ActualCost, + &row.AccountCost, + ); err != nil { + return nil, err + } + results = append(results, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + return results, nil +} diff --git a/backend/internal/repository/user_msg_queue_cache.go b/backend/internal/repository/user_msg_queue_cache.go index 67e75a87cf..9b7707614c 100644 --- a/backend/internal/repository/user_msg_queue_cache.go +++ b/backend/internal/repository/user_msg_queue_cache.go @@ -6,6 +6,7 @@ import ( "fmt" "strconv" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" ) @@ -25,6 +26,9 @@ const ( // Lua 脚本:原子获取串行锁(SET NX PX + 重入安全) // 返回 {是否获取成功, 锁预计过期时间毫秒},让 Go 侧用同一 Redis 时间源更新索引。 +// 获取失败(锁被他人持有)时也返回观测到的到期时间,供 Go 侧回填锁索引: +// 这让升级窗口遗留、索引写失败、释放竞态误删索引的存量锁在下一次被争用时自动重新入索引, +// 是替代旧 SCAN 兜底的自愈机制。PTTL == -1 的异常锁返回当前时间,使其立即成为 reconcile 候选。 var acquireLockScript = redis.NewScript(` redis.replicate_commands() local cur = redis.call('GET', KEYS[1]) @@ -35,7 +39,15 @@ if cur == ARGV[1] then local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000) return {1, ms + ttl} end -if cur ~= false then return {0, 0} end +if cur ~= false then + local t = redis.call('TIME') + local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000) + local pttl = redis.call('PTTL', KEYS[1]) + if pttl and pttl > 0 then + return {0, ms + pttl} + end + return {0, ms} +end redis.call('SET', KEYS[1], ARGV[1], 'PX', ttl) local t = redis.call('TIME') local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000) @@ -92,7 +104,8 @@ func umqLastKey(accountID int64) string { } // AcquireLock 尝试获取账号级串行锁 -// 成功后尽力写入锁索引,后台清理只需要看“到期候选”而不是扫描所有锁 key。 +// 无论成功与否都尽力写入锁索引:成功时登记自己的锁,失败时回填观测到的持有者锁, +// 保证任何被争用的锁都能被后台 reconcile 发现,无需扫描所有锁 key。 func (c *userMsgQueueCache) AcquireLock(ctx context.Context, accountID int64, requestID string, lockTtlMs int) (bool, error) { key := umqLockKey(accountID) result, err := acquireLockScript.Run(ctx, c.rdb, []string{key}, requestID, lockTtlMs).Result() @@ -107,11 +120,13 @@ func (c *userMsgQueueCache) AcquireLock(ctx context.Context, accountID int64, re if err != nil { return false, fmt.Errorf("umq parse acquire lock expire: %w", err) } - if acquired == 1 { - _ = c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{ + if expireAtMs > 0 { + if err := c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{ Score: float64(expireAtMs), Member: strconv.FormatInt(accountID, 10), - }).Err() + }).Err(); err != nil { + logger.LegacyPrintf("repository.umq", "Warning: update lock index for account %d failed: %v", accountID, err) + } } return acquired == 1, nil } @@ -126,7 +141,11 @@ func (c *userMsgQueueCache) ReleaseLock(ctx context.Context, accountID int64, re return false, fmt.Errorf("umq release lock: %w", err) } if result == 1 { - _ = c.rdb.ZRem(ctx, umqLockIndexKey, strconv.FormatInt(accountID, 10)).Err() + // 与下一个 AcquireLock 的 ZAdd 存在竞态:可能误删新持有者刚写入的索引项。 + // 该锁下次被争用时 AcquireLock 的回填路径会重新登记,无需在此加锁。 + if err := c.rdb.ZRem(ctx, umqLockIndexKey, strconv.FormatInt(accountID, 10)).Err(); err != nil { + logger.LegacyPrintf("repository.umq", "Warning: remove lock index for account %d failed: %v", accountID, err) + } } return result == 1, nil } @@ -180,7 +199,7 @@ func (c *userMsgQueueCache) ReconcileExpiredLockCandidates(ctx context.Context, for _, member := range members { accountID, err := strconv.ParseInt(member, 10, 64) if err != nil || accountID <= 0 { - _ = c.rdb.ZRem(ctx, umqLockIndexKey, member).Err() + c.removeLockIndexMember(ctx, member) continue } @@ -200,22 +219,31 @@ func (c *userMsgQueueCache) ReconcileExpiredLockCandidates(ctx context.Context, switch status { case -2: // 锁自然过期或已释放,只需移除索引残留。 - _ = c.rdb.ZRem(ctx, umqLockIndexKey, member).Err() + c.removeLockIndexMember(ctx, member) case -1: // 无 TTL 的锁会永久阻塞队列,Lua 已原子删除它,这里统计一次清理。 - _ = c.rdb.ZRem(ctx, umqLockIndexKey, member).Err() + c.removeLockIndexMember(ctx, member) cleaned++ case 1: // 锁仍存活,说明索引过期时间滞后;按剩余 PTTL 重新排期。 - _ = c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{ + if err := c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{ Score: float64(nowMs + pttl), Member: member, - }).Err() + }).Err(); err != nil { + logger.LegacyPrintf("repository.umq", "Warning: reschedule lock index member %s failed: %v", member, err) + } } } return cleaned, nil } +// removeLockIndexMember 移除锁索引残留;索引维护是 best-effort,失败只记日志。 +func (c *userMsgQueueCache) removeLockIndexMember(ctx context.Context, member string) { + if err := c.rdb.ZRem(ctx, umqLockIndexKey, member).Err(); err != nil { + logger.LegacyPrintf("repository.umq", "Warning: remove lock index member %s failed: %v", member, err) + } +} + // redisScriptInt64At 兼容 go-redis 对 Lua 数组元素的不同返回类型。 func redisScriptInt64At(result any, index int) (int64, error) { values, ok := result.([]any) diff --git a/backend/internal/repository/user_msg_queue_cache_integration_test.go b/backend/internal/repository/user_msg_queue_cache_integration_test.go index c61b658357..e683b44aa3 100644 --- a/backend/internal/repository/user_msg_queue_cache_integration_test.go +++ b/backend/internal/repository/user_msg_queue_cache_integration_test.go @@ -126,3 +126,52 @@ func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRemovesInvali _, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "not-an-account-id").Result() require.True(s.T(), errors.Is(err, redis.Nil)) } + +func (s *UserMsgQueueCacheSuite) TestAcquireLockBusyPathReindexesUnindexedLiveLock() { + // 模拟索引丢失的存量锁(升级窗口/索引写失败/释放竞态误删): + // 锁存在且有 TTL,但索引里没有对应 member。 + accountID := int64(705) + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + require.NoError(s.T(), err) + require.NoError(s.T(), s.rdb.Set(s.ctx, umqLockKey(accountID), "holder-705", time.Minute).Err()) + + // 另一个请求争锁失败,应顺手把观测到的持有者锁回填进索引。 + acquired, err := s.cache.AcquireLock(s.ctx, accountID, "contender-705", 10_000) + require.NoError(s.T(), err) + require.False(s.T(), acquired) + + score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "705").Result() + require.NoError(s.T(), err, "busy acquire should re-index the observed live lock") + require.Greater(s.T(), int64(score), nowMs) + // 锁本身不应被争锁方改动。 + val, err := s.rdb.Get(s.ctx, umqLockKey(accountID)).Result() + require.NoError(s.T(), err) + require.Equal(s.T(), "holder-705", val) +} + +func (s *UserMsgQueueCacheSuite) TestAcquireLockBusyPathMakesNoTTLLockReconcilable() { + // PTTL == -1 的异常锁若不在索引中,永远不会被 reconcile 发现; + // 争锁失败路径必须以“已到期候选”的 score 回填它,形成自愈闭环。 + accountID := int64(706) + require.NoError(s.T(), s.rdb.Set(s.ctx, umqLockKey(accountID), "holder-706", 0).Err()) + + acquired, err := s.cache.AcquireLock(s.ctx, accountID, "contender-706", 10_000) + require.NoError(s.T(), err) + require.False(s.T(), acquired) + + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + require.NoError(s.T(), err) + score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "706").Result() + require.NoError(s.T(), err, "busy acquire should index the anomalous lock") + require.LessOrEqual(s.T(), int64(score), nowMs, "anomalous lock should be an immediately-expired candidate") + + cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000) + require.NoError(s.T(), err) + require.Equal(s.T(), 1, cleaned, "reconcile should delete the no-TTL lock") + + exists, err := s.rdb.Exists(s.ctx, umqLockKey(accountID)).Result() + require.NoError(s.T(), err) + require.EqualValues(s.T(), 0, exists, "queue is unblocked after reconcile") + _, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "706").Result() + require.ErrorIs(s.T(), err, redis.Nil) +} diff --git a/backend/internal/repository/wire.go b/backend/internal/repository/wire.go index 37f8e9bd2f..ec5078eac8 100644 --- a/backend/internal/repository/wire.go +++ b/backend/internal/repository/wire.go @@ -77,6 +77,7 @@ var ProviderSet = wire.NewSet( NewAnnouncementReadRepository, NewUsageLogRepository, NewUsageBillingRepository, + NewBatchImageRepository, NewIdempotencyRepository, NewUsageCleanupRepository, NewDashboardAggregationRepository, @@ -115,6 +116,8 @@ var ProviderSet = wire.NewSet( NewRedeemCache, NewUpdateCache, NewGeminiTokenCache, + NewBatchImageQueue, + NewBatchImageDownloadLimiter, NewLeaderLockCache, ProvideSchedulerCache, NewSchedulerOutboxRepository, diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index e3e68bb6e7..f48ecb060a 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -52,9 +52,10 @@ func TestAPIContracts(t *testing.T) { "email": "alice@example.com", "email_bound": true, "username": "alice", - "role": "user", - "balance": 12.5, - "concurrency": 5, + "role": "user", + "balance": 12.5, + "frozen_balance": 0, + "concurrency": 5, "rpm_limit": 0, "status": "active", "allowed_groups": null, @@ -233,6 +234,7 @@ func TestAPIContracts(t *testing.T) { "ip_whitelist": null, "ip_blacklist": null, "last_used_at": null, + "last_used_ip": null, "current_concurrency": 0, "quota": 0, "quota_used": 0, @@ -283,6 +285,7 @@ func TestAPIContracts(t *testing.T) { "ip_whitelist": null, "ip_blacklist": null, "last_used_at": null, + "last_used_ip": null, "current_concurrency": 0, "quota": 0, "quota_used": 0, @@ -360,9 +363,17 @@ func TestAPIContracts(t *testing.T) { "image_price_1k": null, "image_price_2k": null, "image_price_4k": null, + "video_price_480p": null, + "video_price_720p": null, + "video_price_1080p": null, "allow_image_generation": false, + "allow_batch_image_generation": false, + "batch_image_discount_multiplier": 0, + "batch_image_hold_multiplier": 0, "image_rate_independent": false, "image_rate_multiplier": 0, + "video_rate_independent": false, + "video_rate_multiplier": 0, "claude_code_only": false, "allow_messages_dispatch": false, "fallback_group_id": null, diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index 2f0a3f1cf7..1610390440 100644 --- a/backend/internal/server/middleware/api_key_auth.go +++ b/backend/internal/server/middleware/api_key_auth.go @@ -213,7 +213,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti } } else { // 非订阅模式 或 订阅模式但 subscriptionService 未注入:回退到余额检查 - if apiKey.User.Balance <= 0 { + if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) { AbortWithError(c, 403, "INSUFFICIENT_BALANCE", "Insufficient account balance") return } @@ -289,6 +289,13 @@ func setGroupContext(c *gin.Context, group *service.Group) { c.Request = c.Request.WithContext(ctx) } +// apiKeyBalanceBelowAuthThreshold 保持鉴权层的历史语义:仅在余额耗尽(<=0)时拒绝。 +// MinimumBalanceReserve 只作为 billing-cache 预检的保守下限,不得复用为鉴权硬门槛, +// 否则已配置该值的存量部署升级后,0 < balance < reserve 的用户会在所有端点被静默 403。 +func apiKeyBalanceBelowAuthThreshold(balance float64, _ *config.Config) bool { + return balance <= 0 +} + func abortIfAPIKeyGroupUnavailable(c *gin.Context, apiKey *service.APIKey) bool { code, message, ok := validateAPIKeyGroupAvailable(apiKey) if ok { diff --git a/backend/internal/server/middleware/api_key_auth_google.go b/backend/internal/server/middleware/api_key_auth_google.go index 97f3936c0c..c75d5b99f2 100644 --- a/backend/internal/server/middleware/api_key_auth_google.go +++ b/backend/internal/server/middleware/api_key_auth_google.go @@ -2,10 +2,12 @@ package middleware import ( "errors" + "fmt" "strings" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/googleapi" + "github.com/Wei-Shaw/sub2api/internal/pkg/ip" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" @@ -46,10 +48,32 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs // user/group/platform。 SetOpsFallbackAPIKey(c, apiKey) - if !apiKey.IsActive() { + // disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段, + // 与主中间件 api_key_auth.go 保持一致)。 + if !apiKey.IsActive() && + apiKey.Status != service.StatusAPIKeyExpired && + apiKey.Status != service.StatusAPIKeyQuotaExhausted { abortWithGoogleError(c, 401, "API key is disabled") return } + + // 检查 IP 限制(白名单/黑名单)。与主中间件保持一致,避免 Gemini 端点绕过 Key 的 IP ACL。 + if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 { + clientIP := ip.GetTrustedClientIP(c) + if cfg.TrustForwardedIPForAPIKeyACL() { + clientIP = ip.GetClientIP(c) + } + allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist) + if !allowed { + if clientIP == "" { + clientIP = "unknown" + } + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction) + abortWithGoogleError(c, 403, fmt.Sprintf("Access denied. Your IP is %s", clientIP)) + return + } + } + if apiKey.User == nil { abortWithGoogleError(c, 401, "User associated with API key not found") return @@ -63,6 +87,12 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs abortWithGoogleError(c, 403, message) return } + // 专属分组授权校验:用户对该专属分组的授权被撤销后应拒绝(与主中间件一致,防止越权)。 + if !validateAPIKeyGroupAllowed(apiKey) { + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable) + abortWithGoogleError(c, 403, "API Key 所属专属分组不再允许当前用户使用") + return + } // 简易模式:跳过余额和订阅检查 if cfg.RunMode == config.RunModeSimple { @@ -78,6 +108,26 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs return } + // Key 状态检查(状态字段可能因后台异步刷新而滞后,故显式拦截)。 + switch apiKey.Status { + case service.StatusAPIKeyQuotaExhausted: + abortWithGoogleError(c, 429, "API key 额度已用完") + return + case service.StatusAPIKeyExpired: + abortWithGoogleError(c, 403, "API key 已过期") + return + } + + // 运行时过期/配额检查(即使状态是 active,也要检查时间和用量,与主中间件一致)。 + if apiKey.IsExpired() { + abortWithGoogleError(c, 403, "API key 已过期") + return + } + if apiKey.IsQuotaExhausted() { + abortWithGoogleError(c, 429, "API key 额度已用完") + return + } + isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType() if isSubscriptionType && subscriptionService != nil { subscription, err := subscriptionService.GetActiveSubscription( @@ -109,7 +159,7 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs subscriptionService.DoWindowMaintenance(&maintenanceCopy) } } else { - if apiKey.User.Balance <= 0 { + if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) { abortWithGoogleError(c, 403, "Insufficient account balance") return } diff --git a/backend/internal/server/middleware/api_key_auth_google_test.go b/backend/internal/server/middleware/api_key_auth_google_test.go index bf3909fcd4..45ddb0bf93 100644 --- a/backend/internal/server/middleware/api_key_auth_google_test.go +++ b/backend/internal/server/middleware/api_key_auth_google_test.go @@ -539,6 +539,74 @@ func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) { require.Equal(t, "PERMISSION_DENIED", resp.Error.Status) } +func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.T) { + gin.SetMode(gin.TestMode) + + // 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检, + // 0 < balance < reserve 的用户不得在鉴权中间件被硬 403。 + r := gin.New() + apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{ + getByKey: func(ctx context.Context, key string) (*service.APIKey, error) { + return &service.APIKey{ + ID: 1, + Key: key, + Status: service.StatusActive, + User: &service.User{ + ID: 123, + Status: service.StatusActive, + Balance: 0.005, + }, + }, nil + }, + }) + cfg := &config.Config{} + cfg.Billing.MinimumBalanceReserve = 0.01 + r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)) + r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) + + req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil) + req.Header.Set("Authorization", "Bearer ok") + rec := httptest.NewRecorder() + r.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) +} + +func TestApiKeyAuthWithSubscriptionGoogle_RejectsExhaustedBalance(t *testing.T) { + gin.SetMode(gin.TestMode) + + r := gin.New() + apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{ + getByKey: func(ctx context.Context, key string) (*service.APIKey, error) { + return &service.APIKey{ + ID: 1, + Key: key, + Status: service.StatusActive, + User: &service.User{ + ID: 123, + Status: service.StatusActive, + Balance: 0, + }, + }, nil + }, + }) + cfg := &config.Config{} + r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)) + r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) + + req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil) + req.Header.Set("Authorization", "Bearer ok") + rec := httptest.NewRecorder() + r.ServeHTTP(rec, req) + + require.Equal(t, http.StatusForbidden, rec.Code) + var resp googleErrorResponse + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + require.Equal(t, http.StatusForbidden, resp.Error.Code) + require.Equal(t, "Insufficient account balance", resp.Error.Message) + require.Equal(t, "PERMISSION_DENIED", resp.Error.Status) +} + func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedOnSuccess(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/server/middleware/api_key_auth_test.go b/backend/internal/server/middleware/api_key_auth_test.go index 25c7db0aac..d5fbc46098 100644 --- a/backend/internal/server/middleware/api_key_auth_test.go +++ b/backend/internal/server/middleware/api_key_auth_test.go @@ -1000,6 +1000,92 @@ func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) { require.Equal(t, 1, touchCalls) } +func TestAPIKeyAuthAllowsBalanceBelowMinimumReserve(t *testing.T) { + gin.SetMode(gin.TestMode) + + user := &service.User{ + ID: 10, + Role: service.RoleUser, + Status: service.StatusActive, + Balance: 0.005, + Concurrency: 3, + } + apiKey := &service.APIKey{ + ID: 103, + UserID: user.ID, + Key: "held-balance-low", + Status: service.StatusActive, + User: user, + } + apiKeyRepo := &stubApiKeyRepo{ + getByKey: func(ctx context.Context, key string) (*service.APIKey, error) { + if key != apiKey.Key { + return nil, service.ErrAPIKeyNotFound + } + clone := *apiKey + userClone := *user + clone.User = &userClone + return &clone, nil + }, + } + + cfg := &config.Config{RunMode: config.RunModeStandard} + cfg.Billing.MinimumBalanceReserve = 0.01 + apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) + router := newAuthTestRouter(apiKeyService, nil, cfg) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/t", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + // 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检, + // 0 < balance < reserve 不得被鉴权中间件硬 403(存量部署静默行为变更)。 + require.Equal(t, http.StatusOK, w.Code) +} + +func TestAPIKeyAuthRejectsExhaustedBalance(t *testing.T) { + gin.SetMode(gin.TestMode) + + user := &service.User{ + ID: 10, + Role: service.RoleUser, + Status: service.StatusActive, + Balance: 0, + Concurrency: 3, + } + apiKey := &service.APIKey{ + ID: 104, + UserID: user.ID, + Key: "held-balance-zero", + Status: service.StatusActive, + User: user, + } + apiKeyRepo := &stubApiKeyRepo{ + getByKey: func(ctx context.Context, key string) (*service.APIKey, error) { + if key != apiKey.Key { + return nil, service.ErrAPIKeyNotFound + } + clone := *apiKey + userClone := *user + clone.User = &userClone + return &clone, nil + }, + } + + cfg := &config.Config{RunMode: config.RunModeStandard} + apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) + router := newAuthTestRouter(apiKeyService, nil, cfg) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/t", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusForbidden, w.Code) + requireAPIKeyAuthError(t, w, "INSUFFICIENT_BALANCE", "Insufficient account balance") +} + func newAuthTestRouter(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) *gin.Engine { router := gin.New() router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, cfg))) diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index a456c8c370..0d7e2a505a 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -545,6 +545,7 @@ func registerSystemRoutes(admin *gin.RouterGroup, h *handler.Handlers) { { system.GET("/version", h.Admin.System.GetVersion) system.GET("/check-updates", h.Admin.System.CheckUpdates) + system.GET("/rollback-versions", h.Admin.System.GetRollbackVersions) system.POST("/update", h.Admin.System.PerformUpdate) system.POST("/rollback", h.Admin.System.Rollback) system.POST("/restart", h.Admin.System.RestartService) diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index 9522578051..ba5b4f61d1 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -121,7 +121,16 @@ func RegisterGatewayRoutes( } h.Gateway.CountTokens(c) }) - gateway.GET("/models", h.Gateway.Models) + // Codex CLI / Codex app refresh their model picker from the provider's + // /models endpoint with a client_version query and expect the ChatGPT + // Codex manifest format; other clients keep the OpenAI-style list. + gateway.GET("/models", func(c *gin.Context) { + if isOpenAIGatewayPlatform(c) && c.Query("client_version") != "" { + h.OpenAIGateway.CodexModels(c) + return + } + h.Gateway.Models(c) + }) gateway.GET("/usage", h.Gateway.Usage) // OpenAI Responses API: auto-route based on group platform gateway.POST("/responses", func(c *gin.Context) { @@ -164,6 +173,16 @@ func RegisterGatewayRoutes( }) gateway.POST("/images/generations", imagesHandler) gateway.POST("/images/edits", imagesHandler) + gateway.POST("/images/batches", h.BatchImage.Submit) + gateway.GET("/images/batches", h.BatchImage.List) + gateway.GET("/images/batches/models", h.BatchImage.Models) + gateway.GET("/images/batches/:id", h.BatchImage.Get) + gateway.GET("/images/batches/:id/items", h.BatchImage.Items) + gateway.GET("/images/batches/:id/items/:custom_id/content", h.BatchImage.ItemContent) + gateway.GET("/images/batches/:id/download", h.BatchImage.Download) + gateway.POST("/images/batches/:id/cancel", h.BatchImage.Cancel) + gateway.DELETE("/images/batches/:id", h.BatchImage.DeleteRecord) + gateway.DELETE("/images/batches/:id/outputs", h.BatchImage.DeleteOutputs) gateway.POST("/videos/generations", videoGenerationHandler) gateway.GET("/videos/:request_id", videoStatusHandler) } @@ -204,6 +223,7 @@ func RegisterGatewayRoutes( codexDirect.GET("/responses", func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) + codexDirect.GET("/models", h.OpenAIGateway.CodexModels) } // OpenAI Chat Completions API(不带v1前缀的别名)— auto-route based on group platform r.POST("/chat/completions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { diff --git a/backend/internal/server/routes/gateway_codex_models_test.go b/backend/internal/server/routes/gateway_codex_models_test.go new file mode 100644 index 0000000000..74af755919 --- /dev/null +++ b/backend/internal/server/routes/gateway_codex_models_test.go @@ -0,0 +1,22 @@ +package routes + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGatewayRoutesCodexModelsManifestPathIsRegistered(t *testing.T) { + router := newGatewayRoutesTestRouter() + + registered := make(map[string]bool) + for _, route := range router.Routes() { + if route.Method == http.MethodGet { + registered[route.Path] = true + } + } + + require.True(t, registered["/backend-api/codex/models"], "GET /backend-api/codex/models should be registered") + require.True(t, registered["/v1/models"], "GET /v1/models should be registered") +} diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index ae25bf387d..d099f93979 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -750,10 +750,19 @@ func resolveRequestedModelInMapping(mapping map[string]string, requestedModel st } // IsModelSupported 检查模型是否在 model_mapping 中(支持通配符) -// 如果未配置 mapping,返回 true(允许所有模型) +// 如果未配置 mapping,返回 true(允许所有模型)。 +// +// 例外:OpenAI OAuth 账号(Codex 上游)的空映射会排除明确属于其他厂商 +// 家族的模型(deepseek-*/glm-* 等)——转发阶段 normalizeOpenAIModelForUpstream +// 会把未知模型原样透传,Codex 上游对这类模型必然返回不可重试的 400,导致 +// 请求卡死在该账号上、无法 failover 到真正支持该模型的 API Key 账号(#3662)。 +// 未知/自定义别名仍保持允许(兼容渠道级映射),见 isOpenAIOAuthServableModel。 func (a *Account) IsModelSupported(requestedModel string) bool { mapping := a.GetModelMapping() if len(mapping) == 0 { + if a.IsOpenAIOAuth() && !a.IsOpenAIPassthroughEnabled() { + return isOpenAIOAuthServableModel(requestedModel) + } return true // 无映射 = 允许所有 } if mappingSupportsRequestedModel(mapping, requestedModel) { diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go new file mode 100644 index 0000000000..52e5ce719b --- /dev/null +++ b/backend/internal/service/admin_account.go @@ -0,0 +1,1075 @@ +package service + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/http" + "strconv" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" +) + +// Account management implementations +func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string, sortBy, sortOrder string) ([]Account, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + accounts, result, err := s.accountRepo.ListWithFilters(ctx, params, platform, accountType, status, search, groupID, privacyMode) + if err != nil { + return nil, 0, err + } + return accounts, result.Total, nil +} + +func (s *adminServiceImpl) ListAccountsForSchedulerScoreFilter(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) { + if s == nil || s.accountRepo == nil { + return nil, nil + } + return s.accountRepo.ListAllWithFilters(ctx, platform, accountType, status, search, groupID, privacyMode) +} + +func (s *adminServiceImpl) ListOpenAISchedulableAccountsForSchedulerScore(ctx context.Context, groupID *int64) ([]Account, error) { + if s == nil || s.accountRepo == nil { + return nil, nil + } + if groupID != nil { + return s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, PlatformOpenAI) + } + return s.accountRepo.ListSchedulableUngroupedByPlatform(ctx, PlatformOpenAI) +} + +func (s *adminServiceImpl) GetAccount(ctx context.Context, id int64) (*Account, error) { + return s.accountRepo.GetByID(ctx, id) +} + +func (s *adminServiceImpl) GetAccountsByIDs(ctx context.Context, ids []int64) ([]*Account, error) { + if len(ids) == 0 { + return []*Account{}, nil + } + + accounts, err := s.accountRepo.GetByIDs(ctx, ids) + if err != nil { + return nil, fmt.Errorf("failed to get accounts by IDs: %w", err) + } + + return accounts, nil +} + +func normalizeAccountConcurrency(platform, accountType string, concurrency int) int { + if platform == PlatformGrok && accountType == AccountTypeOAuth { + if concurrency <= 0 { + return 1 + } + } + return concurrency +} + +func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error) { + // 绑定分组 + groupIDs := input.GroupIDs + // 如果没有指定分组,自动绑定对应平台的默认分组 + if len(groupIDs) == 0 && !input.SkipDefaultGroupBind { + defaultGroupName := input.Platform + "-default" + groups, err := s.groupRepo.ListActiveByPlatform(ctx, input.Platform) + if err == nil { + for _, g := range groups { + if g.Name == defaultGroupName { + groupIDs = []int64{g.ID} + break + } + } + } + } + + // 检查混合渠道风险(除非用户已确认) + if len(groupIDs) > 0 && !input.SkipMixedChannelCheck { + if err := s.checkMixedChannelRisk(ctx, 0, input.Platform, groupIDs); err != nil { + return nil, err + } + } + + // 校验并规范化请求头覆写配置(header 名小写化、格式检查) + if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { + return nil, err + } + + account := &Account{ + Name: input.Name, + Notes: normalizeAccountNotes(input.Notes), + Platform: input.Platform, + Type: input.Type, + Credentials: input.Credentials, + Extra: input.Extra, + ProxyID: input.ProxyID, + Concurrency: normalizeAccountConcurrency(input.Platform, input.Type, input.Concurrency), + Priority: input.Priority, + Status: StatusActive, + Schedulable: true, + } + // 预计算固定时间重置的下次重置时间 + if account.Extra != nil { + if err := ValidateQuotaResetConfig(account.Extra); err != nil { + return nil, err + } + ComputeQuotaResetAt(account.Extra) + NormalizeFixedQuotaWindows(account.Extra) + } + if input.ExpiresAt != nil && *input.ExpiresAt > 0 { + expiresAt := time.Unix(*input.ExpiresAt, 0) + account.ExpiresAt = &expiresAt + } + if input.AutoPauseOnExpired != nil { + account.AutoPauseOnExpired = *input.AutoPauseOnExpired + } else { + account.AutoPauseOnExpired = true + } + if input.RateMultiplier != nil { + if *input.RateMultiplier < 0 { + return nil, errors.New("rate_multiplier must be >= 0") + } + account.RateMultiplier = input.RateMultiplier + } + if input.LoadFactor != nil && *input.LoadFactor > 0 { + if *input.LoadFactor > 10000 { + return nil, errors.New("load_factor must be <= 10000") + } + account.LoadFactor = input.LoadFactor + } + if err := s.accountRepo.Create(ctx, account); err != nil { + return nil, err + } + + // 绑定分组 + if len(groupIDs) > 0 { + if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { + return nil, err + } + } + + // OAuth 账号:创建后异步设置隐私。 + // 使用 Ensure(幂等)而非 Force:新建账号 Extra 为空时效果相同,但更安全。 + if account.Type == AccountTypeOAuth { + switch account.Platform { + case PlatformOpenAI: + go func() { + defer func() { + if r := recover(); r != nil { + slog.Error("create_account_openai_privacy_panic", "account_id", account.ID, "recover", r) + } + }() + s.EnsureOpenAIPrivacy(context.Background(), account) + }() + case PlatformAntigravity: + go func() { + defer func() { + if r := recover(); r != nil { + slog.Error("create_account_antigravity_privacy_panic", "account_id", account.ID, "recover", r) + } + }() + s.EnsureAntigravityPrivacy(context.Background(), account) + }() + } + } + + return account, nil +} + +func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *UpdateAccountInput) (*Account, error) { + account, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + // 安全/身份不变量(影子账号):通用更新路径被 edit/re-auth/refresh/batch 共用, + // 必须在此守住,否则仅在创建时的保证可被这些路径绕过。 + if account.IsCredentialShadow() { + // 影子绝不持有凭据(凭据只在母账号)——外审 F5。 + if !isAllowedSparkShadowCredentialsUpdate(input.Credentials) { + return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_NO_CREDENTIALS", + "spark shadow accounts do not hold auth credentials; only model mapping can be configured on the shadow account") + } + // 影子 type 不可变——很多上游逻辑按 account.Type 分支(OAuth transform / ChatGPT + // header 注入 / WS OAuth 决策),改成 apikey 会让 spark 影子被选中后按错误协议转发(外审 G7)。 + if input.Type != "" && input.Type != account.Type { + return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_IMMUTABLE_TYPE", + "spark shadow account type cannot be changed; it must remain an OpenAI OAuth shadow") + } + } else if input.Type != "" && input.Type != account.Type && input.Type != AccountTypeOAuth { + // 母账号守卫(外审 D/P1):有 spark 影子的账号不能把 type 改出 OpenAI OAuth——影子读透母 + // 凭据,母变成 apikey/setup_token 会让影子被调度后按错协议失败(resolveCredentialAccount + // 必报错)。须先删影子再改 type。 + shadows, serr := s.accountRepo.ListShadowsByParent(ctx, id) + if serr != nil { + return nil, serr + } + if len(shadows) > 0 { + return nil, infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_PARENT_IMMUTABLE_TYPE", + "cannot change account type while it has a spark shadow; delete the shadow first") + } + } + wasOveragesEnabled := account.IsOveragesEnabled() + + if input.Name != "" { + account.Name = input.Name + } + if input.Type != "" { + account.Type = input.Type + } + if input.Notes != nil { + account.Notes = normalizeAccountNotes(input.Notes) + } + if account.IsCredentialShadow() && input.Credentials != nil { + account.Credentials = sanitizeSparkShadowCredentials(input.Credentials) + } else if len(input.Credentials) > 0 { + // 敏感子键采用"incoming 没提供就保留"的合并语义:前端响应已脱敏, + // 全对象 PUT 编辑时不会再带回 token,避免覆盖时清空已有凭证。 + account.Credentials = MergePreservingSensitiveCreds(account.Credentials, input.Credentials) + // 校验并规范化请求头覆写配置(header 名小写化、格式检查) + if err := NormalizeHeaderOverrideCredentials(account.Credentials); err != nil { + return nil, err + } + } + // Extra 使用 map:需要区分“未提供(nil)”与“显式清空({})”。 + // 关闭配额限制时前端会删除 quota_* 键并提交 extra:{},此时也必须落库。 + if input.Extra != nil { + // 保留配额用量字段,防止编辑账号时意外重置 + for _, key := range []string{"quota_used", "quota_daily_used", "quota_daily_start", "quota_weekly_used", "quota_weekly_start"} { + if v, ok := account.Extra[key]; ok { + input.Extra[key] = v + } + } + account.Extra = input.Extra + if account.Platform == PlatformAntigravity && wasOveragesEnabled && !account.IsOveragesEnabled() { + delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 + // 清除 AICredits 限流 key + if rawLimits, ok := account.Extra[modelRateLimitsKey].(map[string]any); ok { + delete(rawLimits, creditsExhaustedKey) + } + } + if account.Platform == PlatformAntigravity && !wasOveragesEnabled && account.IsOveragesEnabled() { + delete(account.Extra, modelRateLimitsKey) + delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 + } + // 校验并预计算固定时间重置的下次重置时间 + if err := ValidateQuotaResetConfig(account.Extra); err != nil { + return nil, err + } + ComputeQuotaResetAt(account.Extra) + NormalizeFixedQuotaWindows(account.Extra) + } + // 影子代理恒继承母账号(由 propagateProxyToShadows 同步),不接受独立编辑——外审 B/P1; + // 否则要等母账号下次改 proxy 才被覆盖,期间影子会出现"有时继承、有时独立"的漂移。 + if input.ProxyID != nil && !account.IsCredentialShadow() { + // 0 表示清除代理(前端发送 0 而不是 null 来表达清除意图) + if *input.ProxyID == 0 { + account.ProxyID = nil + } else { + account.ProxyID = input.ProxyID + } + account.Proxy = nil // 清除关联对象,防止 GORM Save 时根据 Proxy.ID 覆盖 ProxyID + } + // 只在指针非 nil 时更新 Concurrency(支持设置为 0) + if input.Concurrency != nil { + account.Concurrency = normalizeAccountConcurrency(account.Platform, account.Type, *input.Concurrency) + } + // 只在指针非 nil 时更新 Priority(支持设置为 0) + if input.Priority != nil { + account.Priority = *input.Priority + } + if input.RateMultiplier != nil { + if *input.RateMultiplier < 0 { + return nil, errors.New("rate_multiplier must be >= 0") + } + account.RateMultiplier = input.RateMultiplier + } + if input.LoadFactor != nil { + if *input.LoadFactor <= 0 { + account.LoadFactor = nil // 0 或负数表示清除 + } else if *input.LoadFactor > 10000 { + return nil, errors.New("load_factor must be <= 10000") + } else { + account.LoadFactor = input.LoadFactor + } + } + if input.Status != "" { + account.Status = input.Status + } + if input.ExpiresAt != nil { + if *input.ExpiresAt <= 0 { + account.ExpiresAt = nil + } else { + expiresAt := time.Unix(*input.ExpiresAt, 0) + account.ExpiresAt = &expiresAt + } + } + if input.AutoPauseOnExpired != nil { + account.AutoPauseOnExpired = *input.AutoPauseOnExpired + } + + // 先验证分组是否存在(在任何写操作之前) + if input.GroupIDs != nil { + if err := s.validateGroupIDsExist(ctx, *input.GroupIDs); err != nil { + return nil, err + } + + // 检查混合渠道风险(除非用户已确认) + if !input.SkipMixedChannelCheck { + if err := s.checkMixedChannelRisk(ctx, account.ID, account.Platform, *input.GroupIDs); err != nil { + return nil, err + } + } + } + + if err := s.accountRepo.Update(ctx, account); err != nil { + return nil, err + } + + // 将 proxy 变更传播到 spark 影子账号(同步;Update 内部已触发调度快照)。 + // 影子自身 proxy 不可独立编辑(见上),故对影子的更新不触发传播。 + if input.ProxyID != nil && !account.IsCredentialShadow() { + if err := s.propagateProxyToShadows(ctx, id, account.ProxyID); err != nil { + return nil, err + } + } + + // 绑定分组 + if input.GroupIDs != nil { + if err := s.accountRepo.BindGroups(ctx, account.ID, *input.GroupIDs); err != nil { + return nil, err + } + } + + // 重新查询以确保返回完整数据(包括正确的 Proxy 关联对象) + updated, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + return updated, nil +} + +// UpdateAccountExtra 仅对 Extra JSONB 做 key 级合并,避免覆盖其它运行态键 +// (如 model_rate_limits / passive_usage_* 等)。 +func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error { + if len(updates) == 0 { + return nil + } + return s.accountRepo.UpdateExtra(ctx, id, updates) +} + +// BulkUpdateAccounts updates multiple accounts in one request. +// It merges credentials/extra keys instead of overwriting the whole object. +func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) { + if len(input.AccountIDs) == 0 && input.Filters != nil { + accountIDs, err := s.resolveBulkUpdateTargetIDs(ctx, input.Filters) + if err != nil { + return nil, err + } + input.AccountIDs = accountIDs + } + + result := &BulkUpdateAccountsResult{ + SuccessIDs: make([]int64, 0, len(input.AccountIDs)), + FailedIDs: make([]int64, 0, len(input.AccountIDs)), + Results: make([]BulkUpdateAccountResult, 0, len(input.AccountIDs)), + } + + if len(input.AccountIDs) == 0 { + return result, nil + } + if input.GroupIDs != nil { + if err := s.validateGroupIDsExist(ctx, *input.GroupIDs); err != nil { + return nil, err + } + } + + needMixedChannelCheck := input.GroupIDs != nil && !input.SkipMixedChannelCheck + + // 预取所有目标账号,供凭据守卫/代理守卫/混合渠道检查共用,避免多次 DB 查询。 + var cachedTargets []*Account + if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck { + loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) + if err != nil { + return nil, err + } + cachedTargets = loaded + } + + // 影子账号绝不持有凭据:批量更新携带凭据时,目标中不得含影子(外审 G5,与单账号 + // UpdateAccount 守卫对齐)。覆盖显式 IDs 与 filter 解析出的 IDs(此处 AccountIDs 已解析完成)。 + if len(input.Credentials) > 0 { + for _, acc := range cachedTargets { + if acc != nil && acc.IsCredentialShadow() { + return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_NO_CREDENTIALS", + "spark shadow account %d cannot hold credentials; manage credentials on the parent account", acc.ID) + } + } + } + + // 影子账号 proxy 恒继承母账号(与单账号 UpdateAccount 守卫对齐——外审第4轮 P1):批量携带 proxy + // 时目标不得含影子,否则影子会获得独立 proxy、破坏继承不变量(网关按所选影子自身 proxy 出站, + // 要等母账号下次改 proxy 才覆盖→漂移)。含影子即整体拒绝,提示从选择中剔除影子。 + if input.ProxyID != nil { + for _, acc := range cachedTargets { + if acc != nil && acc.IsCredentialShadow() { + return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_PROXY_INHERITED", + "spark shadow account %d proxy is inherited from its parent and cannot be set in bulk; manage it on the parent account", acc.ID) + } + } + } + + // 预加载账号平台信息(混合渠道检查需要)。 + platformByID := map[int64]string{} + if needMixedChannelCheck { + for _, account := range cachedTargets { + if account != nil { + platformByID[account.ID] = account.Platform + } + } + } + + // 预检查混合渠道风险:在任何写操作之前,若发现风险立即返回错误。 + if needMixedChannelCheck { + for _, accountID := range input.AccountIDs { + platform := platformByID[accountID] + if platform == "" { + continue + } + if err := s.checkMixedChannelRisk(ctx, accountID, platform, *input.GroupIDs); err != nil { + return nil, err + } + } + } + + if input.RateMultiplier != nil { + if *input.RateMultiplier < 0 { + return nil, errors.New("rate_multiplier must be >= 0") + } + } + + // 校验并规范化请求头覆写配置(批量路径为 JSONB 顶层 key 合并,直接校验增量即可) + if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { + return nil, err + } + + // Prepare bulk updates for columns and JSONB fields. + repoUpdates := AccountBulkUpdate{ + Credentials: input.Credentials, + Extra: input.Extra, + } + if input.Name != "" { + repoUpdates.Name = &input.Name + } + if input.ProxyID != nil { + repoUpdates.ProxyID = input.ProxyID + } + if input.Concurrency != nil { + repoUpdates.Concurrency = input.Concurrency + } + if input.Priority != nil { + repoUpdates.Priority = input.Priority + } + if input.RateMultiplier != nil { + repoUpdates.RateMultiplier = input.RateMultiplier + } + if input.LoadFactor != nil { + if *input.LoadFactor <= 0 { + repoUpdates.LoadFactor = nil // 0 或负数表示清除 + } else if *input.LoadFactor > 10000 { + return nil, errors.New("load_factor must be <= 10000") + } else { + repoUpdates.LoadFactor = input.LoadFactor + } + } + if input.Status != "" { + repoUpdates.Status = &input.Status + } + if input.Schedulable != nil { + repoUpdates.Schedulable = input.Schedulable + } + + // Run bulk update for column/jsonb fields first. + if _, err := s.accountRepo.BulkUpdate(ctx, input.AccountIDs, repoUpdates); err != nil { + return nil, err + } + + // 将 proxy 变更传播到每个目标账号的 spark 影子账号 + if repoUpdates.ProxyID != nil { + var effectiveProxyID *int64 + if *repoUpdates.ProxyID != 0 { + effectiveProxyID = repoUpdates.ProxyID + } + for _, accountID := range input.AccountIDs { + if err := s.propagateProxyToShadows(ctx, accountID, effectiveProxyID); err != nil { + return nil, err + } + } + } + + // Handle group bindings per account (requires individual operations). + for _, accountID := range input.AccountIDs { + entry := BulkUpdateAccountResult{AccountID: accountID} + + if input.GroupIDs != nil { + if err := s.accountRepo.BindGroups(ctx, accountID, *input.GroupIDs); err != nil { + entry.Success = false + entry.Error = err.Error() + result.Failed++ + result.FailedIDs = append(result.FailedIDs, accountID) + result.Results = append(result.Results, entry) + continue + } + } + + entry.Success = true + result.Success++ + result.SuccessIDs = append(result.SuccessIDs, accountID) + result.Results = append(result.Results, entry) + } + + return result, nil +} + +func (s *adminServiceImpl) resolveBulkUpdateTargetIDs(ctx context.Context, filters *BulkUpdateAccountFilters) ([]int64, error) { + if filters == nil { + return nil, nil + } + + groupID := int64(0) + switch strings.TrimSpace(filters.Group) { + case "": + case "ungrouped": + groupID = AccountListGroupUngrouped + default: + parsedGroupID, err := strconv.ParseInt(strings.TrimSpace(filters.Group), 10, 64) + if err != nil { + return nil, fmt.Errorf("invalid group filter: %w", err) + } + groupID = parsedGroupID + } + + const pageSize = 500 + page := 1 + accountIDs := make([]int64, 0, pageSize) + + for { + accounts, total, err := s.ListAccounts( + ctx, + page, + pageSize, + filters.Platform, + filters.Type, + filters.Status, + filters.Search, + groupID, + filters.PrivacyMode, + "", + "", + ) + if err != nil { + return nil, err + } + for _, account := range accounts { + accountIDs = append(accountIDs, account.ID) + } + if int64(len(accountIDs)) >= total || len(accounts) == 0 { + return accountIDs, nil + } + page++ + } +} + +func (s *adminServiceImpl) DeleteAccount(ctx context.Context, id int64) error { + // 级联删除 spark 影子账号(先删影子,再删母账号) + shadows, err := s.accountRepo.ListShadowsByParent(ctx, id) + if err != nil { + return fmt.Errorf("list spark shadows for cascade delete: %w", err) + } + for _, shadow := range shadows { + if err := s.accountRepo.Delete(ctx, shadow.ID); err != nil { + return fmt.Errorf("cascade delete spark shadow %d: %w", shadow.ID, err) + } + } + if err := s.accountRepo.Delete(ctx, id); err != nil { + return err + } + return nil +} + +func (s *adminServiceImpl) RefreshAccountCredentials(ctx context.Context, id int64) (*Account, error) { + account, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + // TODO: Implement refresh logic + return account, nil +} + +func (s *adminServiceImpl) ClearAccountError(ctx context.Context, id int64) (*Account, error) { + if err := s.accountRepo.ClearError(ctx, id); err != nil { + return nil, err + } + if err := s.accountRepo.ClearRateLimit(ctx, id); err != nil { + return nil, err + } + if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, id); err != nil { + return nil, err + } + if err := s.accountRepo.ClearModelRateLimits(ctx, id); err != nil { + return nil, err + } + if err := s.accountRepo.ClearTempUnschedulable(ctx, id); err != nil { + return nil, err + } + if s.runtimeBlocker != nil { + s.runtimeBlocker.ClearAccountSchedulingBlock(id) + } + return s.accountRepo.GetByID(ctx, id) +} + +func (s *adminServiceImpl) SetAccountError(ctx context.Context, id int64, errorMsg string) error { + return s.accountRepo.SetError(ctx, id, errorMsg) +} + +func (s *adminServiceImpl) SetAccountSchedulable(ctx context.Context, id int64, schedulable bool) (*Account, error) { + if err := s.accountRepo.SetSchedulable(ctx, id, schedulable); err != nil { + return nil, err + } + updated, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + return updated, nil +} + +func (s *adminServiceImpl) RevertAccountProxyFallback(ctx context.Context, id int64) error { + if err := s.accountRepo.RevertProxyFallback(ctx, id); err != nil { + return err + } + // 加载回退后的账号以获取实际 ProxyID,再传播到影子账号 + account, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return fmt.Errorf("get account after proxy revert: %w", err) + } + return s.propagateProxyToShadows(ctx, id, account.ProxyID) +} + +// CreateShadow 为指定 OpenAI OAuth 母账号创建 spark 维度影子账号(一母一影)。 +// 安全不变量:Credentials 恒不含 auth token(仅 model_mapping,守卫 isAllowedSparkShadowCredentialsUpdate 放行)。 +func (s *adminServiceImpl) CreateShadow(ctx context.Context, parentID int64, opts ShadowOptions) (*Account, error) { + // 1. 加载母账号并校验平台/类型 + parent, err := s.accountRepo.GetByID(ctx, parentID) + if err != nil { + return nil, fmt.Errorf("get parent account: %w", err) + } + if !parent.IsOpenAIOAuth() { + return nil, infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_INVALID_PARENT", + "spark shadow requires an OpenAI OAuth parent account") + } + // G6:母账号本身不能是影子,否则会建出二级影子——resolveCredentialAccount 只解一层, + // 会解析到无凭据的一级影子,进入坏调度/上游失败。 + if parent.IsCredentialShadow() { + return nil, infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_PARENT_IS_SHADOW", + "spark shadow parent must be a real account, not another spark shadow") + } + + // 2. 一母一影校验 + shadows, err := s.accountRepo.ListShadowsByParent(ctx, parentID) + if err != nil { + return nil, fmt.Errorf("check existing spark shadows: %w", err) + } + if len(shadows) > 0 { + return nil, infraerrors.New(http.StatusConflict, "SPARK_SHADOW_ALREADY_EXISTS", + "parent account already has a spark shadow account") + } + + // 3. 解析分组。未指定 GroupIDs 时:优先**继承母账号当前分组**(影子与母同路由域,母在自定义 + // 组时该组的 spark 请求也能选到影子;G1 决策);母无分组再回落 openai-default(F4)。 + // 显式指定 GroupIDs 时,与 UpdateAccount 对齐先校验存在性(创建前),避免建出影子后再因无效组 + // 失败而留下孤儿影子(一母一影唯一索引会挡住重试)——外审 C/P1。 + groupIDs := opts.GroupIDs + if len(groupIDs) > 0 { + if s.groupRepo != nil { + if err := s.validateGroupIDsExist(ctx, groupIDs); err != nil { + return nil, err + } + } + } else if len(parent.GroupIDs) > 0 { + groupIDs = append([]int64(nil), parent.GroupIDs...) + } else if s.groupRepo != nil { + defaultGroupName := PlatformOpenAI + "-default" + if groups, gerr := s.groupRepo.ListActiveByPlatform(ctx, PlatformOpenAI); gerr == nil { + for _, g := range groups { + if g.Name == defaultGroupName { + groupIDs = []int64{g.ID} + break + } + } + } + } + + // 4. 构造影子账号(安全不变量:Credentials 恒不含 auth token,仅含 model_mapping)。 + // name 为空时默认 "<母账号名> (Spark)"——否则空 name 会在 ent(name NotEmpty)处变成裸 500 + // (外审 E/P2);并 rune 安全截断到 ent MaxLen(100)。 + name := strings.TrimSpace(opts.Name) + if name == "" { + name = parent.Name + " (Spark)" + } + if runes := []rune(name); len(runes) > 100 { + name = string(runes[:100]) + } + // 并发未指定(<=0)时继承母账号,避免 0 被限流器解读为"无限并发"(外审 F3)。 + concurrency := opts.Concurrency + if concurrency <= 0 { + concurrency = parent.Concurrency + } + // 优先级未指定(<=0)时继承母账号——前端一键创建只传 name,opts.Priority 省略即 0,而调度 + // 比较是「数值越小越优先」(openai_account_scheduler.isOpenAIAccountCandidateBetter),且 repo + // 显式 SetPriority 会绕过 ent 默认 50,直写 0 会让影子意外抢到最高优先级(外审第5轮 P1)。 + // 与上方 Concurrency 一致采用「省略继承母账号」语义(影子的 proxy/分组/并发亦全部继承母账号)。 + priority := opts.Priority + if priority <= 0 { + priority = parent.Priority + } + shadow := &Account{ + Name: name, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{"model_mapping": defaultSparkShadowModelMapping()}, + ParentAccountID: &parentID, + QuotaDimension: QuotaDimensionSpark, + ProxyID: parent.ProxyID, + Priority: priority, + Concurrency: concurrency, + Schedulable: true, + } + + // 5. 持久化(Create 填充 shadow.ID)。并发竞态:预查(步骤2)放行后另一请求抢先建成,本次会撞 + // 一母一影唯一索引。复查确认确为"已存在"竞态时返回结构化 409 而非裸 500——外审 A/P1。 + if err := s.accountRepo.Create(ctx, shadow); err != nil { + if existing, qerr := s.accountRepo.ListShadowsByParent(ctx, parentID); qerr == nil && len(existing) > 0 { + return nil, infraerrors.New(http.StatusConflict, "SPARK_SHADOW_ALREADY_EXISTS", + "parent account already has a spark shadow account") + } + return nil, fmt.Errorf("create spark shadow: %w", err) + } + + // 6. 绑定分组。注意:create+bind 非单一 DB 事务(通用 Create 走 r.client、outbox 走 r.sql, + // 无现成共享事务路径),故绑组失败时做 best-effort 补偿删除刚建的影子,避免半成品影子(否则 + // 一母一影唯一索引会挡住重试)——外审 C/P1。补偿删除用 detached ctx,即便请求 ctx 已取消/超时 + // 仍能完成清理(外审第4轮);进程崩溃这种极端仍可能残留,属已知权衡。 + if len(groupIDs) > 0 { + if err := s.accountRepo.BindGroups(ctx, shadow.ID, groupIDs); err != nil { + if delErr := s.accountRepo.Delete(context.WithoutCancel(ctx), shadow.ID); delErr != nil { + slog.Error("spark_shadow_bind_groups_rollback_failed", + "shadow_id", shadow.ID, "parent_id", parentID, "delete_err", delErr) + } + return nil, fmt.Errorf("bind groups for spark shadow: %w", err) + } + shadow.GroupIDs = groupIDs + } + + return shadow, nil +} + +// propagateProxyToShadows syncs proxyID to all spark shadow accounts of parentID. +// It is called synchronously so that proxy changes are immediately consistent; +// accountRepo.Update triggers the scheduler outbox + cache propagation internally. +// Calling this for a non-parent account is a harmless no-op. +func (s *adminServiceImpl) propagateProxyToShadows(ctx context.Context, parentID int64, proxyID *int64) error { + return propagateAccountProxyToShadows(ctx, s.accountRepo, parentID, proxyID) +} + +// propagateAccountProxyToShadows 把母账号的 proxy 同步到其所有 spark 影子(影子 proxy 恒继承母账号)。 +// 供 AdminService 编辑路径与 CRS 同步路径共用——后者改动母账号 proxy 后必须同样传播,否则影子保留 +// 旧 proxy 出现出站漂移(外审第8轮)。 +func propagateAccountProxyToShadows(ctx context.Context, repo AccountRepository, parentID int64, proxyID *int64) error { + shadows, err := repo.ListShadowsByParent(ctx, parentID) + if err != nil { + return fmt.Errorf("list spark shadows for proxy propagation: %w", err) + } + for _, shadow := range shadows { + shadow.ProxyID = proxyID + if err := repo.Update(ctx, shadow); err != nil { + return fmt.Errorf("update spark shadow %d proxy: %w", shadow.ID, err) + } + } + return nil +} + +// checkMixedChannelRisk 检查分组中是否存在混合渠道(Antigravity + Anthropic) +// 如果存在混合,返回错误提示用户确认 +func (s *adminServiceImpl) checkMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error { + // 判断当前账号的渠道类型(基于 platform 字段,而不是 type 字段) + currentPlatform := getAccountPlatform(currentAccountPlatform) + if currentPlatform == "" { + // 不是 Antigravity 或 Anthropic,无需检查 + return nil + } + + // 检查每个分组中的其他账号 + for _, groupID := range groupIDs { + accounts, err := s.accountRepo.ListByGroup(ctx, groupID) + if err != nil { + return fmt.Errorf("get accounts in group %d: %w", groupID, err) + } + + // 检查是否存在不同渠道的账号 + for _, account := range accounts { + if currentAccountID > 0 && account.ID == currentAccountID { + continue // 跳过当前账号 + } + + otherPlatform := getAccountPlatform(account.Platform) + if otherPlatform == "" { + continue // 不是 Antigravity 或 Anthropic,跳过 + } + + // 检测混合渠道 + if currentPlatform != otherPlatform { + group, _ := s.groupRepo.GetByID(ctx, groupID) + groupName := fmt.Sprintf("Group %d", groupID) + if group != nil { + groupName = group.Name + } + + return &MixedChannelError{ + GroupID: groupID, + GroupName: groupName, + CurrentPlatform: currentPlatform, + OtherPlatform: otherPlatform, + } + } + } + } + + return nil +} + +func (s *adminServiceImpl) validateGroupIDsExist(ctx context.Context, groupIDs []int64) error { + if len(groupIDs) == 0 { + return nil + } + if s.groupRepo == nil { + return errors.New("group repository not configured") + } + + if batchReader, ok := s.groupRepo.(groupExistenceBatchReader); ok { + existsByID, err := batchReader.ExistsByIDs(ctx, groupIDs) + if err != nil { + return fmt.Errorf("check groups exists: %w", err) + } + for _, groupID := range groupIDs { + if groupID <= 0 || !existsByID[groupID] { + return fmt.Errorf("get group: %w", ErrGroupNotFound) + } + } + return nil + } + + for _, groupID := range groupIDs { + if _, err := s.groupRepo.GetByID(ctx, groupID); err != nil { + return fmt.Errorf("get group: %w", err) + } + } + return nil +} + +// CheckMixedChannelRisk checks whether target groups contain mixed channels for the current account platform. +func (s *adminServiceImpl) CheckMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error { + return s.checkMixedChannelRisk(ctx, currentAccountID, currentAccountPlatform, groupIDs) +} + +// getAccountPlatform 根据账号 platform 判断混合渠道检查用的平台标识 +func getAccountPlatform(accountPlatform string) string { + switch strings.ToLower(strings.TrimSpace(accountPlatform)) { + case PlatformAntigravity: + return "Antigravity" + case PlatformAnthropic, "claude": + return "Anthropic" + default: + return "" + } +} + +// MixedChannelError 混合渠道错误 +type MixedChannelError struct { + GroupID int64 + GroupName string + CurrentPlatform string + OtherPlatform string +} + +func (e *MixedChannelError) Error() string { + return fmt.Sprintf("mixed_channel_warning: Group '%s' contains both %s and %s accounts. Using mixed channels in the same context may cause thinking block signature validation issues, which will fallback to non-thinking mode for historical messages.", + e.GroupName, e.CurrentPlatform, e.OtherPlatform) +} + +func (s *adminServiceImpl) ResetAccountQuota(ctx context.Context, id int64) error { + account, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return err + } + // spark 影子账号不持自有配额(凭据透传母账号、spark 用量走独立 codex_* 维度由 QueryUsage 维护), + // 通用 quota 重置对其无意义且语义不一致——明确 400 拒绝(与 OpenAI reset-credit 对影子一致)(外审第7轮 P2)。 + if account.IsCredentialShadow() { + return infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_NO_QUOTA_RESET", + "cannot reset quota for a spark shadow account; manage it on the parent account") + } + return s.accountRepo.ResetQuotaUsed(ctx, id) +} + +// EnsureOpenAIPrivacy 检查 OpenAI OAuth 账号是否已设置 privacy_mode, +// 未设置则调用 disableOpenAITraining 并持久化到 Extra,返回设置的 mode 值。 +func (s *adminServiceImpl) EnsureOpenAIPrivacy(ctx context.Context, account *Account) string { + // 影子账号不持凭据,隐私设置由母账号管理,直接跳过。 + if account.IsCredentialShadow() { + return "" + } + if account.Platform != PlatformOpenAI || account.Type != AccountTypeOAuth { + return "" + } + if s.privacyClientFactory == nil { + return "" + } + if shouldSkipOpenAIPrivacyEnsure(account.Extra) { + return "" + } + + token, _ := account.Credentials["access_token"].(string) + if token == "" { + return "" + } + + var proxyURL string + if account.ProxyID != nil { + if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { + proxyURL = p.URL() + } + } + + mode := disableOpenAITraining(ctx, s.privacyClientFactory, token, proxyURL) + if mode == "" { + return "" + } + + _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}) + return mode +} + +// ForceOpenAIPrivacy 强制重新设置 OpenAI OAuth 账号隐私,无论当前状态。 +func (s *adminServiceImpl) ForceOpenAIPrivacy(ctx context.Context, account *Account) string { + // 影子账号不持凭据,隐私由母账号管理,直接跳过(与 EnsureOpenAIPrivacy 一致——外审第4轮)。 + if account.IsCredentialShadow() { + return "" + } + if account.Platform != PlatformOpenAI || account.Type != AccountTypeOAuth { + return "" + } + if s.privacyClientFactory == nil { + return "" + } + + token, _ := account.Credentials["access_token"].(string) + if token == "" { + return "" + } + + var proxyURL string + if account.ProxyID != nil { + if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { + proxyURL = p.URL() + } + } + + mode := disableOpenAITraining(ctx, s.privacyClientFactory, token, proxyURL) + if mode == "" { + return "" + } + + if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil { + logger.LegacyPrintf("service.admin", "force_update_openai_privacy_mode_failed: account_id=%d err=%v", account.ID, err) + return mode + } + if account.Extra == nil { + account.Extra = make(map[string]any) + } + account.Extra["privacy_mode"] = mode + return mode +} + +// EnsureAntigravityPrivacy 检查 Antigravity OAuth 账号隐私状态。 +// 仅当 privacy_mode 已成功设置("privacy_set")时跳过; +// 未设置或之前失败("privacy_set_failed")均会重试。 +func (s *adminServiceImpl) EnsureAntigravityPrivacy(ctx context.Context, account *Account) string { + if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth { + return "" + } + if account.Extra != nil { + if existing, ok := account.Extra["privacy_mode"].(string); ok && existing == AntigravityPrivacySet { + return existing + } + } + + token, _ := account.Credentials["access_token"].(string) + if token == "" { + return "" + } + + projectID, _ := account.Credentials["project_id"].(string) + + var proxyURL string + if account.ProxyID != nil { + if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { + proxyURL = p.URL() + } + } + + mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL) + if mode == "" { + return "" + } + + if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil { + logger.LegacyPrintf("service.admin", "update_antigravity_privacy_mode_failed: account_id=%d err=%v", account.ID, err) + return mode + } + applyAntigravityPrivacyMode(account, mode) + return mode +} + +// ForceAntigravityPrivacy 强制重新设置 Antigravity OAuth 账号隐私,无论当前状态。 +func (s *adminServiceImpl) ForceAntigravityPrivacy(ctx context.Context, account *Account) string { + if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth { + return "" + } + + token, _ := account.Credentials["access_token"].(string) + if token == "" { + return "" + } + + projectID, _ := account.Credentials["project_id"].(string) + + var proxyURL string + if account.ProxyID != nil { + if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { + proxyURL = p.URL() + } + } + + mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL) + if mode == "" { + return "" + } + + if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil { + logger.LegacyPrintf("service.admin", "force_update_antigravity_privacy_mode_failed: account_id=%d err=%v", account.ID, err) + return mode + } + applyAntigravityPrivacyMode(account, mode) + return mode +} diff --git a/backend/internal/service/admin_group.go b/backend/internal/service/admin_group.go new file mode 100644 index 0000000000..c85056d623 --- /dev/null +++ b/backend/internal/service/admin_group.go @@ -0,0 +1,998 @@ +package service + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" +) + +// Group management implementations +func (s *adminServiceImpl) ListGroups(ctx context.Context, page, pageSize int, platform, status, search string, isExclusive *bool, sortBy, sortOrder string) ([]Group, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + groups, result, err := s.groupRepo.ListWithFilters(ctx, params, platform, status, search, isExclusive) + if err != nil { + return nil, 0, err + } + return groups, result.Total, nil +} + +func (s *adminServiceImpl) GetAllGroups(ctx context.Context) ([]Group, error) { + return s.groupRepo.ListActive(ctx) +} + +func (s *adminServiceImpl) GetAllGroupsByPlatform(ctx context.Context, platform string) ([]Group, error) { + return s.groupRepo.ListActiveByPlatform(ctx, platform) +} + +func (s *adminServiceImpl) GetAllGroupsIncludingInactive(ctx context.Context) ([]Group, error) { + // ListWithFilters with empty status = no status filter, so active + disabled groups are returned. + // PageSize 10000 is intentionally large; group count is O(dozens) in practice. + groups, _, err := s.groupRepo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10000}, "", "", "", nil) + return groups, err +} + +func (s *adminServiceImpl) GetGroup(ctx context.Context, id int64) (*Group, error) { + return s.groupRepo.GetByID(ctx, id) +} + +func (s *adminServiceImpl) GetGroupModelsListCandidates(ctx context.Context, id int64, platform string) ([]string, error) { + platform = strings.TrimSpace(platform) + if id > 0 { + group, err := s.groupRepo.GetByIDLite(ctx, id) + if err != nil { + return nil, err + } + if platform == "" { + platform = group.Platform + } + } + if platform == "" { + platform = PlatformAnthropic + } + + candidates := defaultModelsListCandidateIDs(platform) + if id <= 0 || s.accountRepo == nil { + return candidates, nil + } + + accounts, err := s.accountRepo.ListSchedulableByGroupID(ctx, id) + if err != nil { + return nil, err + } + + seen := make(map[string]struct{}, len(candidates)) + for _, model := range candidates { + seen[model] = struct{}{} + } + for _, acc := range accounts { + if acc.Platform != platform { + continue + } + for model := range acc.GetModelMapping() { + model = strings.TrimSpace(model) + if model == "" { + continue + } + if _, ok := seen[model]; ok { + continue + } + seen[model] = struct{}{} + candidates = append(candidates, model) + } + } + return candidates, nil +} + +func defaultModelsListCandidateIDs(platform string) []string { + switch platform { + case PlatformOpenAI: + return openai.DefaultModelIDs() + case PlatformGemini: + ids := make([]string, 0, len(geminicli.DefaultModels)) + for _, model := range geminicli.DefaultModels { + ids = append(ids, model.ID) + } + return ids + case PlatformAntigravity: + models := antigravity.DefaultModels() + ids := make([]string, 0, len(models)) + for _, model := range models { + ids = append(ids, model.ID) + } + return ids + case PlatformGrok: + return xai.DefaultModelIDs() + default: + ids := make([]string, 0, len(claude.DefaultModels)) + for _, model := range claude.DefaultModels { + ids = append(ids, model.ID) + } + return ids + } +} + +func defaultAllowImageGenerationForPlatform(platform string) bool { + // Grok image and video generation routes share the legacy image-generation gate. + // Older clients send the false zero value, so Grok groups must default enabled. + return platform == PlatformGrok +} + +func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupInput) (*Group, error) { + if input.RateMultiplier <= 0 { + return nil, errors.New("rate_multiplier must be > 0") + } + + platform := input.Platform + if platform == "" { + platform = PlatformAnthropic + } + + subscriptionType := input.SubscriptionType + if subscriptionType == "" { + subscriptionType = SubscriptionTypeStandard + } + + // 限额字段:nil/负数 表示"无限制",0 表示"不允许用量",正数表示具体限额 + dailyLimit := normalizeLimit(input.DailyLimitUSD) + weeklyLimit := normalizeLimit(input.WeeklyLimitUSD) + monthlyLimit := normalizeLimit(input.MonthlyLimitUSD) + + // 图片价格:负数表示清除(使用默认价格),0 保留(表示免费) + imagePrice1K := normalizePrice(input.ImagePrice1K) + imagePrice2K := normalizePrice(input.ImagePrice2K) + imagePrice4K := normalizePrice(input.ImagePrice4K) + videoPrice480P := normalizePrice(input.VideoPrice480P) + videoPrice720P := normalizePrice(input.VideoPrice720P) + videoPrice1080P := normalizePrice(input.VideoPrice1080P) + imageRateMultiplier := 1.0 + if input.ImageRateMultiplier != nil { + if *input.ImageRateMultiplier < 0 { + return nil, errors.New("image_rate_multiplier must be >= 0") + } + imageRateMultiplier = *input.ImageRateMultiplier + } + batchImageDiscountMultiplier := defaultBatchImageDiscountMultiplier + if input.BatchImageDiscountMultiplier != nil { + if *input.BatchImageDiscountMultiplier < 0 { + return nil, errors.New("batch_image_discount_multiplier must be >= 0") + } + batchImageDiscountMultiplier = *input.BatchImageDiscountMultiplier + } + batchImageHoldMultiplier := defaultBatchImageHoldMultiplier + if input.BatchImageHoldMultiplier != nil { + if *input.BatchImageHoldMultiplier < 0 { + return nil, errors.New("batch_image_hold_multiplier must be >= 0") + } + batchImageHoldMultiplier = *input.BatchImageHoldMultiplier + } + // 不变式:hold 比例 >= discount 比例。否则批量任务成功率足够高时 + // 实际成本会超过冻结额,结算永远失败、用户冻结余额无法解冻。 + if batchImageHoldMultiplier < batchImageDiscountMultiplier { + return nil, errors.New("batch_image_hold_multiplier must be >= batch_image_discount_multiplier") + } + videoRateMultiplier := 1.0 + if input.VideoRateMultiplier != nil { + if *input.VideoRateMultiplier < 0 { + return nil, errors.New("video_rate_multiplier must be >= 0") + } + videoRateMultiplier = *input.VideoRateMultiplier + } + + peakRateMultiplier := 1.0 + if input.PeakRateMultiplier != nil { + peakRateMultiplier = *input.PeakRateMultiplier + } + // 先归一化(非订阅分组清空高峰配置、清洗停用状态下的脏字段)再校验,与 UpdateGroup 同一收口。 + peakRateEnabled, peakStart, peakEnd, peakRateMultiplier := NormalizePeakRateConfig(subscriptionType, input.PeakRateEnabled, input.PeakStart, input.PeakEnd, peakRateMultiplier) + if err := ValidatePeakRateConfig(subscriptionType, peakRateEnabled, peakStart, peakEnd, peakRateMultiplier); err != nil { + return nil, err + } + + // 校验降级分组 + if input.FallbackGroupID != nil { + if err := s.validateFallbackGroup(ctx, 0, *input.FallbackGroupID); err != nil { + return nil, err + } + } + fallbackOnInvalidRequest := input.FallbackGroupIDOnInvalidRequest + if fallbackOnInvalidRequest != nil && *fallbackOnInvalidRequest <= 0 { + fallbackOnInvalidRequest = nil + } + // 校验无效请求兜底分组 + if fallbackOnInvalidRequest != nil { + if err := s.validateFallbackGroupOnInvalidRequest(ctx, 0, platform, subscriptionType, *fallbackOnInvalidRequest); err != nil { + return nil, err + } + } + + // MCPXMLInject:默认为 true,仅当显式传入 false 时关闭 + mcpXMLInject := true + if input.MCPXMLInject != nil { + mcpXMLInject = *input.MCPXMLInject + } + + allowImageGeneration := input.AllowImageGeneration || defaultAllowImageGenerationForPlatform(platform) + allowBatchImageGeneration := input.AllowBatchImageGeneration && allowImageGeneration && platform == PlatformGemini + + // 如果指定了复制账号的源分组,先获取账号 ID 列表 + var accountIDsToCopy []int64 + if len(input.CopyAccountsFromGroupIDs) > 0 { + // 去重源分组 IDs + seen := make(map[int64]struct{}) + uniqueSourceGroupIDs := make([]int64, 0, len(input.CopyAccountsFromGroupIDs)) + for _, srcGroupID := range input.CopyAccountsFromGroupIDs { + if _, exists := seen[srcGroupID]; !exists { + seen[srcGroupID] = struct{}{} + uniqueSourceGroupIDs = append(uniqueSourceGroupIDs, srcGroupID) + } + } + + // 校验源分组的平台是否与新分组一致 + for _, srcGroupID := range uniqueSourceGroupIDs { + srcGroup, err := s.groupRepo.GetByIDLite(ctx, srcGroupID) + if err != nil { + return nil, fmt.Errorf("source group %d not found: %w", srcGroupID, err) + } + if srcGroup.Platform != platform { + return nil, fmt.Errorf("source group %d platform mismatch: expected %s, got %s", srcGroupID, platform, srcGroup.Platform) + } + } + + // 获取所有源分组的账号(去重) + var err error + accountIDsToCopy, err = s.groupRepo.GetAccountIDsByGroupIDs(ctx, uniqueSourceGroupIDs) + if err != nil { + return nil, fmt.Errorf("failed to get accounts from source groups: %w", err) + } + } + + group := &Group{ + Name: input.Name, + Description: input.Description, + Platform: platform, + RateMultiplier: input.RateMultiplier, + IsExclusive: input.IsExclusive, + Status: StatusActive, + SubscriptionType: subscriptionType, + DailyLimitUSD: dailyLimit, + WeeklyLimitUSD: weeklyLimit, + MonthlyLimitUSD: monthlyLimit, + AllowImageGeneration: allowImageGeneration, + AllowBatchImageGeneration: allowBatchImageGeneration, + ImageRateIndependent: input.ImageRateIndependent, + ImageRateMultiplier: imageRateMultiplier, + BatchImageDiscountMultiplier: batchImageDiscountMultiplier, + BatchImageHoldMultiplier: batchImageHoldMultiplier, + VideoRateIndependent: input.VideoRateIndependent, + VideoRateMultiplier: videoRateMultiplier, + PeakRateEnabled: peakRateEnabled, + PeakStart: peakStart, + PeakEnd: peakEnd, + PeakRateMultiplier: peakRateMultiplier, + ImagePrice1K: imagePrice1K, + ImagePrice2K: imagePrice2K, + ImagePrice4K: imagePrice4K, + VideoPrice480P: videoPrice480P, + VideoPrice720P: videoPrice720P, + VideoPrice1080P: videoPrice1080P, + ClaudeCodeOnly: input.ClaudeCodeOnly, + FallbackGroupID: input.FallbackGroupID, + FallbackGroupIDOnInvalidRequest: fallbackOnInvalidRequest, + ModelRouting: input.ModelRouting, + MCPXMLInject: mcpXMLInject, + SupportedModelScopes: input.SupportedModelScopes, + AllowMessagesDispatch: input.AllowMessagesDispatch, + RequireOAuthOnly: input.RequireOAuthOnly, + RequirePrivacySet: input.RequirePrivacySet, + DefaultMappedModel: input.DefaultMappedModel, + MessagesDispatchModelConfig: normalizeOpenAIMessagesDispatchModelConfig(input.MessagesDispatchModelConfig), + ModelsListConfig: normalizeGroupModelsListConfig(input.ModelsListConfig), + RPMLimit: input.RPMLimit, + } + sanitizeGroupMessagesDispatchFields(group) + if err := s.groupRepo.Create(ctx, group); err != nil { + return nil, err + } + + // require_oauth_only: 过滤掉 apikey 类型账号 + if group.RequireOAuthOnly && (group.Platform == PlatformOpenAI || group.Platform == PlatformAntigravity || group.Platform == PlatformAnthropic || group.Platform == PlatformGemini || group.Platform == PlatformGrok) && len(accountIDsToCopy) > 0 { + accounts, err := s.accountRepo.GetByIDs(ctx, accountIDsToCopy) + if err != nil { + return nil, fmt.Errorf("failed to fetch accounts for oauth filter: %w", err) + } + oauthIDs := make(map[int64]struct{}, len(accounts)) + for _, acc := range accounts { + if acc.Type != AccountTypeAPIKey { + oauthIDs[acc.ID] = struct{}{} + } + } + var filtered []int64 + for _, aid := range accountIDsToCopy { + if _, ok := oauthIDs[aid]; ok { + filtered = append(filtered, aid) + } + } + accountIDsToCopy = filtered + } + + // 如果有需要复制的账号,绑定到新分组 + if len(accountIDsToCopy) > 0 { + if err := s.groupRepo.BindAccountsToGroup(ctx, group.ID, accountIDsToCopy); err != nil { + return nil, fmt.Errorf("failed to bind accounts to new group: %w", err) + } + group.AccountCount = int64(len(accountIDsToCopy)) + } + + return group, nil +} + +// normalizeLimit 将负数转换为 nil(表示无限制),0 保留(表示限额为零) +func normalizeLimit(limit *float64) *float64 { + if limit == nil || *limit < 0 { + return nil + } + return limit +} + +// normalizePrice 将负数转换为 nil(表示使用默认价格),0 保留(表示免费) +func normalizePrice(price *float64) *float64 { + if price == nil || *price < 0 { + return nil + } + return price +} + +// validateFallbackGroup 校验降级分组的有效性 +// currentGroupID: 当前分组 ID(新建时为 0) +// fallbackGroupID: 降级分组 ID +func (s *adminServiceImpl) validateFallbackGroup(ctx context.Context, currentGroupID, fallbackGroupID int64) error { + // 不能将自己设置为降级分组 + if currentGroupID > 0 && currentGroupID == fallbackGroupID { + return fmt.Errorf("cannot set self as fallback group") + } + + visited := map[int64]struct{}{} + nextID := fallbackGroupID + for { + if _, seen := visited[nextID]; seen { + return fmt.Errorf("fallback group cycle detected") + } + visited[nextID] = struct{}{} + if currentGroupID > 0 && nextID == currentGroupID { + return fmt.Errorf("fallback group cycle detected") + } + + // 检查降级分组是否存在 + fallbackGroup, err := s.groupRepo.GetByIDLite(ctx, nextID) + if err != nil { + return fmt.Errorf("fallback group not found: %w", err) + } + + // 降级分组不能启用 claude_code_only,否则会造成死循环 + if nextID == fallbackGroupID && fallbackGroup.ClaudeCodeOnly { + return fmt.Errorf("fallback group cannot have claude_code_only enabled") + } + + if fallbackGroup.FallbackGroupID == nil { + return nil + } + nextID = *fallbackGroup.FallbackGroupID + } +} + +// validateFallbackGroupOnInvalidRequest 校验无效请求兜底分组的有效性 +// currentGroupID: 当前分组 ID(新建时为 0) +// platform/subscriptionType: 当前分组的有效平台/订阅类型 +// fallbackGroupID: 兜底分组 ID +func (s *adminServiceImpl) validateFallbackGroupOnInvalidRequest(ctx context.Context, currentGroupID int64, platform, subscriptionType string, fallbackGroupID int64) error { + if platform != PlatformAnthropic && platform != PlatformAntigravity { + return fmt.Errorf("invalid request fallback only supported for anthropic or antigravity groups") + } + if subscriptionType == SubscriptionTypeSubscription { + return fmt.Errorf("subscription groups cannot set invalid request fallback") + } + if currentGroupID > 0 && currentGroupID == fallbackGroupID { + return fmt.Errorf("cannot set self as invalid request fallback group") + } + + fallbackGroup, err := s.groupRepo.GetByIDLite(ctx, fallbackGroupID) + if err != nil { + return fmt.Errorf("fallback group not found: %w", err) + } + if fallbackGroup.Platform != PlatformAnthropic { + return fmt.Errorf("fallback group must be anthropic platform") + } + if fallbackGroup.SubscriptionType == SubscriptionTypeSubscription { + return fmt.Errorf("fallback group cannot be subscription type") + } + if fallbackGroup.FallbackGroupIDOnInvalidRequest != nil { + return fmt.Errorf("fallback group cannot have invalid request fallback configured") + } + return nil +} + +func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error) { + group, err := s.groupRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + + if input.Name != "" { + group.Name = input.Name + } + if input.Description != nil { + group.Description = *input.Description + } + if input.Platform != "" { + group.Platform = input.Platform + } + if input.RateMultiplier != nil { + if *input.RateMultiplier <= 0 { + return nil, errors.New("rate_multiplier must be > 0") + } + group.RateMultiplier = *input.RateMultiplier + } + if input.IsExclusive != nil { + group.IsExclusive = *input.IsExclusive + } + if input.Status != "" { + group.Status = input.Status + } + + // 订阅相关字段 + if input.SubscriptionType != "" { + group.SubscriptionType = input.SubscriptionType + } + // 限额字段:nil/负数 表示"无限制",0 表示"不允许用量",正数表示具体限额 + // 前端始终发送这三个字段,无需 nil 守卫 + group.DailyLimitUSD = normalizeLimit(input.DailyLimitUSD) + group.WeeklyLimitUSD = normalizeLimit(input.WeeklyLimitUSD) + group.MonthlyLimitUSD = normalizeLimit(input.MonthlyLimitUSD) + // 图片生成计费配置:负数表示清除(使用默认价格) + if input.AllowImageGeneration != nil { + group.AllowImageGeneration = *input.AllowImageGeneration + } + if input.AllowBatchImageGeneration != nil { + group.AllowBatchImageGeneration = *input.AllowBatchImageGeneration + } + if !group.AllowImageGeneration || group.Platform != PlatformGemini { + group.AllowBatchImageGeneration = false + } + if input.ImageRateIndependent != nil { + group.ImageRateIndependent = *input.ImageRateIndependent + } + if input.ImageRateMultiplier != nil { + if *input.ImageRateMultiplier < 0 { + return nil, errors.New("image_rate_multiplier must be >= 0") + } + group.ImageRateMultiplier = *input.ImageRateMultiplier + } + if input.BatchImageDiscountMultiplier != nil { + if *input.BatchImageDiscountMultiplier < 0 { + return nil, errors.New("batch_image_discount_multiplier must be >= 0") + } + group.BatchImageDiscountMultiplier = *input.BatchImageDiscountMultiplier + } + if input.BatchImageHoldMultiplier != nil { + if *input.BatchImageHoldMultiplier < 0 { + return nil, errors.New("batch_image_hold_multiplier must be >= 0") + } + group.BatchImageHoldMultiplier = *input.BatchImageHoldMultiplier + } + // 仅在本次更新显式触碰任一比例时校验合并后的不变式(hold >= discount), + // 避免存量脏数据阻塞其他字段的正常更新(提交侧另有钳制兜底)。 + if (input.BatchImageDiscountMultiplier != nil || input.BatchImageHoldMultiplier != nil) && + group.BatchImageHoldMultiplier < group.BatchImageDiscountMultiplier { + return nil, errors.New("batch_image_hold_multiplier must be >= batch_image_discount_multiplier") + } + if input.VideoRateIndependent != nil { + group.VideoRateIndependent = *input.VideoRateIndependent + } + if input.VideoRateMultiplier != nil { + if *input.VideoRateMultiplier < 0 { + return nil, errors.New("video_rate_multiplier must be >= 0") + } + group.VideoRateMultiplier = *input.VideoRateMultiplier + } + if input.PeakRateEnabled != nil { + group.PeakRateEnabled = *input.PeakRateEnabled + } + if input.PeakStart != nil { + group.PeakStart = *input.PeakStart + } + if input.PeakEnd != nil { + group.PeakEnd = *input.PeakEnd + } + if input.PeakRateMultiplier != nil { + group.PeakRateMultiplier = *input.PeakRateMultiplier + } + // 先归一化(非订阅分组——含本次更新转为非订阅——静默清空高峰配置,清洗停用状态下的脏字段), + // 再收敛校验:Update 可能只传部分 peak 字段,需对合并后的最终配置统一校验, + // 防止单独修改 start/end 导致最终 start>=end 等非法配置入库。与 CreateGroup 同一收口。 + group.PeakRateEnabled, group.PeakStart, group.PeakEnd, group.PeakRateMultiplier = NormalizePeakRateConfig(group.SubscriptionType, group.PeakRateEnabled, group.PeakStart, group.PeakEnd, group.PeakRateMultiplier) + if err := ValidatePeakRateConfig(group.SubscriptionType, group.PeakRateEnabled, group.PeakStart, group.PeakEnd, group.PeakRateMultiplier); err != nil { + return nil, err + } + if input.ImagePrice1K != nil { + group.ImagePrice1K = normalizePrice(input.ImagePrice1K) + } + if input.ImagePrice2K != nil { + group.ImagePrice2K = normalizePrice(input.ImagePrice2K) + } + if input.ImagePrice4K != nil { + group.ImagePrice4K = normalizePrice(input.ImagePrice4K) + } + if input.VideoPrice480P != nil { + group.VideoPrice480P = normalizePrice(input.VideoPrice480P) + } + if input.VideoPrice720P != nil { + group.VideoPrice720P = normalizePrice(input.VideoPrice720P) + } + if input.VideoPrice1080P != nil { + group.VideoPrice1080P = normalizePrice(input.VideoPrice1080P) + } + + // Claude Code 客户端限制 + if input.ClaudeCodeOnly != nil { + group.ClaudeCodeOnly = *input.ClaudeCodeOnly + } + if input.FallbackGroupID != nil { + // 校验降级分组 + if *input.FallbackGroupID > 0 { + if err := s.validateFallbackGroup(ctx, id, *input.FallbackGroupID); err != nil { + return nil, err + } + group.FallbackGroupID = input.FallbackGroupID + } else { + // 传入 0 或负数表示清除降级分组 + group.FallbackGroupID = nil + } + } + fallbackOnInvalidRequest := group.FallbackGroupIDOnInvalidRequest + if input.FallbackGroupIDOnInvalidRequest != nil { + if *input.FallbackGroupIDOnInvalidRequest > 0 { + fallbackOnInvalidRequest = input.FallbackGroupIDOnInvalidRequest + } else { + fallbackOnInvalidRequest = nil + } + } + if fallbackOnInvalidRequest != nil { + if err := s.validateFallbackGroupOnInvalidRequest(ctx, id, group.Platform, group.SubscriptionType, *fallbackOnInvalidRequest); err != nil { + return nil, err + } + } + group.FallbackGroupIDOnInvalidRequest = fallbackOnInvalidRequest + + // 模型路由配置 + if input.ModelRouting != nil { + group.ModelRouting = input.ModelRouting + } + if input.ModelRoutingEnabled != nil { + group.ModelRoutingEnabled = *input.ModelRoutingEnabled + } + if input.MCPXMLInject != nil { + group.MCPXMLInject = *input.MCPXMLInject + } + + // 支持的模型系列(仅 antigravity 平台使用) + if input.SupportedModelScopes != nil { + group.SupportedModelScopes = *input.SupportedModelScopes + } + + // OpenAI Messages 调度配置 + if input.AllowMessagesDispatch != nil { + group.AllowMessagesDispatch = *input.AllowMessagesDispatch + } + if input.RequireOAuthOnly != nil { + group.RequireOAuthOnly = *input.RequireOAuthOnly + } + if input.RequirePrivacySet != nil { + group.RequirePrivacySet = *input.RequirePrivacySet + } + if input.DefaultMappedModel != nil { + group.DefaultMappedModel = *input.DefaultMappedModel + } + if input.MessagesDispatchModelConfig != nil { + group.MessagesDispatchModelConfig = normalizeOpenAIMessagesDispatchModelConfig(*input.MessagesDispatchModelConfig) + } + if input.ModelsListConfig != nil { + group.ModelsListConfig = normalizeGroupModelsListConfig(*input.ModelsListConfig) + } + if input.RPMLimit != nil { + group.RPMLimit = *input.RPMLimit + } + sanitizeGroupMessagesDispatchFields(group) + + if err := s.groupRepo.Update(ctx, group); err != nil { + return nil, err + } + + if s.authCacheInvalidator != nil { + s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, id) + } + + // 如果指定了复制账号的源分组,同步绑定(替换当前分组的账号) + if len(input.CopyAccountsFromGroupIDs) > 0 { + // 去重源分组 IDs + seen := make(map[int64]struct{}) + uniqueSourceGroupIDs := make([]int64, 0, len(input.CopyAccountsFromGroupIDs)) + for _, srcGroupID := range input.CopyAccountsFromGroupIDs { + // 校验:源分组不能是自身 + if srcGroupID == id { + return nil, fmt.Errorf("cannot copy accounts from self") + } + // 去重 + if _, exists := seen[srcGroupID]; !exists { + seen[srcGroupID] = struct{}{} + uniqueSourceGroupIDs = append(uniqueSourceGroupIDs, srcGroupID) + } + } + + // 校验源分组的平台是否与当前分组一致 + for _, srcGroupID := range uniqueSourceGroupIDs { + srcGroup, err := s.groupRepo.GetByIDLite(ctx, srcGroupID) + if err != nil { + return nil, fmt.Errorf("source group %d not found: %w", srcGroupID, err) + } + if srcGroup.Platform != group.Platform { + return nil, fmt.Errorf("source group %d platform mismatch: expected %s, got %s", srcGroupID, group.Platform, srcGroup.Platform) + } + } + + // 获取所有源分组的账号(去重) + accountIDsToCopy, err := s.groupRepo.GetAccountIDsByGroupIDs(ctx, uniqueSourceGroupIDs) + if err != nil { + return nil, fmt.Errorf("failed to get accounts from source groups: %w", err) + } + + // 先清空当前分组的所有账号绑定 + if _, err := s.groupRepo.DeleteAccountGroupsByGroupID(ctx, id); err != nil { + return nil, fmt.Errorf("failed to clear existing account bindings: %w", err) + } + + // require_oauth_only: 过滤掉 apikey 类型账号 + if group.RequireOAuthOnly && (group.Platform == PlatformOpenAI || group.Platform == PlatformAntigravity || group.Platform == PlatformAnthropic || group.Platform == PlatformGemini || group.Platform == PlatformGrok) && len(accountIDsToCopy) > 0 { + accounts, err := s.accountRepo.GetByIDs(ctx, accountIDsToCopy) + if err != nil { + return nil, fmt.Errorf("failed to fetch accounts for oauth filter: %w", err) + } + oauthIDs := make(map[int64]struct{}, len(accounts)) + for _, acc := range accounts { + if acc.Type != AccountTypeAPIKey { + oauthIDs[acc.ID] = struct{}{} + } + } + var filtered []int64 + for _, aid := range accountIDsToCopy { + if _, ok := oauthIDs[aid]; ok { + filtered = append(filtered, aid) + } + } + accountIDsToCopy = filtered + } + + // 再绑定源分组的账号 + if len(accountIDsToCopy) > 0 { + if err := s.groupRepo.BindAccountsToGroup(ctx, id, accountIDsToCopy); err != nil { + return nil, fmt.Errorf("failed to bind accounts to group: %w", err) + } + } + } + + return group, nil +} + +func (s *adminServiceImpl) DeleteGroup(ctx context.Context, id int64) error { + var groupKeys []string + if s.authCacheInvalidator != nil { + keys, err := s.apiKeyRepo.ListKeysByGroupID(ctx, id) + if err == nil { + groupKeys = keys + } + } + + affectedUserIDs, err := s.groupRepo.DeleteCascade(ctx, id) + if err != nil { + return err + } + // 注意:user_group_rate_multipliers 表通过外键 ON DELETE CASCADE 自动清理 + + // 事务成功后,异步失效受影响用户的订阅缓存 + if len(affectedUserIDs) > 0 && s.billingCacheService != nil { + groupID := id + go func() { + cacheCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + for _, userID := range affectedUserIDs { + if err := s.billingCacheService.InvalidateSubscription(cacheCtx, userID, groupID); err != nil { + logger.LegacyPrintf("service.admin", "invalidate subscription cache failed: user_id=%d group_id=%d err=%v", userID, groupID, err) + } + } + }() + } + if s.authCacheInvalidator != nil { + for _, key := range groupKeys { + s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, key) + } + } + + return nil +} + +func (s *adminServiceImpl) GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]APIKey, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize} + keys, result, err := s.apiKeyRepo.ListByGroupID(ctx, groupID, params) + if err != nil { + return nil, 0, err + } + return keys, result.Total, nil +} + +func (s *adminServiceImpl) GetGroupRateMultipliers(ctx context.Context, groupID int64) ([]UserGroupRateEntry, error) { + if s.userGroupRateRepo == nil { + return nil, nil + } + return s.userGroupRateRepo.GetByGroupID(ctx, groupID) +} + +func (s *adminServiceImpl) ClearGroupRateMultipliers(ctx context.Context, groupID int64) error { + if s.userGroupRateRepo == nil { + return nil + } + return s.userGroupRateRepo.DeleteByGroupID(ctx, groupID) +} + +func (s *adminServiceImpl) BatchSetGroupRateMultipliers(ctx context.Context, groupID int64, entries []GroupRateMultiplierInput) error { + if s.userGroupRateRepo == nil { + return nil + } + for _, e := range entries { + if e.RateMultiplier <= 0 { + return fmt.Errorf("rate_multiplier must be > 0 (user_id=%d)", e.UserID) + } + } + return s.userGroupRateRepo.SyncGroupRateMultipliers(ctx, groupID, entries) +} + +func (s *adminServiceImpl) ClearGroupRPMOverrides(ctx context.Context, groupID int64) error { + if s.userGroupRateRepo == nil { + return nil + } + if err := s.userGroupRateRepo.ClearGroupRPMOverrides(ctx, groupID); err != nil { + return err + } + // RPM override 已嵌入 auth cache snapshot (v7),变更后必须失效相关缓存。 + if s.authCacheInvalidator != nil { + s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, groupID) + } + return nil +} + +func (s *adminServiceImpl) BatchSetGroupRPMOverrides(ctx context.Context, groupID int64, entries []GroupRPMOverrideInput) error { + if s.userGroupRateRepo == nil { + return nil + } + for _, e := range entries { + if e.RPMOverride != nil && *e.RPMOverride < 0 { + return infraerrors.BadRequest("INVALID_RPM_OVERRIDE", fmt.Sprintf("rpm_override must be >= 0 (user_id=%d)", e.UserID)) + } + } + if err := s.userGroupRateRepo.SyncGroupRPMOverrides(ctx, groupID, entries); err != nil { + return err + } + // RPM override 已嵌入 auth cache snapshot (v7),变更后必须失效相关缓存。 + if s.authCacheInvalidator != nil { + s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, groupID) + } + return nil +} + +func (s *adminServiceImpl) UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error { + return s.groupRepo.UpdateSortOrders(ctx, updates) +} + +// AdminUpdateAPIKeyGroupID 管理员修改 API Key 分组绑定 +// groupID: nil=不修改, 指向0=解绑, 指向正整数=绑定到目标分组 +func (s *adminServiceImpl) AdminUpdateAPIKeyGroupID(ctx context.Context, keyID int64, groupID *int64) (*AdminUpdateAPIKeyGroupIDResult, error) { + apiKey, err := s.apiKeyRepo.GetByID(ctx, keyID) + if err != nil { + return nil, err + } + + if groupID == nil { + // nil 表示不修改,直接返回 + return &AdminUpdateAPIKeyGroupIDResult{APIKey: apiKey}, nil + } + + if *groupID < 0 { + return nil, infraerrors.BadRequest("INVALID_GROUP_ID", "group_id must be non-negative") + } + + result := &AdminUpdateAPIKeyGroupIDResult{} + + if *groupID == 0 { + // 0 表示解绑分组(不修改 user_allowed_groups,避免影响用户其他 Key) + apiKey.GroupID = nil + apiKey.Group = nil + } else { + // 验证目标分组存在且状态为 active + group, err := s.groupRepo.GetByID(ctx, *groupID) + if err != nil { + return nil, err + } + if group.Status != StatusActive { + return nil, infraerrors.BadRequest("GROUP_NOT_ACTIVE", "target group is not active") + } + // 订阅类型分组:用户须持有该分组的有效订阅才可绑定 + if group.IsSubscriptionType() { + if s.userSubRepo == nil { + return nil, infraerrors.InternalServer("SUBSCRIPTION_REPOSITORY_UNAVAILABLE", "subscription repository is not configured") + } + if _, err := s.userSubRepo.GetActiveByUserIDAndGroupID(ctx, apiKey.UserID, *groupID); err != nil { + if errors.Is(err, ErrSubscriptionNotFound) { + return nil, infraerrors.BadRequest("SUBSCRIPTION_REQUIRED", "user does not have an active subscription for this group") + } + return nil, err + } + } + + gid := *groupID + apiKey.GroupID = &gid + apiKey.Group = group + + // 专属标准分组:使用事务保证「添加分组权限」与「更新 API Key」的原子性 + if group.IsExclusive && !group.IsSubscriptionType() { + opCtx := ctx + var tx *dbent.Tx + if s.entClient == nil { + logger.LegacyPrintf("service.admin", "Warning: entClient is nil, skipping transaction protection for exclusive group binding") + } else { + var txErr error + tx, txErr = s.entClient.Tx(ctx) + if txErr != nil { + return nil, fmt.Errorf("begin transaction: %w", txErr) + } + defer func() { _ = tx.Rollback() }() + opCtx = dbent.NewTxContext(ctx, tx) + } + + if addErr := s.userRepo.AddGroupToAllowedGroups(opCtx, apiKey.UserID, gid); addErr != nil { + return nil, fmt.Errorf("add group to user allowed groups: %w", addErr) + } + if err := s.apiKeyRepo.Update(opCtx, apiKey); err != nil { + return nil, fmt.Errorf("update api key: %w", err) + } + if tx != nil { + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit transaction: %w", err) + } + } + + result.AutoGrantedGroupAccess = true + result.GrantedGroupID = &gid + result.GrantedGroupName = group.Name + + // 失效认证缓存(在事务提交后执行) + if s.authCacheInvalidator != nil { + s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, apiKey.Key) + } + + result.APIKey = apiKey + return result, nil + } + } + + // 非专属分组 / 解绑:无需事务,单步更新即可 + if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil { + return nil, fmt.Errorf("update api key: %w", err) + } + + // 失效认证缓存 + if s.authCacheInvalidator != nil { + s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, apiKey.Key) + } + + result.APIKey = apiKey + return result, nil +} + +// AdminResetAPIKeyRateLimitUsage resets all API key rate-limit usage windows. +func (s *adminServiceImpl) AdminResetAPIKeyRateLimitUsage(ctx context.Context, keyID int64) (*APIKey, error) { + apiKey, err := s.apiKeyRepo.GetByID(ctx, keyID) + if err != nil { + return nil, err + } + apiKey.Usage5h = 0 + apiKey.Usage1d = 0 + apiKey.Usage7d = 0 + apiKey.Window5hStart = nil + apiKey.Window1dStart = nil + apiKey.Window7dStart = nil + if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil { + return nil, fmt.Errorf("reset api key rate limit usage: %w", err) + } + if s.authCacheInvalidator != nil { + s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, apiKey.Key) + } + if s.billingCacheService != nil { + _ = s.billingCacheService.InvalidateAPIKeyRateLimit(ctx, apiKey.ID) + } + return apiKey, nil +} + +// ReplaceUserGroup 替换用户的专属分组 +func (s *adminServiceImpl) ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*ReplaceUserGroupResult, error) { + if oldGroupID == newGroupID { + return nil, infraerrors.BadRequest("SAME_GROUP", "old and new group must be different") + } + + // 验证新分组存在且为活跃的专属标准分组 + newGroup, err := s.groupRepo.GetByID(ctx, newGroupID) + if err != nil { + return nil, err + } + if newGroup.Status != StatusActive { + return nil, infraerrors.BadRequest("GROUP_NOT_ACTIVE", "target group is not active") + } + if !newGroup.IsExclusive { + return nil, infraerrors.BadRequest("GROUP_NOT_EXCLUSIVE", "target group is not exclusive") + } + if newGroup.IsSubscriptionType() { + return nil, infraerrors.BadRequest("GROUP_IS_SUBSCRIPTION", "subscription groups are not supported for replacement") + } + + // 事务保证原子性 + if s.entClient == nil { + return nil, fmt.Errorf("entClient is nil, cannot perform group replacement") + } + tx, err := s.entClient.Tx(ctx) + if err != nil { + return nil, fmt.Errorf("begin transaction: %w", err) + } + defer func() { _ = tx.Rollback() }() + opCtx := dbent.NewTxContext(ctx, tx) + + // 1. 授予新分组权限 + if err := s.userRepo.AddGroupToAllowedGroups(opCtx, userID, newGroupID); err != nil { + return nil, fmt.Errorf("add new group to allowed groups: %w", err) + } + + // 2. 迁移绑定旧分组的 Key 到新分组 + migrated, err := s.apiKeyRepo.UpdateGroupIDByUserAndGroup(opCtx, userID, oldGroupID, newGroupID) + if err != nil { + return nil, fmt.Errorf("migrate api keys: %w", err) + } + + // 3. 移除旧分组权限 + if err := s.userRepo.RemoveGroupFromUserAllowedGroups(opCtx, userID, oldGroupID); err != nil { + return nil, fmt.Errorf("remove old group from allowed groups: %w", err) + } + + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit transaction: %w", err) + } + + // 失效该用户所有 Key 的认证缓存 + if s.authCacheInvalidator != nil { + keys, keyErr := s.apiKeyRepo.ListKeysByUserID(ctx, userID) + if keyErr == nil { + for _, k := range keys { + s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, k) + } + } + } + + return &ReplaceUserGroupResult{MigratedKeys: migrated}, nil +} diff --git a/backend/internal/service/admin_proxy.go b/backend/internal/service/admin_proxy.go new file mode 100644 index 0000000000..9b5a9223c7 --- /dev/null +++ b/backend/internal/service/admin_proxy.go @@ -0,0 +1,608 @@ +package service + +import ( + "context" + "fmt" + "io" + "net/http" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/util/httputil" +) + +// Proxy management implementations +func (s *adminServiceImpl) ListProxies(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]Proxy, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + proxies, result, err := s.proxyRepo.ListWithFilters(ctx, params, protocol, status, search) + if err != nil { + return nil, 0, err + } + return proxies, result.Total, nil +} + +func (s *adminServiceImpl) ListProxiesWithAccountCount(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]ProxyWithAccountCount, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + proxies, result, err := s.proxyRepo.ListWithFiltersAndAccountCount(ctx, params, protocol, status, search) + if err != nil { + return nil, 0, err + } + s.attachProxyLatency(ctx, proxies) + return proxies, result.Total, nil +} + +func (s *adminServiceImpl) GetAllProxies(ctx context.Context) ([]Proxy, error) { + return s.proxyRepo.ListActive(ctx) +} + +func (s *adminServiceImpl) GetAllProxiesWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error) { + proxies, err := s.proxyRepo.ListActiveWithAccountCount(ctx) + if err != nil { + return nil, err + } + s.attachProxyLatency(ctx, proxies) + return proxies, nil +} + +func (s *adminServiceImpl) GetProxy(ctx context.Context, id int64) (*Proxy, error) { + return s.proxyRepo.GetByID(ctx, id) +} + +func (s *adminServiceImpl) GetProxiesByIDs(ctx context.Context, ids []int64) ([]Proxy, error) { + return s.proxyRepo.ListByIDs(ctx, ids) +} + +func (s *adminServiceImpl) CreateProxy(ctx context.Context, input *CreateProxyInput) (*Proxy, error) { + // 规范化 fallback_mode + mode := input.FallbackMode + if mode == "" { + mode = FallbackModeNone + } + // 校验:mode=proxy 必须有 backup + if mode == FallbackModeProxy && input.BackupProxyID == nil { + return nil, infraerrors.BadRequest("PROXY_BACKUP_REQUIRED", "backup proxy required when fallback_mode=proxy") + } + if input.ExpiryWarnDays < 0 { + return nil, infraerrors.BadRequest("PROXY_WARN_DAYS_INVALID", "expiry_warn_days must be >= 0") + } + + proxy := &Proxy{ + Name: input.Name, + Protocol: input.Protocol, + Host: input.Host, + Port: input.Port, + Username: input.Username, + Password: input.Password, + Status: StatusActive, + ExpiresAt: input.ExpiresAt, + FallbackMode: mode, + BackupProxyID: input.BackupProxyID, + ExpiryWarnDays: input.ExpiryWarnDays, + } + if err := s.proxyRepo.Create(ctx, proxy); err != nil { + return nil, err + } + // Probe latency asynchronously so creation isn't blocked by network timeout. + go s.probeProxyLatency(context.Background(), proxy) + return proxy, nil +} + +func (s *adminServiceImpl) UpdateProxy(ctx context.Context, id int64, input *UpdateProxyInput) (*Proxy, error) { + // 校验:backup_proxy_id 不能是自身 + if input.BackupProxyID != nil && *input.BackupProxyID == id { + return nil, infraerrors.BadRequest("PROXY_BACKUP_SELF", "backup proxy cannot be itself") + } + // 规范化 fallback_mode + mode := input.FallbackMode + if mode == "" { + mode = FallbackModeNone + } + // 校验:mode=proxy 必须有 backup + if mode == FallbackModeProxy && input.BackupProxyID == nil { + return nil, infraerrors.BadRequest("PROXY_BACKUP_REQUIRED", "backup proxy required when fallback_mode=proxy") + } + if input.ExpiryWarnDays < 0 { + return nil, infraerrors.BadRequest("PROXY_WARN_DAYS_INVALID", "expiry_warn_days must be >= 0") + } + + proxy, err := s.proxyRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + + if input.Name != "" { + proxy.Name = input.Name + } + if input.Protocol != "" { + proxy.Protocol = input.Protocol + } + if input.Host != "" { + proxy.Host = input.Host + } + if input.Port != 0 { + proxy.Port = input.Port + } + if input.Username != "" { + proxy.Username = input.Username + } + if input.Password != "" { + proxy.Password = input.Password + } + if input.Status != "" { + proxy.Status = input.Status + } + // 透传有效期与回退字段 + proxy.ExpiresAt = input.ExpiresAt + proxy.FallbackMode = mode + proxy.BackupProxyID = input.BackupProxyID + proxy.ExpiryWarnDays = input.ExpiryWarnDays + + if err := s.proxyRepo.Update(ctx, proxy); err != nil { + return nil, err + } + return proxy, nil +} + +func (s *adminServiceImpl) DeleteProxy(ctx context.Context, id int64) error { + count, err := s.proxyRepo.CountAccountsByProxyID(ctx, id) + if err != nil { + return err + } + if count > 0 { + return ErrProxyInUse + } + return s.proxyRepo.Delete(ctx, id) +} + +func (s *adminServiceImpl) BatchDeleteProxies(ctx context.Context, ids []int64) (*ProxyBatchDeleteResult, error) { + result := &ProxyBatchDeleteResult{} + if len(ids) == 0 { + return result, nil + } + + for _, id := range ids { + count, err := s.proxyRepo.CountAccountsByProxyID(ctx, id) + if err != nil { + result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{ + ID: id, + Reason: err.Error(), + }) + continue + } + if count > 0 { + result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{ + ID: id, + Reason: ErrProxyInUse.Error(), + }) + continue + } + if err := s.proxyRepo.Delete(ctx, id); err != nil { + result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{ + ID: id, + Reason: err.Error(), + }) + continue + } + result.DeletedIDs = append(result.DeletedIDs, id) + } + + return result, nil +} + +func (s *adminServiceImpl) GetProxyAccounts(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error) { + return s.proxyRepo.ListAccountSummariesByProxyID(ctx, proxyID) +} + +func (s *adminServiceImpl) CheckProxyExists(ctx context.Context, host string, port int, username, password string) (bool, error) { + return s.proxyRepo.ExistsByHostPortAuth(ctx, host, port, username, password) +} + +func (s *adminServiceImpl) TestProxy(ctx context.Context, id int64) (*ProxyTestResult, error) { + proxy, err := s.proxyRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + + proxyURL := proxy.URL() + exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxyURL) + if err != nil { + s.saveProxyLatency(ctx, id, &ProxyLatencyInfo{ + Success: false, + Message: err.Error(), + UpdatedAt: time.Now(), + }) + return &ProxyTestResult{ + Success: false, + Message: err.Error(), + }, nil + } + + latency := latencyMs + s.saveProxyLatency(ctx, id, &ProxyLatencyInfo{ + Success: true, + LatencyMs: &latency, + Message: "Proxy is accessible", + IPAddress: exitInfo.IP, + Country: exitInfo.Country, + CountryCode: exitInfo.CountryCode, + Region: exitInfo.Region, + City: exitInfo.City, + UpdatedAt: time.Now(), + }) + return &ProxyTestResult{ + Success: true, + Message: "Proxy is accessible", + LatencyMs: latencyMs, + IPAddress: exitInfo.IP, + City: exitInfo.City, + Region: exitInfo.Region, + Country: exitInfo.Country, + CountryCode: exitInfo.CountryCode, + }, nil +} + +func (s *adminServiceImpl) CheckProxyQuality(ctx context.Context, id int64) (*ProxyQualityCheckResult, error) { + proxy, err := s.proxyRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + + result := &ProxyQualityCheckResult{ + ProxyID: id, + Score: 100, + Grade: "A", + CheckedAt: time.Now().Unix(), + Items: make([]ProxyQualityCheckItem, 0, len(proxyQualityTargets)+1), + } + + proxyURL := proxy.URL() + if s.proxyProber == nil { + result.Items = append(result.Items, ProxyQualityCheckItem{ + Target: "base_connectivity", + Status: "fail", + Message: "代理探测服务未配置", + }) + result.FailedCount++ + finalizeProxyQualityResult(result) + s.saveProxyQualitySnapshot(ctx, id, result, nil) + return result, nil + } + + exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxyURL) + if err != nil { + result.Items = append(result.Items, ProxyQualityCheckItem{ + Target: "base_connectivity", + Status: "fail", + LatencyMs: latencyMs, + Message: err.Error(), + }) + result.FailedCount++ + finalizeProxyQualityResult(result) + s.saveProxyQualitySnapshot(ctx, id, result, nil) + return result, nil + } + + result.ExitIP = exitInfo.IP + result.Country = exitInfo.Country + result.CountryCode = exitInfo.CountryCode + result.BaseLatencyMs = latencyMs + result.Items = append(result.Items, ProxyQualityCheckItem{ + Target: "base_connectivity", + Status: "pass", + LatencyMs: latencyMs, + Message: "代理出口连通正常", + }) + result.PassedCount++ + + client, err := httpclient.GetClient(httpclient.Options{ + ProxyURL: proxyURL, + Timeout: proxyQualityRequestTimeout, + ResponseHeaderTimeout: proxyQualityResponseHeaderTimeout, + }) + if err != nil { + result.Items = append(result.Items, ProxyQualityCheckItem{ + Target: "http_client", + Status: "fail", + Message: fmt.Sprintf("创建检测客户端失败: %v", err), + }) + result.FailedCount++ + finalizeProxyQualityResult(result) + s.saveProxyQualitySnapshot(ctx, id, result, exitInfo) + return result, nil + } + + for _, target := range proxyQualityTargets { + item := runProxyQualityTarget(ctx, client, target) + result.Items = append(result.Items, item) + switch item.Status { + case "pass": + result.PassedCount++ + case "warn": + result.WarnCount++ + case "challenge": + result.ChallengeCount++ + default: + result.FailedCount++ + } + } + + finalizeProxyQualityResult(result) + s.saveProxyQualitySnapshot(ctx, id, result, exitInfo) + return result, nil +} + +func runProxyQualityTarget(ctx context.Context, client *http.Client, target proxyQualityTarget) ProxyQualityCheckItem { + item := ProxyQualityCheckItem{ + Target: target.Target, + } + + req, err := http.NewRequestWithContext(ctx, target.Method, target.URL, nil) + if err != nil { + item.Status = "fail" + item.Message = fmt.Sprintf("构建请求失败: %v", err) + return item + } + req.Header.Set("Accept", "application/json,text/html,*/*") + req.Header.Set("User-Agent", proxyQualityClientUserAgent) + + start := time.Now() + resp, err := client.Do(req) + if err != nil { + item.Status = "fail" + item.LatencyMs = time.Since(start).Milliseconds() + item.Message = fmt.Sprintf("请求失败: %v", err) + return item + } + defer func() { _ = resp.Body.Close() }() + item.LatencyMs = time.Since(start).Milliseconds() + item.HTTPStatus = resp.StatusCode + + body, readErr := io.ReadAll(io.LimitReader(resp.Body, proxyQualityMaxBodyBytes+1)) + if readErr != nil { + item.Status = "fail" + item.Message = fmt.Sprintf("读取响应失败: %v", readErr) + return item + } + if int64(len(body)) > proxyQualityMaxBodyBytes { + body = body[:proxyQualityMaxBodyBytes] + } + + // Cloudflare challenge 检测 + if httputil.IsCloudflareChallengeResponse(resp.StatusCode, resp.Header, body) { + item.Status = "challenge" + item.CFRay = httputil.ExtractCloudflareRayID(resp.Header, body) + item.Message = "命中 Cloudflare challenge" + return item + } + + if _, ok := target.AllowedStatuses[resp.StatusCode]; ok { + // 白名单内的状态码均代表目标可达:2xx 表示接口直接可用, + // 401/405 等是无鉴权探测的预期结果,同样视为连通正常,不再扣分。 + item.Status = "pass" + if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices { + item.Message = fmt.Sprintf("HTTP %d", resp.StatusCode) + } else { + item.Message = fmt.Sprintf("HTTP %d(目标可达)", resp.StatusCode) + } + return item + } + + if resp.StatusCode == http.StatusTooManyRequests { + item.Status = "warn" + item.Message = "目标返回 429,可能存在频控" + return item + } + + item.Status = "fail" + item.Message = fmt.Sprintf("非预期状态码: %d", resp.StatusCode) + return item +} + +func finalizeProxyQualityResult(result *ProxyQualityCheckResult) { + if result == nil { + return + } + score := 100 - result.WarnCount*10 - result.FailedCount*22 - result.ChallengeCount*30 + if score < 0 { + score = 0 + } + result.Score = score + result.Grade = proxyQualityGrade(score) + result.Summary = fmt.Sprintf( + "通过 %d 项,告警 %d 项,失败 %d 项,挑战 %d 项", + result.PassedCount, + result.WarnCount, + result.FailedCount, + result.ChallengeCount, + ) +} + +func proxyQualityGrade(score int) string { + switch { + case score >= 90: + return "A" + case score >= 75: + return "B" + case score >= 60: + return "C" + case score >= 40: + return "D" + default: + return "F" + } +} + +func proxyQualityOverallStatus(result *ProxyQualityCheckResult) string { + if result == nil { + return "" + } + if result.ChallengeCount > 0 { + return "challenge" + } + if result.FailedCount > 0 { + return "failed" + } + if result.WarnCount > 0 { + return "warn" + } + if result.PassedCount > 0 { + return "healthy" + } + return "failed" +} + +func proxyQualityFirstCFRay(result *ProxyQualityCheckResult) string { + if result == nil { + return "" + } + for _, item := range result.Items { + if item.CFRay != "" { + return item.CFRay + } + } + return "" +} + +func proxyQualityBaseConnectivityPass(result *ProxyQualityCheckResult) bool { + if result == nil { + return false + } + for _, item := range result.Items { + if item.Target == "base_connectivity" { + return item.Status == "pass" + } + } + return false +} + +func (s *adminServiceImpl) saveProxyQualitySnapshot(ctx context.Context, proxyID int64, result *ProxyQualityCheckResult, exitInfo *ProxyExitInfo) { + if result == nil { + return + } + score := result.Score + checkedAt := result.CheckedAt + info := &ProxyLatencyInfo{ + Success: proxyQualityBaseConnectivityPass(result), + Message: result.Summary, + QualityStatus: proxyQualityOverallStatus(result), + QualityScore: &score, + QualityGrade: result.Grade, + QualitySummary: result.Summary, + QualityCheckedAt: &checkedAt, + QualityCFRay: proxyQualityFirstCFRay(result), + UpdatedAt: time.Now(), + } + if result.BaseLatencyMs > 0 { + latency := result.BaseLatencyMs + info.LatencyMs = &latency + } + if exitInfo != nil { + info.IPAddress = exitInfo.IP + info.Country = exitInfo.Country + info.CountryCode = exitInfo.CountryCode + info.Region = exitInfo.Region + info.City = exitInfo.City + } + s.saveProxyLatency(ctx, proxyID, info) +} + +func (s *adminServiceImpl) probeProxyLatency(ctx context.Context, proxy *Proxy) { + if s.proxyProber == nil || proxy == nil { + return + } + exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxy.URL()) + if err != nil { + s.saveProxyLatency(ctx, proxy.ID, &ProxyLatencyInfo{ + Success: false, + Message: err.Error(), + UpdatedAt: time.Now(), + }) + return + } + + latency := latencyMs + s.saveProxyLatency(ctx, proxy.ID, &ProxyLatencyInfo{ + Success: true, + LatencyMs: &latency, + Message: "Proxy is accessible", + IPAddress: exitInfo.IP, + Country: exitInfo.Country, + CountryCode: exitInfo.CountryCode, + Region: exitInfo.Region, + City: exitInfo.City, + UpdatedAt: time.Now(), + }) +} + +func (s *adminServiceImpl) attachProxyLatency(ctx context.Context, proxies []ProxyWithAccountCount) { + if s.proxyLatencyCache == nil || len(proxies) == 0 { + return + } + + ids := make([]int64, 0, len(proxies)) + for i := range proxies { + ids = append(ids, proxies[i].ID) + } + + latencies, err := s.proxyLatencyCache.GetProxyLatencies(ctx, ids) + if err != nil { + logger.LegacyPrintf("service.admin", "Warning: load proxy latency cache failed: %v", err) + return + } + + for i := range proxies { + info := latencies[proxies[i].ID] + if info == nil { + continue + } + if info.Success { + proxies[i].LatencyStatus = "success" + proxies[i].LatencyMs = info.LatencyMs + } else { + proxies[i].LatencyStatus = "failed" + } + proxies[i].LatencyMessage = info.Message + proxies[i].IPAddress = info.IPAddress + proxies[i].Country = info.Country + proxies[i].CountryCode = info.CountryCode + proxies[i].Region = info.Region + proxies[i].City = info.City + proxies[i].QualityStatus = info.QualityStatus + proxies[i].QualityScore = info.QualityScore + proxies[i].QualityGrade = info.QualityGrade + proxies[i].QualitySummary = info.QualitySummary + proxies[i].QualityChecked = info.QualityCheckedAt + } +} + +func (s *adminServiceImpl) saveProxyLatency(ctx context.Context, proxyID int64, info *ProxyLatencyInfo) { + if s.proxyLatencyCache == nil || info == nil { + return + } + + merged := *info + if latencies, err := s.proxyLatencyCache.GetProxyLatencies(ctx, []int64{proxyID}); err == nil { + if existing := latencies[proxyID]; existing != nil { + if merged.QualityCheckedAt == nil && + merged.QualityScore == nil && + merged.QualityGrade == "" && + merged.QualityStatus == "" && + merged.QualitySummary == "" && + merged.QualityCFRay == "" { + merged.QualityStatus = existing.QualityStatus + merged.QualityScore = existing.QualityScore + merged.QualityGrade = existing.QualityGrade + merged.QualitySummary = existing.QualitySummary + merged.QualityCheckedAt = existing.QualityCheckedAt + merged.QualityCFRay = existing.QualityCFRay + } + } + } + + if err := s.proxyLatencyCache.SetProxyLatency(ctx, proxyID, &merged); err != nil { + logger.LegacyPrintf("service.admin", "Warning: store proxy latency cache failed: %v", err) + } +} diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index f1de60eb47..07b85ab827 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -2,31 +2,11 @@ package service import ( "context" - "database/sql" - "encoding/json" - "errors" - "fmt" - "io" - "log/slog" "net/http" - "sort" - "strconv" - "strings" "time" dbent "github.com/Wei-Shaw/sub2api/ent" - "github.com/Wei-Shaw/sub2api/ent/authidentity" - "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" - "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" - "github.com/Wei-Shaw/sub2api/internal/pkg/claude" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" - "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" - "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient" - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/Wei-Shaw/sub2api/internal/pkg/openai" - "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" - "github.com/Wei-Shaw/sub2api/internal/util/httputil" ) // AdminService interface defines admin management operations @@ -145,10 +125,13 @@ type CreateUserInput struct { Password string Username string Notes string + Role string // 空字符串表示使用默认角色(user);合法值 admin/user Balance *float64 Concurrency int RPMLimit int AllowedGroups []int64 + // ActorAdminID 执行本次操作的管理员ID(来自JWT),仅用于权限敏感操作的审计日志。 + ActorAdminID int64 } type UpdateUserInput struct { @@ -156,6 +139,7 @@ type UpdateUserInput struct { Password string Username *string Notes *string + Role string // 空字符串表示"未提供"(不修改);合法值 admin/user Balance *float64 // 使用指针区分"未提供"和"设置为0" Concurrency *int // 使用指针区分"未提供"和"设置为0" RPMLimit *int // 使用指针区分"未提供"和"设置为0" @@ -164,6 +148,8 @@ type UpdateUserInput struct { // GroupRates 用户专属分组倍率配置 // map[groupID]*rate,nil 表示删除该分组的专属倍率 GroupRates map[int64]*float64 + // ActorAdminID 执行本次操作的管理员ID(来自JWT),仅用于权限敏感操作的审计日志。 + ActorAdminID int64 } type AdminBindAuthIdentityInput struct { @@ -215,9 +201,14 @@ type CreateGroupInput struct { WeeklyLimitUSD *float64 // 周限额 (USD) MonthlyLimitUSD *float64 // 月限额 (USD) // 图片生成计费配置(仅 antigravity 平台使用) - AllowImageGeneration bool - ImageRateIndependent bool - ImageRateMultiplier *float64 + AllowImageGeneration bool + AllowBatchImageGeneration bool + ImageRateIndependent bool + ImageRateMultiplier *float64 + BatchImageDiscountMultiplier *float64 + BatchImageHoldMultiplier *float64 + VideoRateIndependent bool + VideoRateMultiplier *float64 // 高峰时段倍率配置(PeakRateMultiplier 为 nil 时按 1.0 处理) PeakRateEnabled bool PeakStart string @@ -226,6 +217,9 @@ type CreateGroupInput struct { ImagePrice1K *float64 ImagePrice2K *float64 ImagePrice4K *float64 + VideoPrice480P *float64 + VideoPrice720P *float64 + VideoPrice1080P *float64 ClaudeCodeOnly bool // 仅允许 Claude Code 客户端 FallbackGroupID *int64 // 降级分组 ID // 无效请求兜底分组 ID(仅 anthropic 平台使用) @@ -261,9 +255,14 @@ type UpdateGroupInput struct { WeeklyLimitUSD *float64 // 周限额 (USD) MonthlyLimitUSD *float64 // 月限额 (USD) // 图片生成计费配置(仅 antigravity 平台使用) - AllowImageGeneration *bool - ImageRateIndependent *bool - ImageRateMultiplier *float64 + AllowImageGeneration *bool + AllowBatchImageGeneration *bool + ImageRateIndependent *bool + ImageRateMultiplier *float64 + BatchImageDiscountMultiplier *float64 + BatchImageHoldMultiplier *float64 + VideoRateIndependent *bool + VideoRateMultiplier *float64 // 高峰时段倍率配置(nil 表示不修改) PeakRateEnabled *bool PeakStart *string @@ -272,6 +271,9 @@ type UpdateGroupInput struct { ImagePrice1K *float64 ImagePrice2K *float64 ImagePrice4K *float64 + VideoPrice480P *float64 + VideoPrice720P *float64 + VideoPrice1080P *float64 ClaudeCodeOnly *bool // 仅允许 Claude Code 客户端 FallbackGroupID *int64 // 降级分组 ID // 无效请求兜底分组 ID(仅 anthropic 平台使用) @@ -634,3723 +636,3 @@ func NewAdminService( runtimeBlocker: runtimeBlocker, } } - -// User management implementations -func (s *adminServiceImpl) ListUsers(ctx context.Context, page, pageSize int, filters UserListFilters, sortBy, sortOrder string) ([]User, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - users, result, err := s.userRepo.ListWithFilters(ctx, params, filters) - if err != nil { - return nil, 0, err - } - if len(users) > 0 { - userIDs := make([]int64, 0, len(users)) - for i := range users { - userIDs = append(userIDs, users[i].ID) - } - lastUsedByUserID, latestErr := s.userRepo.GetLatestUsedAtByUserIDs(ctx, userIDs) - if latestErr != nil { - logger.LegacyPrintf("service.admin", "failed to load user last_used_at in batch: err=%v", latestErr) - } else { - for i := range users { - users[i].LastUsedAt = lastUsedByUserID[users[i].ID] - } - } - } - // 批量加载用户专属分组倍率 - if s.userGroupRateRepo != nil && len(users) > 0 { - if batchRepo, ok := s.userGroupRateRepo.(userGroupRateBatchReader); ok { - userIDs := make([]int64, 0, len(users)) - for i := range users { - userIDs = append(userIDs, users[i].ID) - } - ratesByUser, err := batchRepo.GetByUserIDs(ctx, userIDs) - if err != nil { - logger.LegacyPrintf("service.admin", "failed to load user group rates in batch: err=%v", err) - s.loadUserGroupRatesOneByOne(ctx, users) - } else { - for i := range users { - if rates, ok := ratesByUser[users[i].ID]; ok { - users[i].GroupRates = rates - } - } - } - } else { - s.loadUserGroupRatesOneByOne(ctx, users) - } - } - return users, result.Total, nil -} - -func (s *adminServiceImpl) loadUserGroupRatesOneByOne(ctx context.Context, users []User) { - if s.userGroupRateRepo == nil { - return - } - for i := range users { - rates, err := s.userGroupRateRepo.GetByUserID(ctx, users[i].ID) - if err != nil { - logger.LegacyPrintf("service.admin", "failed to load user group rates: user_id=%d err=%v", users[i].ID, err) - continue - } - users[i].GroupRates = rates - } -} - -func (s *adminServiceImpl) GetUser(ctx context.Context, id int64) (*User, error) { - user, err := s.userRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - lastUsedAt, latestErr := s.userRepo.GetLatestUsedAtByUserID(ctx, id) - if latestErr != nil { - logger.LegacyPrintf("service.admin", "failed to load user last_used_at: user_id=%d err=%v", id, latestErr) - } else { - user.LastUsedAt = lastUsedAt - } - // 加载用户专属分组倍率 - if s.userGroupRateRepo != nil { - rates, err := s.userGroupRateRepo.GetByUserID(ctx, id) - if err != nil { - logger.LegacyPrintf("service.admin", "failed to load user group rates: user_id=%d err=%v", id, err) - } else { - user.GroupRates = rates - } - } - return user, nil -} - -func (s *adminServiceImpl) GetUserIncludeDeleted(ctx context.Context, id int64) (*User, error) { - return s.userRepo.GetByIDIncludeDeleted(ctx, id) -} - -func (s *adminServiceImpl) CreateUser(ctx context.Context, input *CreateUserInput) (*User, error) { - balance := 0.0 - if input.Balance != nil { - balance = *input.Balance - } else if s.settingService != nil { - balance = s.settingService.GetDefaultBalance(ctx) - } - - user := &User{ - Email: input.Email, - Username: input.Username, - Notes: input.Notes, - Role: RoleUser, // Always create as regular user, never admin - Balance: balance, - Concurrency: input.Concurrency, - RPMLimit: input.RPMLimit, - Status: StatusActive, - AllowedGroups: input.AllowedGroups, - } - if err := user.SetPassword(input.Password); err != nil { - return nil, err - } - if err := s.userRepo.Create(ctx, user); err != nil { - return nil, err - } - s.assignDefaultSubscriptions(ctx, user.ID) - return user, nil -} - -func (s *adminServiceImpl) assignDefaultSubscriptions(ctx context.Context, userID int64) { - if s.settingService == nil || s.defaultSubAssigner == nil || userID <= 0 { - return - } - items := s.settingService.GetDefaultSubscriptions(ctx) - for _, item := range items { - if _, _, err := s.defaultSubAssigner.AssignOrExtendSubscription(ctx, &AssignSubscriptionInput{ - UserID: userID, - GroupID: item.GroupID, - ValidityDays: item.ValidityDays, - Notes: "auto assigned by default user subscriptions setting", - }); err != nil { - logger.LegacyPrintf("service.admin", "failed to assign default subscription: user_id=%d group_id=%d err=%v", userID, item.GroupID, err) - } - } -} - -func (s *adminServiceImpl) UpdateUser(ctx context.Context, id int64, input *UpdateUserInput) (*User, error) { - // 校验用户专属分组倍率:必须 > 0(nil 合法,表示清除专属倍率) - if input.GroupRates != nil { - for groupID, rate := range input.GroupRates { - if rate != nil && *rate <= 0 { - return nil, fmt.Errorf("rate_multiplier must be > 0 (group_id=%d)", groupID) - } - } - } - - user, err := s.userRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - - // Protect admin users: cannot disable admin accounts - if user.Role == "admin" && input.Status == "disabled" { - return nil, errors.New("cannot disable admin user") - } - - oldConcurrency := user.Concurrency - oldStatus := user.Status - oldRole := user.Role - oldRPMLimit := user.RPMLimit - oldAllowedGroups := append([]int64(nil), user.AllowedGroups...) - - if input.Email != "" { - user.Email = input.Email - } - if input.Password != "" { - if err := user.SetPassword(input.Password); err != nil { - return nil, err - } - } - - if input.Username != nil { - user.Username = *input.Username - } - if input.Notes != nil { - user.Notes = *input.Notes - } - - if input.Status != "" { - user.Status = input.Status - } - - if input.Concurrency != nil { - user.Concurrency = *input.Concurrency - } - - if input.RPMLimit != nil { - user.RPMLimit = *input.RPMLimit - } - - if input.AllowedGroups != nil { - user.AllowedGroups = *input.AllowedGroups - } - - if err := s.userRepo.Update(ctx, user); err != nil { - return nil, err - } - - // 同步用户专属分组倍率 - if input.GroupRates != nil && s.userGroupRateRepo != nil { - if err := s.userGroupRateRepo.SyncUserGroupRates(ctx, user.ID, input.GroupRates); err != nil { - logger.LegacyPrintf("service.admin", "failed to sync user group rates: user_id=%d err=%v", user.ID, err) - } - } - - if s.authCacheInvalidator != nil { - // RPMLimit 直接参与 billing_cache_service.checkRPM 的三级级联, - // allowed_groups 参与 API Key 专属分组授权判断;不失效缓存会让修改在一个 L2 TTL 内失去效果。 - if user.Concurrency != oldConcurrency || user.Status != oldStatus || user.Role != oldRole || user.RPMLimit != oldRPMLimit || !sameInt64Set(user.AllowedGroups, oldAllowedGroups) { - s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, user.ID) - } - } - - concurrencyDiff := user.Concurrency - oldConcurrency - if concurrencyDiff != 0 { - code, err := GenerateRedeemCode() - if err != nil { - logger.LegacyPrintf("service.admin", "failed to generate adjustment redeem code: %v", err) - return user, nil - } - adjustmentRecord := &RedeemCode{ - Code: code, - Type: AdjustmentTypeAdminConcurrency, - Value: float64(concurrencyDiff), - Status: StatusUsed, - UsedBy: &user.ID, - } - now := time.Now() - adjustmentRecord.UsedAt = &now - if err := s.redeemCodeRepo.Create(ctx, adjustmentRecord); err != nil { - logger.LegacyPrintf("service.admin", "failed to create concurrency adjustment redeem code: %v", err) - } - } - - return user, nil -} - -func sameInt64Set(a, b []int64) bool { - if len(a) != len(b) { - return false - } - if len(a) == 0 { - return true - } - counts := make(map[int64]int, len(a)) - for _, v := range a { - counts[v]++ - } - for _, v := range b { - if counts[v] == 0 { - return false - } - counts[v]-- - } - return true -} - -func (s *adminServiceImpl) DeleteUser(ctx context.Context, id int64) error { - // Protect admin users: cannot delete admin accounts - user, err := s.userRepo.GetByID(ctx, id) - if err != nil { - return err - } - if user.Role == "admin" { - return errors.New("cannot delete admin user") - } - - apiKeys, err := s.listUserAPIKeysForDeletion(ctx, id) - if err != nil { - return err - } - - if s.entClient != nil { - tx, err := s.entClient.Tx(ctx) - if err != nil { - return err - } - defer func() { _ = tx.Rollback() }() - - opCtx := dbent.NewTxContext(ctx, tx) - if err := s.deleteUserWithAPIKeys(opCtx, id, apiKeys); err != nil { - return err - } - if err := tx.Commit(); err != nil { - return err - } - } else { - if err := s.deleteUserWithAPIKeys(ctx, id, apiKeys); err != nil { - return err - } - } - - if s.authCacheInvalidator != nil { - for _, key := range apiKeys { - if keyValue := strings.TrimSpace(key.Key); keyValue != "" { - s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, keyValue) - } - } - s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, id) - } - return nil -} - -func (s *adminServiceImpl) listUserAPIKeysForDeletion(ctx context.Context, userID int64) ([]APIKey, error) { - if s.apiKeyRepo == nil { - return nil, nil - } - - const pageSize = 1000 - keys := make([]APIKey, 0) - for page := 1; ; page++ { - batch, result, err := s.apiKeyRepo.ListByUserID(ctx, userID, pagination.PaginationParams{ - Page: page, - PageSize: pageSize, - SortBy: "id", - SortOrder: pagination.SortOrderAsc, - }, APIKeyListFilters{}) - if err != nil { - return nil, fmt.Errorf("list user api keys: %w", err) - } - keys = append(keys, batch...) - if len(batch) == 0 || len(batch) < pageSize || result == nil || int64(len(keys)) >= result.Total { - break - } - } - return keys, nil -} - -func (s *adminServiceImpl) deleteUserWithAPIKeys(ctx context.Context, userID int64, apiKeys []APIKey) error { - if s.apiKeyRepo != nil { - for _, key := range apiKeys { - if key.ID <= 0 { - continue - } - if err := s.apiKeyRepo.DeleteWithAudit(ctx, key.ID); err != nil { - logger.LegacyPrintf("service.admin", "delete user api key failed: user_id=%d api_key_id=%d err=%v", userID, key.ID, err) - return fmt.Errorf("delete user api key %d: %w", key.ID, err) - } - } - } - - if err := s.userRepo.Delete(ctx, userID); err != nil { - logger.LegacyPrintf("service.admin", "delete user failed: user_id=%d err=%v", userID, err) - return err - } - return nil -} - -func (s *adminServiceImpl) BatchUpdateConcurrency(ctx context.Context, userIDs []int64, value int, mode string) (int, error) { - cleaned := make([]int64, 0, len(userIDs)) - for _, uid := range userIDs { - if uid > 0 { - cleaned = append(cleaned, uid) - } - } - if len(cleaned) == 0 { - return 0, nil - } - - var affected int - var err error - switch mode { - case "set": - affected, err = s.userRepo.BatchSetConcurrency(ctx, cleaned, value) - case "add": - affected, err = s.userRepo.BatchAddConcurrency(ctx, cleaned, value) - default: - return 0, errors.New("invalid mode: must be 'set' or 'add'") - } - if err != nil { - return 0, err - } - - if s.authCacheInvalidator != nil { - for _, uid := range cleaned { - s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, uid) - } - } - return affected, nil -} - -func (s *adminServiceImpl) UpdateUserBalance(ctx context.Context, userID int64, balance float64, operation string, notes string) (*User, error) { - user, err := s.userRepo.GetByID(ctx, userID) - if err != nil { - return nil, err - } - - oldBalance := user.Balance - - switch operation { - case "set": - user.Balance = balance - case "add": - user.Balance += balance - case "subtract": - user.Balance -= balance - } - - if user.Balance < 0 { - return nil, fmt.Errorf("balance cannot be negative, current balance: %.2f, requested operation would result in: %.2f", oldBalance, user.Balance) - } - - if err := s.userRepo.Update(ctx, user); err != nil { - return nil, err - } - balanceDiff := user.Balance - oldBalance - if s.authCacheInvalidator != nil && balanceDiff != 0 { - s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID) - } - - if s.billingCacheService != nil { - go func() { - cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if err := s.billingCacheService.InvalidateUserBalance(cacheCtx, userID); err != nil { - logger.LegacyPrintf("service.admin", "invalidate user balance cache failed: user_id=%d err=%v", userID, err) - } - }() - } - - if balanceDiff != 0 { - code, err := GenerateRedeemCode() - if err != nil { - logger.LegacyPrintf("service.admin", "failed to generate adjustment redeem code: %v", err) - return user, nil - } - - adjustmentRecord := &RedeemCode{ - Code: code, - Type: AdjustmentTypeAdminBalance, - Value: balanceDiff, - Status: StatusUsed, - UsedBy: &user.ID, - Notes: notes, - } - now := time.Now() - adjustmentRecord.UsedAt = &now - - if err := s.redeemCodeRepo.Create(ctx, adjustmentRecord); err != nil { - logger.LegacyPrintf("service.admin", "failed to create balance adjustment redeem code: %v", err) - } - } - - return user, nil -} - -func (s *adminServiceImpl) GetUserAPIKeys(ctx context.Context, userID int64, page, pageSize int, sortBy, sortOrder string) ([]APIKey, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - keys, result, err := s.apiKeyRepo.ListByUserID(ctx, userID, params, APIKeyListFilters{}) - if err != nil { - return nil, 0, err - } - return keys, result.Total, nil -} - -func (s *adminServiceImpl) GetUserRPMStatus(ctx context.Context, userID int64) (*UserRPMStatus, error) { - if s.userRPMCache == nil { - return nil, ErrRPMStatusUnavailable - } - - user, err := s.userRepo.GetByID(ctx, userID) - if err != nil { - return nil, err - } - - userRPMUsed, err := s.userRPMCache.GetUserRPM(ctx, userID) - if err != nil { - logger.LegacyPrintf("service.admin", "failed to get user rpm: user_id=%d err=%v", userID, err) - } - - keys, _, err := s.GetUserAPIKeys(ctx, userID, 1, 1000, "", "") - if err != nil { - return nil, err - } - - groupIDSet := make(map[int64]struct{}) - for _, key := range keys { - if key.GroupID != nil && *key.GroupID > 0 { - groupIDSet[*key.GroupID] = struct{}{} - } - } - - groupIDs := make([]int64, 0, len(groupIDSet)) - for groupID := range groupIDSet { - groupIDs = append(groupIDs, groupID) - } - sort.Slice(groupIDs, func(i, j int) bool { return groupIDs[i] < groupIDs[j] }) - - var perGroup []UserGroupRPMStatus - for _, groupID := range groupIDs { - used, getErr := s.userRPMCache.GetUserGroupRPM(ctx, userID, groupID) - if getErr != nil { - logger.LegacyPrintf("service.admin", "failed to get user group rpm: user_id=%d group_id=%d err=%v", userID, groupID, getErr) - } - - entry := UserGroupRPMStatus{ - GroupID: groupID, - Used: used, - } - - if s.groupRepo != nil { - if group, groupErr := s.groupRepo.GetByIDLite(ctx, groupID); groupErr == nil && group != nil { - entry.GroupName = group.Name - entry.Limit = group.RPMLimit - entry.Source = "group" - } else if groupErr != nil { - logger.LegacyPrintf("service.admin", "failed to get group rpm status metadata: group_id=%d err=%v", groupID, groupErr) - } - } - - if s.userGroupRateRepo != nil { - override, overrideErr := s.userGroupRateRepo.GetRPMOverrideByUserAndGroup(ctx, userID, groupID) - if overrideErr != nil { - logger.LegacyPrintf("service.admin", "failed to get rpm override: user_id=%d group_id=%d err=%v", userID, groupID, overrideErr) - } else if override != nil { - entry.Limit = *override - entry.Source = "override" - } - } - - perGroup = append(perGroup, entry) - } - - return &UserRPMStatus{ - UserRPMUsed: userRPMUsed, - UserRPMLimit: user.RPMLimit, - PerGroup: perGroup, - }, nil -} - -func (s *adminServiceImpl) GetUserUsageStats(ctx context.Context, userID int64, period string) (any, error) { - // Return mock data for now - return map[string]any{ - "period": period, - "total_requests": 0, - "total_cost": 0.0, - "total_tokens": 0, - "avg_duration_ms": 0, - }, nil -} - -// GetUserBalanceHistory returns paginated balance/concurrency change records for a user. -func (s *adminServiceImpl) GetUserBalanceHistory(ctx context.Context, userID int64, page, pageSize int, codeType string) ([]RedeemCode, int64, float64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize} - if codeType == RedeemTypeAffiliateBalance { - codes, total, err := s.listAffiliateBalanceHistory(ctx, userID, params) - if err != nil { - return nil, 0, 0, err - } - totalRecharged, err := s.redeemCodeRepo.SumPositiveBalanceByUser(ctx, userID) - if err != nil { - return nil, 0, 0, err - } - return codes, total, totalRecharged, nil - } - - if codeType == "" { - return s.getAllUserBalanceHistory(ctx, userID, params) - } - - codes, result, err := s.redeemCodeRepo.ListByUserPaginated(ctx, userID, params, codeType) - if err != nil { - return nil, 0, 0, err - } - total := result.Total - // Aggregate total recharged amount (only once, regardless of type filter) - totalRecharged, err := s.redeemCodeRepo.SumPositiveBalanceByUser(ctx, userID) - if err != nil { - return nil, 0, 0, err - } - return codes, total, totalRecharged, nil -} - -func (s *adminServiceImpl) getAllUserBalanceHistory(ctx context.Context, userID int64, params pagination.PaginationParams) ([]RedeemCode, int64, float64, error) { - needed := params.Offset() + params.Limit() - if needed < params.Limit() { - needed = params.Limit() - } - - redeemCodes, redeemTotal, err := s.listRedeemBalanceHistoryForMerge(ctx, userID, needed) - if err != nil { - return nil, 0, 0, err - } - affiliateCodes, affiliateTotal, err := s.listAffiliateBalanceHistoryForMerge(ctx, userID, needed) - if err != nil { - return nil, 0, 0, err - } - codes := mergeBalanceHistoryCodes(redeemCodes, affiliateCodes, params) - - totalRecharged, err := s.redeemCodeRepo.SumPositiveBalanceByUser(ctx, userID) - if err != nil { - return nil, 0, 0, err - } - return codes, redeemTotal + affiliateTotal, totalRecharged, nil -} - -func (s *adminServiceImpl) listRedeemBalanceHistoryForMerge(ctx context.Context, userID int64, needed int) ([]RedeemCode, int64, error) { - if needed <= 0 { - return nil, 0, nil - } - - var ( - out []RedeemCode - total int64 - ) - for page := 1; len(out) < needed; page++ { - params := pagination.PaginationParams{Page: page, PageSize: 1000} - codes, result, err := s.redeemCodeRepo.ListByUserPaginated(ctx, userID, params, "") - if err != nil { - return nil, 0, err - } - if result != nil { - total = result.Total - } - out = append(out, codes...) - if len(codes) < params.Limit() || int64(len(out)) >= total { - break - } - } - if len(out) > needed { - out = out[:needed] - } - return out, total, nil -} - -func (s *adminServiceImpl) listAffiliateBalanceHistoryForMerge(ctx context.Context, userID int64, needed int) ([]RedeemCode, int64, error) { - if needed <= 0 { - return nil, 0, nil - } - - var ( - out []RedeemCode - total int64 - ) - for page := 1; len(out) < needed; page++ { - params := pagination.PaginationParams{Page: page, PageSize: 1000} - codes, currentTotal, err := s.listAffiliateBalanceHistory(ctx, userID, params) - if err != nil { - return nil, 0, err - } - total = currentTotal - out = append(out, codes...) - if len(codes) < params.Limit() || int64(len(out)) >= total { - break - } - } - if len(out) > needed { - out = out[:needed] - } - return out, total, nil -} - -func (s *adminServiceImpl) listAffiliateBalanceHistory(ctx context.Context, userID int64, params pagination.PaginationParams) ([]RedeemCode, int64, error) { - if s == nil || s.entClient == nil || userID <= 0 { - return nil, 0, nil - } - - rows, err := s.entClient.QueryContext(ctx, ` -SELECT id, - amount::double precision, - created_at -FROM user_affiliate_ledger -WHERE user_id = $1 - AND action = 'transfer' -ORDER BY created_at DESC, id DESC -OFFSET $2 -LIMIT $3`, userID, params.Offset(), params.Limit()) - if err != nil { - return nil, 0, err - } - defer func() { _ = rows.Close() }() - - codes := make([]RedeemCode, 0, params.Limit()) - for rows.Next() { - var id int64 - var amount float64 - var createdAt time.Time - if err := rows.Scan(&id, &amount, &createdAt); err != nil { - return nil, 0, err - } - usedBy := userID - usedAt := createdAt - codes = append(codes, RedeemCode{ - ID: -id, - Code: fmt.Sprintf("AFF-%d", id), - Type: RedeemTypeAffiliateBalance, - Value: amount, - Status: StatusUsed, - UsedBy: &usedBy, - UsedAt: &usedAt, - CreatedAt: createdAt, - }) - } - if err := rows.Err(); err != nil { - return nil, 0, err - } - - total, err := countAffiliateBalanceHistory(ctx, s.entClient, userID) - if err != nil { - return nil, 0, err - } - return codes, total, nil -} - -func countAffiliateBalanceHistory(ctx context.Context, client *dbent.Client, userID int64) (int64, error) { - rows, err := client.QueryContext(ctx, ` -SELECT COUNT(*) -FROM user_affiliate_ledger -WHERE user_id = $1 - AND action = 'transfer'`, userID) - if err != nil { - return 0, err - } - defer func() { _ = rows.Close() }() - - var total sql.NullInt64 - if rows.Next() { - if err := rows.Scan(&total); err != nil { - return 0, err - } - } - if err := rows.Err(); err != nil { - return 0, err - } - if !total.Valid { - return 0, nil - } - return total.Int64, nil -} - -func mergeBalanceHistoryCodes(redeemCodes, affiliateCodes []RedeemCode, params pagination.PaginationParams) []RedeemCode { - combined := append(append([]RedeemCode{}, redeemCodes...), affiliateCodes...) - sort.SliceStable(combined, func(i, j int) bool { - return redeemCodeHistoryTime(combined[i]).After(redeemCodeHistoryTime(combined[j])) - }) - offset := params.Offset() - if offset >= len(combined) { - return []RedeemCode{} - } - end := offset + params.Limit() - if end > len(combined) { - end = len(combined) - } - return combined[offset:end] -} - -func redeemCodeHistoryTime(code RedeemCode) time.Time { - if code.UsedAt != nil { - return *code.UsedAt - } - return code.CreatedAt -} - -func (s *adminServiceImpl) BindUserAuthIdentity(ctx context.Context, userID int64, input AdminBindAuthIdentityInput) (*AdminBoundAuthIdentity, error) { - if userID <= 0 { - return nil, infraerrors.BadRequest("INVALID_INPUT", "user_id must be greater than 0") - } - if s == nil || s.entClient == nil || s.userRepo == nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_UNAVAILABLE", "auth identity binding service is unavailable") - } - if _, err := s.userRepo.GetByID(ctx, userID); err != nil { - return nil, err - } - - providerType := normalizeAdminAuthIdentityProviderType(input.ProviderType) - providerKey := strings.TrimSpace(input.ProviderKey) - providerSubject := strings.TrimSpace(input.ProviderSubject) - if providerType == "" { - return nil, infraerrors.BadRequest("INVALID_INPUT", "provider_type must be one of email, linuxdo, oidc, wechat, or dingtalk") - } - if providerKey == "" || providerSubject == "" { - return nil, infraerrors.BadRequest("INVALID_INPUT", "provider_type, provider_key, and provider_subject are required") - } - canonicalProviderKey := canonicalAdminAuthIdentityProviderKey(providerType, "", providerKey) - compatibleProviderKeys := compatibleAdminAuthIdentityProviderKeys(providerType, providerKey) - - var issuer *string - if input.Issuer != nil { - trimmed := strings.TrimSpace(*input.Issuer) - if trimmed != "" { - issuer = &trimmed - } - } - - channelInput := normalizeAdminBindChannelInput(input.Channel) - if input.Channel != nil && channelInput == nil { - return nil, infraerrors.BadRequest("INVALID_INPUT", "channel, channel_app_id, and channel_subject are required when channel binding is provided") - } - - verifiedAt := time.Now().UTC() - tx, err := s.entClient.Tx(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_TX_FAILED", "failed to start auth identity bind transaction").WithCause(err) - } - defer func() { _ = tx.Rollback() }() - - identityRecords, err := tx.AuthIdentity.Query(). - Where( - authidentity.ProviderTypeEQ(providerType), - authidentity.ProviderKeyIn(compatibleProviderKeys...), - authidentity.ProviderSubjectEQ(providerSubject), - ). - All(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_LOOKUP_FAILED", "failed to inspect auth identity ownership").WithCause(err) - } - if hasAdminAuthIdentityOwnershipConflict(identityRecords, userID) { - return nil, infraerrors.Conflict("AUTH_IDENTITY_OWNERSHIP_CONFLICT", "auth identity already belongs to another user") - } - identity := selectOwnedAdminAuthIdentity(identityRecords, userID) - - if identity == nil { - create := tx.AuthIdentity.Create(). - SetUserID(userID). - SetProviderType(providerType). - SetProviderKey(canonicalProviderKey). - SetProviderSubject(providerSubject). - SetVerifiedAt(verifiedAt) - if issuer != nil { - create = create.SetIssuer(*issuer) - } - if input.Metadata != nil { - create = create.SetMetadata(cloneAdminAuthIdentityMetadata(input.Metadata)) - } - identity, err = create.Save(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_SAVE_FAILED", "failed to save auth identity").WithCause(err) - } - } else { - update := tx.AuthIdentity.UpdateOneID(identity.ID). - SetVerifiedAt(verifiedAt). - SetProviderKey(canonicalProviderKey) - if issuer != nil { - update = update.SetIssuer(*issuer) - } - if input.Metadata != nil { - update = update.SetMetadata(cloneAdminAuthIdentityMetadata(input.Metadata)) - } - identity, err = update.Save(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_SAVE_FAILED", "failed to save auth identity").WithCause(err) - } - } - - var channel *dbent.AuthIdentityChannel - if channelInput != nil { - channelRecords, err := tx.AuthIdentityChannel.Query(). - Where( - authidentitychannel.ProviderTypeEQ(providerType), - authidentitychannel.ProviderKeyIn(compatibleProviderKeys...), - authidentitychannel.ChannelEQ(channelInput.Channel), - authidentitychannel.ChannelAppIDEQ(channelInput.ChannelAppID), - authidentitychannel.ChannelSubjectEQ(channelInput.ChannelSubject), - ). - WithIdentity(). - All(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_CHANNEL_LOOKUP_FAILED", "failed to inspect auth identity channel ownership").WithCause(err) - } - if hasAdminAuthIdentityChannelOwnershipConflict(channelRecords, userID) { - return nil, infraerrors.Conflict("AUTH_IDENTITY_CHANNEL_OWNERSHIP_CONFLICT", "auth identity channel already belongs to another user") - } - channel = selectOwnedAdminAuthIdentityChannel(channelRecords, userID) - if channel == nil { - create := tx.AuthIdentityChannel.Create(). - SetIdentityID(identity.ID). - SetProviderType(providerType). - SetProviderKey(canonicalProviderKey). - SetChannel(channelInput.Channel). - SetChannelAppID(channelInput.ChannelAppID). - SetChannelSubject(channelInput.ChannelSubject) - if channelInput.Metadata != nil { - create = create.SetMetadata(cloneAdminAuthIdentityMetadata(channelInput.Metadata)) - } - channel, err = create.Save(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_CHANNEL_SAVE_FAILED", "failed to save auth identity channel").WithCause(err) - } - } else { - update := tx.AuthIdentityChannel.UpdateOneID(channel.ID). - SetIdentityID(identity.ID). - SetProviderKey(canonicalProviderKey) - if channelInput.Metadata != nil { - update = update.SetMetadata(cloneAdminAuthIdentityMetadata(channelInput.Metadata)) - } - channel, err = update.Save(ctx) - if err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_CHANNEL_SAVE_FAILED", "failed to save auth identity channel").WithCause(err) - } - } - } - - if err := tx.Commit(); err != nil { - return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_COMMIT_FAILED", "failed to commit auth identity bind").WithCause(err) - } - return buildAdminBoundAuthIdentity(identity, channel), nil -} - -func compatibleAdminAuthIdentityProviderKeys(providerType, providerKey string) []string { - providerType = strings.TrimSpace(strings.ToLower(providerType)) - providerKey = strings.TrimSpace(providerKey) - if providerKey == "" { - return []string{providerKey} - } - if providerType != "wechat" { - return []string{providerKey} - } - - keys := []string{providerKey} - if !strings.EqualFold(providerKey, "wechat-main") { - keys = append(keys, "wechat-main") - } - if !strings.EqualFold(providerKey, "wechat") { - keys = append(keys, "wechat") - } - return keys -} - -func canonicalAdminAuthIdentityProviderKey(providerType, existingKey, requestedKey string) string { - providerType = strings.TrimSpace(strings.ToLower(providerType)) - existingKey = strings.TrimSpace(existingKey) - requestedKey = strings.TrimSpace(requestedKey) - if providerType != "wechat" { - if requestedKey != "" { - return requestedKey - } - return existingKey - } - if strings.EqualFold(existingKey, "wechat") || strings.EqualFold(existingKey, "wechat-main") || strings.EqualFold(requestedKey, "wechat-main") { - return "wechat-main" - } - if requestedKey != "" { - return requestedKey - } - return existingKey -} - -func adminAuthIdentityProviderKeyRank(providerType, providerKey string) int { - providerType = strings.TrimSpace(strings.ToLower(providerType)) - providerKey = strings.TrimSpace(providerKey) - if providerType != "wechat" { - return 0 - } - switch { - case strings.EqualFold(providerKey, "wechat-main"): - return 0 - case strings.EqualFold(providerKey, "wechat"): - return 2 - default: - return 1 - } -} - -func selectOwnedAdminAuthIdentity(records []*dbent.AuthIdentity, userID int64) *dbent.AuthIdentity { - var selected *dbent.AuthIdentity - for _, record := range records { - if record.UserID != userID { - continue - } - if selected == nil || adminAuthIdentityProviderKeyRank(record.ProviderType, record.ProviderKey) < adminAuthIdentityProviderKeyRank(selected.ProviderType, selected.ProviderKey) { - selected = record - } - } - return selected -} - -func hasAdminAuthIdentityOwnershipConflict(records []*dbent.AuthIdentity, userID int64) bool { - for _, record := range records { - if record.UserID != userID { - return true - } - } - return false -} - -func selectOwnedAdminAuthIdentityChannel(records []*dbent.AuthIdentityChannel, userID int64) *dbent.AuthIdentityChannel { - var selected *dbent.AuthIdentityChannel - for _, record := range records { - if record.Edges.Identity == nil || record.Edges.Identity.UserID != userID { - continue - } - if selected == nil || adminAuthIdentityProviderKeyRank(record.ProviderType, record.ProviderKey) < adminAuthIdentityProviderKeyRank(selected.ProviderType, selected.ProviderKey) { - selected = record - } - } - return selected -} - -func hasAdminAuthIdentityChannelOwnershipConflict(records []*dbent.AuthIdentityChannel, userID int64) bool { - for _, record := range records { - if record.Edges.Identity != nil && record.Edges.Identity.UserID != userID { - return true - } - } - return false -} - -func normalizeAdminBindChannelInput(input *AdminBindAuthIdentityChannelInput) *AdminBindAuthIdentityChannelInput { - if input == nil { - return nil - } - channel := &AdminBindAuthIdentityChannelInput{ - Channel: strings.TrimSpace(input.Channel), - ChannelAppID: strings.TrimSpace(input.ChannelAppID), - ChannelSubject: strings.TrimSpace(input.ChannelSubject), - Metadata: cloneAdminAuthIdentityMetadata(input.Metadata), - } - if channel.Channel == "" || channel.ChannelAppID == "" || channel.ChannelSubject == "" { - return nil - } - return channel -} - -func normalizeAdminAuthIdentityProviderType(input string) string { - switch strings.ToLower(strings.TrimSpace(input)) { - case "email": - return "email" - case "linuxdo": - return "linuxdo" - case "oidc": - return "oidc" - case "wechat": - return "wechat" - case "dingtalk": - return "dingtalk" - default: - return "" - } -} - -func buildAdminBoundAuthIdentity(identity *dbent.AuthIdentity, channel *dbent.AuthIdentityChannel) *AdminBoundAuthIdentity { - if identity == nil { - return nil - } - result := &AdminBoundAuthIdentity{ - UserID: identity.UserID, - ProviderType: strings.TrimSpace(identity.ProviderType), - ProviderKey: strings.TrimSpace(identity.ProviderKey), - ProviderSubject: strings.TrimSpace(identity.ProviderSubject), - VerifiedAt: identity.VerifiedAt, - Issuer: identity.Issuer, - Metadata: cloneAdminAuthIdentityMetadata(identity.Metadata), - CreatedAt: identity.CreatedAt, - UpdatedAt: identity.UpdatedAt, - } - if channel != nil { - result.Channel = &AdminBoundAuthIdentityChannel{ - Channel: strings.TrimSpace(channel.Channel), - ChannelAppID: strings.TrimSpace(channel.ChannelAppID), - ChannelSubject: strings.TrimSpace(channel.ChannelSubject), - Metadata: cloneAdminAuthIdentityMetadata(channel.Metadata), - CreatedAt: channel.CreatedAt, - UpdatedAt: channel.UpdatedAt, - } - } - return result -} - -func cloneAdminAuthIdentityMetadata(input map[string]any) map[string]any { - if input == nil { - return nil - } - if len(input) == 0 { - return map[string]any{} - } - data, err := json.Marshal(input) - if err != nil { - out := make(map[string]any, len(input)) - for key, value := range input { - out[key] = value - } - return out - } - var out map[string]any - if err := json.Unmarshal(data, &out); err != nil { - out = make(map[string]any, len(input)) - for key, value := range input { - out[key] = value - } - } - return out -} - -// Group management implementations -func (s *adminServiceImpl) ListGroups(ctx context.Context, page, pageSize int, platform, status, search string, isExclusive *bool, sortBy, sortOrder string) ([]Group, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - groups, result, err := s.groupRepo.ListWithFilters(ctx, params, platform, status, search, isExclusive) - if err != nil { - return nil, 0, err - } - return groups, result.Total, nil -} - -func (s *adminServiceImpl) GetAllGroups(ctx context.Context) ([]Group, error) { - return s.groupRepo.ListActive(ctx) -} - -func (s *adminServiceImpl) GetAllGroupsByPlatform(ctx context.Context, platform string) ([]Group, error) { - return s.groupRepo.ListActiveByPlatform(ctx, platform) -} - -func (s *adminServiceImpl) GetAllGroupsIncludingInactive(ctx context.Context) ([]Group, error) { - // ListWithFilters with empty status = no status filter, so active + disabled groups are returned. - // PageSize 10000 is intentionally large; group count is O(dozens) in practice. - groups, _, err := s.groupRepo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10000}, "", "", "", nil) - return groups, err -} - -func (s *adminServiceImpl) GetGroup(ctx context.Context, id int64) (*Group, error) { - return s.groupRepo.GetByID(ctx, id) -} - -func (s *adminServiceImpl) GetGroupModelsListCandidates(ctx context.Context, id int64, platform string) ([]string, error) { - platform = strings.TrimSpace(platform) - if id > 0 { - group, err := s.groupRepo.GetByIDLite(ctx, id) - if err != nil { - return nil, err - } - if platform == "" { - platform = group.Platform - } - } - if platform == "" { - platform = PlatformAnthropic - } - - candidates := defaultModelsListCandidateIDs(platform) - if id <= 0 || s.accountRepo == nil { - return candidates, nil - } - - accounts, err := s.accountRepo.ListSchedulableByGroupID(ctx, id) - if err != nil { - return nil, err - } - - seen := make(map[string]struct{}, len(candidates)) - for _, model := range candidates { - seen[model] = struct{}{} - } - for _, acc := range accounts { - if acc.Platform != platform { - continue - } - for model := range acc.GetModelMapping() { - model = strings.TrimSpace(model) - if model == "" { - continue - } - if _, ok := seen[model]; ok { - continue - } - seen[model] = struct{}{} - candidates = append(candidates, model) - } - } - return candidates, nil -} - -func defaultModelsListCandidateIDs(platform string) []string { - switch platform { - case PlatformOpenAI: - return openai.DefaultModelIDs() - case PlatformGemini: - ids := make([]string, 0, len(geminicli.DefaultModels)) - for _, model := range geminicli.DefaultModels { - ids = append(ids, model.ID) - } - return ids - case PlatformAntigravity: - models := antigravity.DefaultModels() - ids := make([]string, 0, len(models)) - for _, model := range models { - ids = append(ids, model.ID) - } - return ids - case PlatformGrok: - return xai.DefaultModelIDs() - default: - ids := make([]string, 0, len(claude.DefaultModels)) - for _, model := range claude.DefaultModels { - ids = append(ids, model.ID) - } - return ids - } -} - -func defaultAllowImageGenerationForPlatform(platform string) bool { - // Grok image and video generation routes share the legacy image-generation gate. - // Older clients send the false zero value, so Grok groups must default enabled. - return platform == PlatformGrok -} - -func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupInput) (*Group, error) { - if input.RateMultiplier <= 0 { - return nil, errors.New("rate_multiplier must be > 0") - } - - platform := input.Platform - if platform == "" { - platform = PlatformAnthropic - } - - subscriptionType := input.SubscriptionType - if subscriptionType == "" { - subscriptionType = SubscriptionTypeStandard - } - - // 限额字段:nil/负数 表示"无限制",0 表示"不允许用量",正数表示具体限额 - dailyLimit := normalizeLimit(input.DailyLimitUSD) - weeklyLimit := normalizeLimit(input.WeeklyLimitUSD) - monthlyLimit := normalizeLimit(input.MonthlyLimitUSD) - - // 图片价格:负数表示清除(使用默认价格),0 保留(表示免费) - imagePrice1K := normalizePrice(input.ImagePrice1K) - imagePrice2K := normalizePrice(input.ImagePrice2K) - imagePrice4K := normalizePrice(input.ImagePrice4K) - imageRateMultiplier := 1.0 - if input.ImageRateMultiplier != nil { - if *input.ImageRateMultiplier < 0 { - return nil, errors.New("image_rate_multiplier must be >= 0") - } - imageRateMultiplier = *input.ImageRateMultiplier - } - - peakRateMultiplier := 1.0 - if input.PeakRateMultiplier != nil { - peakRateMultiplier = *input.PeakRateMultiplier - } - // 先归一化(非订阅分组清空高峰配置、清洗停用状态下的脏字段)再校验,与 UpdateGroup 同一收口。 - peakRateEnabled, peakStart, peakEnd, peakRateMultiplier := NormalizePeakRateConfig(subscriptionType, input.PeakRateEnabled, input.PeakStart, input.PeakEnd, peakRateMultiplier) - if err := ValidatePeakRateConfig(subscriptionType, peakRateEnabled, peakStart, peakEnd, peakRateMultiplier); err != nil { - return nil, err - } - - // 校验降级分组 - if input.FallbackGroupID != nil { - if err := s.validateFallbackGroup(ctx, 0, *input.FallbackGroupID); err != nil { - return nil, err - } - } - fallbackOnInvalidRequest := input.FallbackGroupIDOnInvalidRequest - if fallbackOnInvalidRequest != nil && *fallbackOnInvalidRequest <= 0 { - fallbackOnInvalidRequest = nil - } - // 校验无效请求兜底分组 - if fallbackOnInvalidRequest != nil { - if err := s.validateFallbackGroupOnInvalidRequest(ctx, 0, platform, subscriptionType, *fallbackOnInvalidRequest); err != nil { - return nil, err - } - } - - // MCPXMLInject:默认为 true,仅当显式传入 false 时关闭 - mcpXMLInject := true - if input.MCPXMLInject != nil { - mcpXMLInject = *input.MCPXMLInject - } - - allowImageGeneration := input.AllowImageGeneration || defaultAllowImageGenerationForPlatform(platform) - - // 如果指定了复制账号的源分组,先获取账号 ID 列表 - var accountIDsToCopy []int64 - if len(input.CopyAccountsFromGroupIDs) > 0 { - // 去重源分组 IDs - seen := make(map[int64]struct{}) - uniqueSourceGroupIDs := make([]int64, 0, len(input.CopyAccountsFromGroupIDs)) - for _, srcGroupID := range input.CopyAccountsFromGroupIDs { - if _, exists := seen[srcGroupID]; !exists { - seen[srcGroupID] = struct{}{} - uniqueSourceGroupIDs = append(uniqueSourceGroupIDs, srcGroupID) - } - } - - // 校验源分组的平台是否与新分组一致 - for _, srcGroupID := range uniqueSourceGroupIDs { - srcGroup, err := s.groupRepo.GetByIDLite(ctx, srcGroupID) - if err != nil { - return nil, fmt.Errorf("source group %d not found: %w", srcGroupID, err) - } - if srcGroup.Platform != platform { - return nil, fmt.Errorf("source group %d platform mismatch: expected %s, got %s", srcGroupID, platform, srcGroup.Platform) - } - } - - // 获取所有源分组的账号(去重) - var err error - accountIDsToCopy, err = s.groupRepo.GetAccountIDsByGroupIDs(ctx, uniqueSourceGroupIDs) - if err != nil { - return nil, fmt.Errorf("failed to get accounts from source groups: %w", err) - } - } - - group := &Group{ - Name: input.Name, - Description: input.Description, - Platform: platform, - RateMultiplier: input.RateMultiplier, - IsExclusive: input.IsExclusive, - Status: StatusActive, - SubscriptionType: subscriptionType, - DailyLimitUSD: dailyLimit, - WeeklyLimitUSD: weeklyLimit, - MonthlyLimitUSD: monthlyLimit, - AllowImageGeneration: allowImageGeneration, - ImageRateIndependent: input.ImageRateIndependent, - ImageRateMultiplier: imageRateMultiplier, - PeakRateEnabled: peakRateEnabled, - PeakStart: peakStart, - PeakEnd: peakEnd, - PeakRateMultiplier: peakRateMultiplier, - ImagePrice1K: imagePrice1K, - ImagePrice2K: imagePrice2K, - ImagePrice4K: imagePrice4K, - ClaudeCodeOnly: input.ClaudeCodeOnly, - FallbackGroupID: input.FallbackGroupID, - FallbackGroupIDOnInvalidRequest: fallbackOnInvalidRequest, - ModelRouting: input.ModelRouting, - MCPXMLInject: mcpXMLInject, - SupportedModelScopes: input.SupportedModelScopes, - AllowMessagesDispatch: input.AllowMessagesDispatch, - RequireOAuthOnly: input.RequireOAuthOnly, - RequirePrivacySet: input.RequirePrivacySet, - DefaultMappedModel: input.DefaultMappedModel, - MessagesDispatchModelConfig: normalizeOpenAIMessagesDispatchModelConfig(input.MessagesDispatchModelConfig), - ModelsListConfig: normalizeGroupModelsListConfig(input.ModelsListConfig), - RPMLimit: input.RPMLimit, - } - sanitizeGroupMessagesDispatchFields(group) - if err := s.groupRepo.Create(ctx, group); err != nil { - return nil, err - } - - // require_oauth_only: 过滤掉 apikey 类型账号 - if group.RequireOAuthOnly && (group.Platform == PlatformOpenAI || group.Platform == PlatformAntigravity || group.Platform == PlatformAnthropic || group.Platform == PlatformGemini || group.Platform == PlatformGrok) && len(accountIDsToCopy) > 0 { - accounts, err := s.accountRepo.GetByIDs(ctx, accountIDsToCopy) - if err != nil { - return nil, fmt.Errorf("failed to fetch accounts for oauth filter: %w", err) - } - oauthIDs := make(map[int64]struct{}, len(accounts)) - for _, acc := range accounts { - if acc.Type != AccountTypeAPIKey { - oauthIDs[acc.ID] = struct{}{} - } - } - var filtered []int64 - for _, aid := range accountIDsToCopy { - if _, ok := oauthIDs[aid]; ok { - filtered = append(filtered, aid) - } - } - accountIDsToCopy = filtered - } - - // 如果有需要复制的账号,绑定到新分组 - if len(accountIDsToCopy) > 0 { - if err := s.groupRepo.BindAccountsToGroup(ctx, group.ID, accountIDsToCopy); err != nil { - return nil, fmt.Errorf("failed to bind accounts to new group: %w", err) - } - group.AccountCount = int64(len(accountIDsToCopy)) - } - - return group, nil -} - -// normalizeLimit 将负数转换为 nil(表示无限制),0 保留(表示限额为零) -func normalizeLimit(limit *float64) *float64 { - if limit == nil || *limit < 0 { - return nil - } - return limit -} - -// normalizePrice 将负数转换为 nil(表示使用默认价格),0 保留(表示免费) -func normalizePrice(price *float64) *float64 { - if price == nil || *price < 0 { - return nil - } - return price -} - -// validateFallbackGroup 校验降级分组的有效性 -// currentGroupID: 当前分组 ID(新建时为 0) -// fallbackGroupID: 降级分组 ID -func (s *adminServiceImpl) validateFallbackGroup(ctx context.Context, currentGroupID, fallbackGroupID int64) error { - // 不能将自己设置为降级分组 - if currentGroupID > 0 && currentGroupID == fallbackGroupID { - return fmt.Errorf("cannot set self as fallback group") - } - - visited := map[int64]struct{}{} - nextID := fallbackGroupID - for { - if _, seen := visited[nextID]; seen { - return fmt.Errorf("fallback group cycle detected") - } - visited[nextID] = struct{}{} - if currentGroupID > 0 && nextID == currentGroupID { - return fmt.Errorf("fallback group cycle detected") - } - - // 检查降级分组是否存在 - fallbackGroup, err := s.groupRepo.GetByIDLite(ctx, nextID) - if err != nil { - return fmt.Errorf("fallback group not found: %w", err) - } - - // 降级分组不能启用 claude_code_only,否则会造成死循环 - if nextID == fallbackGroupID && fallbackGroup.ClaudeCodeOnly { - return fmt.Errorf("fallback group cannot have claude_code_only enabled") - } - - if fallbackGroup.FallbackGroupID == nil { - return nil - } - nextID = *fallbackGroup.FallbackGroupID - } -} - -// validateFallbackGroupOnInvalidRequest 校验无效请求兜底分组的有效性 -// currentGroupID: 当前分组 ID(新建时为 0) -// platform/subscriptionType: 当前分组的有效平台/订阅类型 -// fallbackGroupID: 兜底分组 ID -func (s *adminServiceImpl) validateFallbackGroupOnInvalidRequest(ctx context.Context, currentGroupID int64, platform, subscriptionType string, fallbackGroupID int64) error { - if platform != PlatformAnthropic && platform != PlatformAntigravity { - return fmt.Errorf("invalid request fallback only supported for anthropic or antigravity groups") - } - if subscriptionType == SubscriptionTypeSubscription { - return fmt.Errorf("subscription groups cannot set invalid request fallback") - } - if currentGroupID > 0 && currentGroupID == fallbackGroupID { - return fmt.Errorf("cannot set self as invalid request fallback group") - } - - fallbackGroup, err := s.groupRepo.GetByIDLite(ctx, fallbackGroupID) - if err != nil { - return fmt.Errorf("fallback group not found: %w", err) - } - if fallbackGroup.Platform != PlatformAnthropic { - return fmt.Errorf("fallback group must be anthropic platform") - } - if fallbackGroup.SubscriptionType == SubscriptionTypeSubscription { - return fmt.Errorf("fallback group cannot be subscription type") - } - if fallbackGroup.FallbackGroupIDOnInvalidRequest != nil { - return fmt.Errorf("fallback group cannot have invalid request fallback configured") - } - return nil -} - -func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error) { - group, err := s.groupRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - - if input.Name != "" { - group.Name = input.Name - } - if input.Description != nil { - group.Description = *input.Description - } - if input.Platform != "" { - group.Platform = input.Platform - } - if input.RateMultiplier != nil { - if *input.RateMultiplier <= 0 { - return nil, errors.New("rate_multiplier must be > 0") - } - group.RateMultiplier = *input.RateMultiplier - } - if input.IsExclusive != nil { - group.IsExclusive = *input.IsExclusive - } - if input.Status != "" { - group.Status = input.Status - } - - // 订阅相关字段 - if input.SubscriptionType != "" { - group.SubscriptionType = input.SubscriptionType - } - // 限额字段:nil/负数 表示"无限制",0 表示"不允许用量",正数表示具体限额 - // 前端始终发送这三个字段,无需 nil 守卫 - group.DailyLimitUSD = normalizeLimit(input.DailyLimitUSD) - group.WeeklyLimitUSD = normalizeLimit(input.WeeklyLimitUSD) - group.MonthlyLimitUSD = normalizeLimit(input.MonthlyLimitUSD) - // 图片生成计费配置:负数表示清除(使用默认价格) - if input.AllowImageGeneration != nil { - group.AllowImageGeneration = *input.AllowImageGeneration - } - if input.ImageRateIndependent != nil { - group.ImageRateIndependent = *input.ImageRateIndependent - } - if input.ImageRateMultiplier != nil { - if *input.ImageRateMultiplier < 0 { - return nil, errors.New("image_rate_multiplier must be >= 0") - } - group.ImageRateMultiplier = *input.ImageRateMultiplier - } - if input.PeakRateEnabled != nil { - group.PeakRateEnabled = *input.PeakRateEnabled - } - if input.PeakStart != nil { - group.PeakStart = *input.PeakStart - } - if input.PeakEnd != nil { - group.PeakEnd = *input.PeakEnd - } - if input.PeakRateMultiplier != nil { - group.PeakRateMultiplier = *input.PeakRateMultiplier - } - // 先归一化(非订阅分组——含本次更新转为非订阅——静默清空高峰配置,清洗停用状态下的脏字段), - // 再收敛校验:Update 可能只传部分 peak 字段,需对合并后的最终配置统一校验, - // 防止单独修改 start/end 导致最终 start>=end 等非法配置入库。与 CreateGroup 同一收口。 - group.PeakRateEnabled, group.PeakStart, group.PeakEnd, group.PeakRateMultiplier = NormalizePeakRateConfig(group.SubscriptionType, group.PeakRateEnabled, group.PeakStart, group.PeakEnd, group.PeakRateMultiplier) - if err := ValidatePeakRateConfig(group.SubscriptionType, group.PeakRateEnabled, group.PeakStart, group.PeakEnd, group.PeakRateMultiplier); err != nil { - return nil, err - } - if input.ImagePrice1K != nil { - group.ImagePrice1K = normalizePrice(input.ImagePrice1K) - } - if input.ImagePrice2K != nil { - group.ImagePrice2K = normalizePrice(input.ImagePrice2K) - } - if input.ImagePrice4K != nil { - group.ImagePrice4K = normalizePrice(input.ImagePrice4K) - } - - // Claude Code 客户端限制 - if input.ClaudeCodeOnly != nil { - group.ClaudeCodeOnly = *input.ClaudeCodeOnly - } - if input.FallbackGroupID != nil { - // 校验降级分组 - if *input.FallbackGroupID > 0 { - if err := s.validateFallbackGroup(ctx, id, *input.FallbackGroupID); err != nil { - return nil, err - } - group.FallbackGroupID = input.FallbackGroupID - } else { - // 传入 0 或负数表示清除降级分组 - group.FallbackGroupID = nil - } - } - fallbackOnInvalidRequest := group.FallbackGroupIDOnInvalidRequest - if input.FallbackGroupIDOnInvalidRequest != nil { - if *input.FallbackGroupIDOnInvalidRequest > 0 { - fallbackOnInvalidRequest = input.FallbackGroupIDOnInvalidRequest - } else { - fallbackOnInvalidRequest = nil - } - } - if fallbackOnInvalidRequest != nil { - if err := s.validateFallbackGroupOnInvalidRequest(ctx, id, group.Platform, group.SubscriptionType, *fallbackOnInvalidRequest); err != nil { - return nil, err - } - } - group.FallbackGroupIDOnInvalidRequest = fallbackOnInvalidRequest - - // 模型路由配置 - if input.ModelRouting != nil { - group.ModelRouting = input.ModelRouting - } - if input.ModelRoutingEnabled != nil { - group.ModelRoutingEnabled = *input.ModelRoutingEnabled - } - if input.MCPXMLInject != nil { - group.MCPXMLInject = *input.MCPXMLInject - } - - // 支持的模型系列(仅 antigravity 平台使用) - if input.SupportedModelScopes != nil { - group.SupportedModelScopes = *input.SupportedModelScopes - } - - // OpenAI Messages 调度配置 - if input.AllowMessagesDispatch != nil { - group.AllowMessagesDispatch = *input.AllowMessagesDispatch - } - if input.RequireOAuthOnly != nil { - group.RequireOAuthOnly = *input.RequireOAuthOnly - } - if input.RequirePrivacySet != nil { - group.RequirePrivacySet = *input.RequirePrivacySet - } - if input.DefaultMappedModel != nil { - group.DefaultMappedModel = *input.DefaultMappedModel - } - if input.MessagesDispatchModelConfig != nil { - group.MessagesDispatchModelConfig = normalizeOpenAIMessagesDispatchModelConfig(*input.MessagesDispatchModelConfig) - } - if input.ModelsListConfig != nil { - group.ModelsListConfig = normalizeGroupModelsListConfig(*input.ModelsListConfig) - } - if input.RPMLimit != nil { - group.RPMLimit = *input.RPMLimit - } - sanitizeGroupMessagesDispatchFields(group) - - if err := s.groupRepo.Update(ctx, group); err != nil { - return nil, err - } - - if s.authCacheInvalidator != nil { - s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, id) - } - - // 如果指定了复制账号的源分组,同步绑定(替换当前分组的账号) - if len(input.CopyAccountsFromGroupIDs) > 0 { - // 去重源分组 IDs - seen := make(map[int64]struct{}) - uniqueSourceGroupIDs := make([]int64, 0, len(input.CopyAccountsFromGroupIDs)) - for _, srcGroupID := range input.CopyAccountsFromGroupIDs { - // 校验:源分组不能是自身 - if srcGroupID == id { - return nil, fmt.Errorf("cannot copy accounts from self") - } - // 去重 - if _, exists := seen[srcGroupID]; !exists { - seen[srcGroupID] = struct{}{} - uniqueSourceGroupIDs = append(uniqueSourceGroupIDs, srcGroupID) - } - } - - // 校验源分组的平台是否与当前分组一致 - for _, srcGroupID := range uniqueSourceGroupIDs { - srcGroup, err := s.groupRepo.GetByIDLite(ctx, srcGroupID) - if err != nil { - return nil, fmt.Errorf("source group %d not found: %w", srcGroupID, err) - } - if srcGroup.Platform != group.Platform { - return nil, fmt.Errorf("source group %d platform mismatch: expected %s, got %s", srcGroupID, group.Platform, srcGroup.Platform) - } - } - - // 获取所有源分组的账号(去重) - accountIDsToCopy, err := s.groupRepo.GetAccountIDsByGroupIDs(ctx, uniqueSourceGroupIDs) - if err != nil { - return nil, fmt.Errorf("failed to get accounts from source groups: %w", err) - } - - // 先清空当前分组的所有账号绑定 - if _, err := s.groupRepo.DeleteAccountGroupsByGroupID(ctx, id); err != nil { - return nil, fmt.Errorf("failed to clear existing account bindings: %w", err) - } - - // require_oauth_only: 过滤掉 apikey 类型账号 - if group.RequireOAuthOnly && (group.Platform == PlatformOpenAI || group.Platform == PlatformAntigravity || group.Platform == PlatformAnthropic || group.Platform == PlatformGemini || group.Platform == PlatformGrok) && len(accountIDsToCopy) > 0 { - accounts, err := s.accountRepo.GetByIDs(ctx, accountIDsToCopy) - if err != nil { - return nil, fmt.Errorf("failed to fetch accounts for oauth filter: %w", err) - } - oauthIDs := make(map[int64]struct{}, len(accounts)) - for _, acc := range accounts { - if acc.Type != AccountTypeAPIKey { - oauthIDs[acc.ID] = struct{}{} - } - } - var filtered []int64 - for _, aid := range accountIDsToCopy { - if _, ok := oauthIDs[aid]; ok { - filtered = append(filtered, aid) - } - } - accountIDsToCopy = filtered - } - - // 再绑定源分组的账号 - if len(accountIDsToCopy) > 0 { - if err := s.groupRepo.BindAccountsToGroup(ctx, id, accountIDsToCopy); err != nil { - return nil, fmt.Errorf("failed to bind accounts to group: %w", err) - } - } - } - - return group, nil -} - -func (s *adminServiceImpl) DeleteGroup(ctx context.Context, id int64) error { - var groupKeys []string - if s.authCacheInvalidator != nil { - keys, err := s.apiKeyRepo.ListKeysByGroupID(ctx, id) - if err == nil { - groupKeys = keys - } - } - - affectedUserIDs, err := s.groupRepo.DeleteCascade(ctx, id) - if err != nil { - return err - } - // 注意:user_group_rate_multipliers 表通过外键 ON DELETE CASCADE 自动清理 - - // 事务成功后,异步失效受影响用户的订阅缓存 - if len(affectedUserIDs) > 0 && s.billingCacheService != nil { - groupID := id - go func() { - cacheCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - for _, userID := range affectedUserIDs { - if err := s.billingCacheService.InvalidateSubscription(cacheCtx, userID, groupID); err != nil { - logger.LegacyPrintf("service.admin", "invalidate subscription cache failed: user_id=%d group_id=%d err=%v", userID, groupID, err) - } - } - }() - } - if s.authCacheInvalidator != nil { - for _, key := range groupKeys { - s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, key) - } - } - - return nil -} - -func (s *adminServiceImpl) GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]APIKey, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize} - keys, result, err := s.apiKeyRepo.ListByGroupID(ctx, groupID, params) - if err != nil { - return nil, 0, err - } - return keys, result.Total, nil -} - -func (s *adminServiceImpl) GetGroupRateMultipliers(ctx context.Context, groupID int64) ([]UserGroupRateEntry, error) { - if s.userGroupRateRepo == nil { - return nil, nil - } - return s.userGroupRateRepo.GetByGroupID(ctx, groupID) -} - -func (s *adminServiceImpl) ClearGroupRateMultipliers(ctx context.Context, groupID int64) error { - if s.userGroupRateRepo == nil { - return nil - } - return s.userGroupRateRepo.DeleteByGroupID(ctx, groupID) -} - -func (s *adminServiceImpl) BatchSetGroupRateMultipliers(ctx context.Context, groupID int64, entries []GroupRateMultiplierInput) error { - if s.userGroupRateRepo == nil { - return nil - } - for _, e := range entries { - if e.RateMultiplier <= 0 { - return fmt.Errorf("rate_multiplier must be > 0 (user_id=%d)", e.UserID) - } - } - return s.userGroupRateRepo.SyncGroupRateMultipliers(ctx, groupID, entries) -} - -func (s *adminServiceImpl) ClearGroupRPMOverrides(ctx context.Context, groupID int64) error { - if s.userGroupRateRepo == nil { - return nil - } - if err := s.userGroupRateRepo.ClearGroupRPMOverrides(ctx, groupID); err != nil { - return err - } - // RPM override 已嵌入 auth cache snapshot (v7),变更后必须失效相关缓存。 - if s.authCacheInvalidator != nil { - s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, groupID) - } - return nil -} - -func (s *adminServiceImpl) BatchSetGroupRPMOverrides(ctx context.Context, groupID int64, entries []GroupRPMOverrideInput) error { - if s.userGroupRateRepo == nil { - return nil - } - for _, e := range entries { - if e.RPMOverride != nil && *e.RPMOverride < 0 { - return infraerrors.BadRequest("INVALID_RPM_OVERRIDE", fmt.Sprintf("rpm_override must be >= 0 (user_id=%d)", e.UserID)) - } - } - if err := s.userGroupRateRepo.SyncGroupRPMOverrides(ctx, groupID, entries); err != nil { - return err - } - // RPM override 已嵌入 auth cache snapshot (v7),变更后必须失效相关缓存。 - if s.authCacheInvalidator != nil { - s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, groupID) - } - return nil -} - -func (s *adminServiceImpl) UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error { - return s.groupRepo.UpdateSortOrders(ctx, updates) -} - -// AdminUpdateAPIKeyGroupID 管理员修改 API Key 分组绑定 -// groupID: nil=不修改, 指向0=解绑, 指向正整数=绑定到目标分组 -func (s *adminServiceImpl) AdminUpdateAPIKeyGroupID(ctx context.Context, keyID int64, groupID *int64) (*AdminUpdateAPIKeyGroupIDResult, error) { - apiKey, err := s.apiKeyRepo.GetByID(ctx, keyID) - if err != nil { - return nil, err - } - - if groupID == nil { - // nil 表示不修改,直接返回 - return &AdminUpdateAPIKeyGroupIDResult{APIKey: apiKey}, nil - } - - if *groupID < 0 { - return nil, infraerrors.BadRequest("INVALID_GROUP_ID", "group_id must be non-negative") - } - - result := &AdminUpdateAPIKeyGroupIDResult{} - - if *groupID == 0 { - // 0 表示解绑分组(不修改 user_allowed_groups,避免影响用户其他 Key) - apiKey.GroupID = nil - apiKey.Group = nil - } else { - // 验证目标分组存在且状态为 active - group, err := s.groupRepo.GetByID(ctx, *groupID) - if err != nil { - return nil, err - } - if group.Status != StatusActive { - return nil, infraerrors.BadRequest("GROUP_NOT_ACTIVE", "target group is not active") - } - // 订阅类型分组:用户须持有该分组的有效订阅才可绑定 - if group.IsSubscriptionType() { - if s.userSubRepo == nil { - return nil, infraerrors.InternalServer("SUBSCRIPTION_REPOSITORY_UNAVAILABLE", "subscription repository is not configured") - } - if _, err := s.userSubRepo.GetActiveByUserIDAndGroupID(ctx, apiKey.UserID, *groupID); err != nil { - if errors.Is(err, ErrSubscriptionNotFound) { - return nil, infraerrors.BadRequest("SUBSCRIPTION_REQUIRED", "user does not have an active subscription for this group") - } - return nil, err - } - } - - gid := *groupID - apiKey.GroupID = &gid - apiKey.Group = group - - // 专属标准分组:使用事务保证「添加分组权限」与「更新 API Key」的原子性 - if group.IsExclusive && !group.IsSubscriptionType() { - opCtx := ctx - var tx *dbent.Tx - if s.entClient == nil { - logger.LegacyPrintf("service.admin", "Warning: entClient is nil, skipping transaction protection for exclusive group binding") - } else { - var txErr error - tx, txErr = s.entClient.Tx(ctx) - if txErr != nil { - return nil, fmt.Errorf("begin transaction: %w", txErr) - } - defer func() { _ = tx.Rollback() }() - opCtx = dbent.NewTxContext(ctx, tx) - } - - if addErr := s.userRepo.AddGroupToAllowedGroups(opCtx, apiKey.UserID, gid); addErr != nil { - return nil, fmt.Errorf("add group to user allowed groups: %w", addErr) - } - if err := s.apiKeyRepo.Update(opCtx, apiKey); err != nil { - return nil, fmt.Errorf("update api key: %w", err) - } - if tx != nil { - if err := tx.Commit(); err != nil { - return nil, fmt.Errorf("commit transaction: %w", err) - } - } - - result.AutoGrantedGroupAccess = true - result.GrantedGroupID = &gid - result.GrantedGroupName = group.Name - - // 失效认证缓存(在事务提交后执行) - if s.authCacheInvalidator != nil { - s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, apiKey.Key) - } - - result.APIKey = apiKey - return result, nil - } - } - - // 非专属分组 / 解绑:无需事务,单步更新即可 - if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil { - return nil, fmt.Errorf("update api key: %w", err) - } - - // 失效认证缓存 - if s.authCacheInvalidator != nil { - s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, apiKey.Key) - } - - result.APIKey = apiKey - return result, nil -} - -// AdminResetAPIKeyRateLimitUsage resets all API key rate-limit usage windows. -func (s *adminServiceImpl) AdminResetAPIKeyRateLimitUsage(ctx context.Context, keyID int64) (*APIKey, error) { - apiKey, err := s.apiKeyRepo.GetByID(ctx, keyID) - if err != nil { - return nil, err - } - apiKey.Usage5h = 0 - apiKey.Usage1d = 0 - apiKey.Usage7d = 0 - apiKey.Window5hStart = nil - apiKey.Window1dStart = nil - apiKey.Window7dStart = nil - if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil { - return nil, fmt.Errorf("reset api key rate limit usage: %w", err) - } - if s.authCacheInvalidator != nil { - s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, apiKey.Key) - } - if s.billingCacheService != nil { - _ = s.billingCacheService.InvalidateAPIKeyRateLimit(ctx, apiKey.ID) - } - return apiKey, nil -} - -// ReplaceUserGroup 替换用户的专属分组 -func (s *adminServiceImpl) ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*ReplaceUserGroupResult, error) { - if oldGroupID == newGroupID { - return nil, infraerrors.BadRequest("SAME_GROUP", "old and new group must be different") - } - - // 验证新分组存在且为活跃的专属标准分组 - newGroup, err := s.groupRepo.GetByID(ctx, newGroupID) - if err != nil { - return nil, err - } - if newGroup.Status != StatusActive { - return nil, infraerrors.BadRequest("GROUP_NOT_ACTIVE", "target group is not active") - } - if !newGroup.IsExclusive { - return nil, infraerrors.BadRequest("GROUP_NOT_EXCLUSIVE", "target group is not exclusive") - } - if newGroup.IsSubscriptionType() { - return nil, infraerrors.BadRequest("GROUP_IS_SUBSCRIPTION", "subscription groups are not supported for replacement") - } - - // 事务保证原子性 - if s.entClient == nil { - return nil, fmt.Errorf("entClient is nil, cannot perform group replacement") - } - tx, err := s.entClient.Tx(ctx) - if err != nil { - return nil, fmt.Errorf("begin transaction: %w", err) - } - defer func() { _ = tx.Rollback() }() - opCtx := dbent.NewTxContext(ctx, tx) - - // 1. 授予新分组权限 - if err := s.userRepo.AddGroupToAllowedGroups(opCtx, userID, newGroupID); err != nil { - return nil, fmt.Errorf("add new group to allowed groups: %w", err) - } - - // 2. 迁移绑定旧分组的 Key 到新分组 - migrated, err := s.apiKeyRepo.UpdateGroupIDByUserAndGroup(opCtx, userID, oldGroupID, newGroupID) - if err != nil { - return nil, fmt.Errorf("migrate api keys: %w", err) - } - - // 3. 移除旧分组权限 - if err := s.userRepo.RemoveGroupFromUserAllowedGroups(opCtx, userID, oldGroupID); err != nil { - return nil, fmt.Errorf("remove old group from allowed groups: %w", err) - } - - if err := tx.Commit(); err != nil { - return nil, fmt.Errorf("commit transaction: %w", err) - } - - // 失效该用户所有 Key 的认证缓存 - if s.authCacheInvalidator != nil { - keys, keyErr := s.apiKeyRepo.ListKeysByUserID(ctx, userID) - if keyErr == nil { - for _, k := range keys { - s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, k) - } - } - } - - return &ReplaceUserGroupResult{MigratedKeys: migrated}, nil -} - -// Account management implementations -func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string, sortBy, sortOrder string) ([]Account, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - accounts, result, err := s.accountRepo.ListWithFilters(ctx, params, platform, accountType, status, search, groupID, privacyMode) - if err != nil { - return nil, 0, err - } - return accounts, result.Total, nil -} - -func (s *adminServiceImpl) ListAccountsForSchedulerScoreFilter(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) { - if s == nil || s.accountRepo == nil { - return nil, nil - } - return s.accountRepo.ListAllWithFilters(ctx, platform, accountType, status, search, groupID, privacyMode) -} - -func (s *adminServiceImpl) ListOpenAISchedulableAccountsForSchedulerScore(ctx context.Context, groupID *int64) ([]Account, error) { - if s == nil || s.accountRepo == nil { - return nil, nil - } - if groupID != nil { - return s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, PlatformOpenAI) - } - return s.accountRepo.ListSchedulableUngroupedByPlatform(ctx, PlatformOpenAI) -} - -func (s *adminServiceImpl) GetAccount(ctx context.Context, id int64) (*Account, error) { - return s.accountRepo.GetByID(ctx, id) -} - -func (s *adminServiceImpl) GetAccountsByIDs(ctx context.Context, ids []int64) ([]*Account, error) { - if len(ids) == 0 { - return []*Account{}, nil - } - - accounts, err := s.accountRepo.GetByIDs(ctx, ids) - if err != nil { - return nil, fmt.Errorf("failed to get accounts by IDs: %w", err) - } - - return accounts, nil -} - -func normalizeAccountConcurrency(platform, accountType string, concurrency int) int { - if platform == PlatformGrok && accountType == AccountTypeOAuth { - if concurrency <= 0 { - return 1 - } - } - return concurrency -} - -func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error) { - // 绑定分组 - groupIDs := input.GroupIDs - // 如果没有指定分组,自动绑定对应平台的默认分组 - if len(groupIDs) == 0 && !input.SkipDefaultGroupBind { - defaultGroupName := input.Platform + "-default" - groups, err := s.groupRepo.ListActiveByPlatform(ctx, input.Platform) - if err == nil { - for _, g := range groups { - if g.Name == defaultGroupName { - groupIDs = []int64{g.ID} - break - } - } - } - } - - // 检查混合渠道风险(除非用户已确认) - if len(groupIDs) > 0 && !input.SkipMixedChannelCheck { - if err := s.checkMixedChannelRisk(ctx, 0, input.Platform, groupIDs); err != nil { - return nil, err - } - } - - // 校验并规范化请求头覆写配置(header 名小写化、格式检查) - if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { - return nil, err - } - - account := &Account{ - Name: input.Name, - Notes: normalizeAccountNotes(input.Notes), - Platform: input.Platform, - Type: input.Type, - Credentials: input.Credentials, - Extra: input.Extra, - ProxyID: input.ProxyID, - Concurrency: normalizeAccountConcurrency(input.Platform, input.Type, input.Concurrency), - Priority: input.Priority, - Status: StatusActive, - Schedulable: true, - } - // 预计算固定时间重置的下次重置时间 - if account.Extra != nil { - if err := ValidateQuotaResetConfig(account.Extra); err != nil { - return nil, err - } - ComputeQuotaResetAt(account.Extra) - NormalizeFixedQuotaWindows(account.Extra) - } - if input.ExpiresAt != nil && *input.ExpiresAt > 0 { - expiresAt := time.Unix(*input.ExpiresAt, 0) - account.ExpiresAt = &expiresAt - } - if input.AutoPauseOnExpired != nil { - account.AutoPauseOnExpired = *input.AutoPauseOnExpired - } else { - account.AutoPauseOnExpired = true - } - if input.RateMultiplier != nil { - if *input.RateMultiplier < 0 { - return nil, errors.New("rate_multiplier must be >= 0") - } - account.RateMultiplier = input.RateMultiplier - } - if input.LoadFactor != nil && *input.LoadFactor > 0 { - if *input.LoadFactor > 10000 { - return nil, errors.New("load_factor must be <= 10000") - } - account.LoadFactor = input.LoadFactor - } - if err := s.accountRepo.Create(ctx, account); err != nil { - return nil, err - } - - // 绑定分组 - if len(groupIDs) > 0 { - if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { - return nil, err - } - } - - // OAuth 账号:创建后异步设置隐私。 - // 使用 Ensure(幂等)而非 Force:新建账号 Extra 为空时效果相同,但更安全。 - if account.Type == AccountTypeOAuth { - switch account.Platform { - case PlatformOpenAI: - go func() { - defer func() { - if r := recover(); r != nil { - slog.Error("create_account_openai_privacy_panic", "account_id", account.ID, "recover", r) - } - }() - s.EnsureOpenAIPrivacy(context.Background(), account) - }() - case PlatformAntigravity: - go func() { - defer func() { - if r := recover(); r != nil { - slog.Error("create_account_antigravity_privacy_panic", "account_id", account.ID, "recover", r) - } - }() - s.EnsureAntigravityPrivacy(context.Background(), account) - }() - } - } - - return account, nil -} - -func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *UpdateAccountInput) (*Account, error) { - account, err := s.accountRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - // 安全/身份不变量(影子账号):通用更新路径被 edit/re-auth/refresh/batch 共用, - // 必须在此守住,否则仅在创建时的保证可被这些路径绕过。 - if account.IsCredentialShadow() { - // 影子绝不持有凭据(凭据只在母账号)——外审 F5。 - if !isAllowedSparkShadowCredentialsUpdate(input.Credentials) { - return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_NO_CREDENTIALS", - "spark shadow accounts do not hold auth credentials; only model mapping can be configured on the shadow account") - } - // 影子 type 不可变——很多上游逻辑按 account.Type 分支(OAuth transform / ChatGPT - // header 注入 / WS OAuth 决策),改成 apikey 会让 spark 影子被选中后按错误协议转发(外审 G7)。 - if input.Type != "" && input.Type != account.Type { - return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_IMMUTABLE_TYPE", - "spark shadow account type cannot be changed; it must remain an OpenAI OAuth shadow") - } - } else if input.Type != "" && input.Type != account.Type && input.Type != AccountTypeOAuth { - // 母账号守卫(外审 D/P1):有 spark 影子的账号不能把 type 改出 OpenAI OAuth——影子读透母 - // 凭据,母变成 apikey/setup_token 会让影子被调度后按错协议失败(resolveCredentialAccount - // 必报错)。须先删影子再改 type。 - shadows, serr := s.accountRepo.ListShadowsByParent(ctx, id) - if serr != nil { - return nil, serr - } - if len(shadows) > 0 { - return nil, infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_PARENT_IMMUTABLE_TYPE", - "cannot change account type while it has a spark shadow; delete the shadow first") - } - } - wasOveragesEnabled := account.IsOveragesEnabled() - - if input.Name != "" { - account.Name = input.Name - } - if input.Type != "" { - account.Type = input.Type - } - if input.Notes != nil { - account.Notes = normalizeAccountNotes(input.Notes) - } - if account.IsCredentialShadow() && input.Credentials != nil { - account.Credentials = sanitizeSparkShadowCredentials(input.Credentials) - } else if len(input.Credentials) > 0 { - // 敏感子键采用"incoming 没提供就保留"的合并语义:前端响应已脱敏, - // 全对象 PUT 编辑时不会再带回 token,避免覆盖时清空已有凭证。 - account.Credentials = MergePreservingSensitiveCreds(account.Credentials, input.Credentials) - // 校验并规范化请求头覆写配置(header 名小写化、格式检查) - if err := NormalizeHeaderOverrideCredentials(account.Credentials); err != nil { - return nil, err - } - } - // Extra 使用 map:需要区分“未提供(nil)”与“显式清空({})”。 - // 关闭配额限制时前端会删除 quota_* 键并提交 extra:{},此时也必须落库。 - if input.Extra != nil { - // 保留配额用量字段,防止编辑账号时意外重置 - for _, key := range []string{"quota_used", "quota_daily_used", "quota_daily_start", "quota_weekly_used", "quota_weekly_start"} { - if v, ok := account.Extra[key]; ok { - input.Extra[key] = v - } - } - account.Extra = input.Extra - if account.Platform == PlatformAntigravity && wasOveragesEnabled && !account.IsOveragesEnabled() { - delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 - // 清除 AICredits 限流 key - if rawLimits, ok := account.Extra[modelRateLimitsKey].(map[string]any); ok { - delete(rawLimits, creditsExhaustedKey) - } - } - if account.Platform == PlatformAntigravity && !wasOveragesEnabled && account.IsOveragesEnabled() { - delete(account.Extra, modelRateLimitsKey) - delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 - } - // 校验并预计算固定时间重置的下次重置时间 - if err := ValidateQuotaResetConfig(account.Extra); err != nil { - return nil, err - } - ComputeQuotaResetAt(account.Extra) - NormalizeFixedQuotaWindows(account.Extra) - } - // 影子代理恒继承母账号(由 propagateProxyToShadows 同步),不接受独立编辑——外审 B/P1; - // 否则要等母账号下次改 proxy 才被覆盖,期间影子会出现"有时继承、有时独立"的漂移。 - if input.ProxyID != nil && !account.IsCredentialShadow() { - // 0 表示清除代理(前端发送 0 而不是 null 来表达清除意图) - if *input.ProxyID == 0 { - account.ProxyID = nil - } else { - account.ProxyID = input.ProxyID - } - account.Proxy = nil // 清除关联对象,防止 GORM Save 时根据 Proxy.ID 覆盖 ProxyID - } - // 只在指针非 nil 时更新 Concurrency(支持设置为 0) - if input.Concurrency != nil { - account.Concurrency = normalizeAccountConcurrency(account.Platform, account.Type, *input.Concurrency) - } - // 只在指针非 nil 时更新 Priority(支持设置为 0) - if input.Priority != nil { - account.Priority = *input.Priority - } - if input.RateMultiplier != nil { - if *input.RateMultiplier < 0 { - return nil, errors.New("rate_multiplier must be >= 0") - } - account.RateMultiplier = input.RateMultiplier - } - if input.LoadFactor != nil { - if *input.LoadFactor <= 0 { - account.LoadFactor = nil // 0 或负数表示清除 - } else if *input.LoadFactor > 10000 { - return nil, errors.New("load_factor must be <= 10000") - } else { - account.LoadFactor = input.LoadFactor - } - } - if input.Status != "" { - account.Status = input.Status - } - if input.ExpiresAt != nil { - if *input.ExpiresAt <= 0 { - account.ExpiresAt = nil - } else { - expiresAt := time.Unix(*input.ExpiresAt, 0) - account.ExpiresAt = &expiresAt - } - } - if input.AutoPauseOnExpired != nil { - account.AutoPauseOnExpired = *input.AutoPauseOnExpired - } - - // 先验证分组是否存在(在任何写操作之前) - if input.GroupIDs != nil { - if err := s.validateGroupIDsExist(ctx, *input.GroupIDs); err != nil { - return nil, err - } - - // 检查混合渠道风险(除非用户已确认) - if !input.SkipMixedChannelCheck { - if err := s.checkMixedChannelRisk(ctx, account.ID, account.Platform, *input.GroupIDs); err != nil { - return nil, err - } - } - } - - if err := s.accountRepo.Update(ctx, account); err != nil { - return nil, err - } - - // 将 proxy 变更传播到 spark 影子账号(同步;Update 内部已触发调度快照)。 - // 影子自身 proxy 不可独立编辑(见上),故对影子的更新不触发传播。 - if input.ProxyID != nil && !account.IsCredentialShadow() { - if err := s.propagateProxyToShadows(ctx, id, account.ProxyID); err != nil { - return nil, err - } - } - - // 绑定分组 - if input.GroupIDs != nil { - if err := s.accountRepo.BindGroups(ctx, account.ID, *input.GroupIDs); err != nil { - return nil, err - } - } - - // 重新查询以确保返回完整数据(包括正确的 Proxy 关联对象) - updated, err := s.accountRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - return updated, nil -} - -// UpdateAccountExtra 仅对 Extra JSONB 做 key 级合并,避免覆盖其它运行态键 -// (如 model_rate_limits / passive_usage_* 等)。 -func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error { - if len(updates) == 0 { - return nil - } - return s.accountRepo.UpdateExtra(ctx, id, updates) -} - -// BulkUpdateAccounts updates multiple accounts in one request. -// It merges credentials/extra keys instead of overwriting the whole object. -func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) { - if len(input.AccountIDs) == 0 && input.Filters != nil { - accountIDs, err := s.resolveBulkUpdateTargetIDs(ctx, input.Filters) - if err != nil { - return nil, err - } - input.AccountIDs = accountIDs - } - - result := &BulkUpdateAccountsResult{ - SuccessIDs: make([]int64, 0, len(input.AccountIDs)), - FailedIDs: make([]int64, 0, len(input.AccountIDs)), - Results: make([]BulkUpdateAccountResult, 0, len(input.AccountIDs)), - } - - if len(input.AccountIDs) == 0 { - return result, nil - } - if input.GroupIDs != nil { - if err := s.validateGroupIDsExist(ctx, *input.GroupIDs); err != nil { - return nil, err - } - } - - needMixedChannelCheck := input.GroupIDs != nil && !input.SkipMixedChannelCheck - - // 预取所有目标账号,供凭据守卫/代理守卫/混合渠道检查共用,避免多次 DB 查询。 - var cachedTargets []*Account - if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck { - loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) - if err != nil { - return nil, err - } - cachedTargets = loaded - } - - // 影子账号绝不持有凭据:批量更新携带凭据时,目标中不得含影子(外审 G5,与单账号 - // UpdateAccount 守卫对齐)。覆盖显式 IDs 与 filter 解析出的 IDs(此处 AccountIDs 已解析完成)。 - if len(input.Credentials) > 0 { - for _, acc := range cachedTargets { - if acc != nil && acc.IsCredentialShadow() { - return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_NO_CREDENTIALS", - "spark shadow account %d cannot hold credentials; manage credentials on the parent account", acc.ID) - } - } - } - - // 影子账号 proxy 恒继承母账号(与单账号 UpdateAccount 守卫对齐——外审第4轮 P1):批量携带 proxy - // 时目标不得含影子,否则影子会获得独立 proxy、破坏继承不变量(网关按所选影子自身 proxy 出站, - // 要等母账号下次改 proxy 才覆盖→漂移)。含影子即整体拒绝,提示从选择中剔除影子。 - if input.ProxyID != nil { - for _, acc := range cachedTargets { - if acc != nil && acc.IsCredentialShadow() { - return nil, infraerrors.Newf(http.StatusBadRequest, "SPARK_SHADOW_PROXY_INHERITED", - "spark shadow account %d proxy is inherited from its parent and cannot be set in bulk; manage it on the parent account", acc.ID) - } - } - } - - // 预加载账号平台信息(混合渠道检查需要)。 - platformByID := map[int64]string{} - if needMixedChannelCheck { - for _, account := range cachedTargets { - if account != nil { - platformByID[account.ID] = account.Platform - } - } - } - - // 预检查混合渠道风险:在任何写操作之前,若发现风险立即返回错误。 - if needMixedChannelCheck { - for _, accountID := range input.AccountIDs { - platform := platformByID[accountID] - if platform == "" { - continue - } - if err := s.checkMixedChannelRisk(ctx, accountID, platform, *input.GroupIDs); err != nil { - return nil, err - } - } - } - - if input.RateMultiplier != nil { - if *input.RateMultiplier < 0 { - return nil, errors.New("rate_multiplier must be >= 0") - } - } - - // 校验并规范化请求头覆写配置(批量路径为 JSONB 顶层 key 合并,直接校验增量即可) - if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { - return nil, err - } - - // Prepare bulk updates for columns and JSONB fields. - repoUpdates := AccountBulkUpdate{ - Credentials: input.Credentials, - Extra: input.Extra, - } - if input.Name != "" { - repoUpdates.Name = &input.Name - } - if input.ProxyID != nil { - repoUpdates.ProxyID = input.ProxyID - } - if input.Concurrency != nil { - repoUpdates.Concurrency = input.Concurrency - } - if input.Priority != nil { - repoUpdates.Priority = input.Priority - } - if input.RateMultiplier != nil { - repoUpdates.RateMultiplier = input.RateMultiplier - } - if input.LoadFactor != nil { - if *input.LoadFactor <= 0 { - repoUpdates.LoadFactor = nil // 0 或负数表示清除 - } else if *input.LoadFactor > 10000 { - return nil, errors.New("load_factor must be <= 10000") - } else { - repoUpdates.LoadFactor = input.LoadFactor - } - } - if input.Status != "" { - repoUpdates.Status = &input.Status - } - if input.Schedulable != nil { - repoUpdates.Schedulable = input.Schedulable - } - - // Run bulk update for column/jsonb fields first. - if _, err := s.accountRepo.BulkUpdate(ctx, input.AccountIDs, repoUpdates); err != nil { - return nil, err - } - - // 将 proxy 变更传播到每个目标账号的 spark 影子账号 - if repoUpdates.ProxyID != nil { - var effectiveProxyID *int64 - if *repoUpdates.ProxyID != 0 { - effectiveProxyID = repoUpdates.ProxyID - } - for _, accountID := range input.AccountIDs { - if err := s.propagateProxyToShadows(ctx, accountID, effectiveProxyID); err != nil { - return nil, err - } - } - } - - // Handle group bindings per account (requires individual operations). - for _, accountID := range input.AccountIDs { - entry := BulkUpdateAccountResult{AccountID: accountID} - - if input.GroupIDs != nil { - if err := s.accountRepo.BindGroups(ctx, accountID, *input.GroupIDs); err != nil { - entry.Success = false - entry.Error = err.Error() - result.Failed++ - result.FailedIDs = append(result.FailedIDs, accountID) - result.Results = append(result.Results, entry) - continue - } - } - - entry.Success = true - result.Success++ - result.SuccessIDs = append(result.SuccessIDs, accountID) - result.Results = append(result.Results, entry) - } - - return result, nil -} - -func (s *adminServiceImpl) resolveBulkUpdateTargetIDs(ctx context.Context, filters *BulkUpdateAccountFilters) ([]int64, error) { - if filters == nil { - return nil, nil - } - - groupID := int64(0) - switch strings.TrimSpace(filters.Group) { - case "": - case "ungrouped": - groupID = AccountListGroupUngrouped - default: - parsedGroupID, err := strconv.ParseInt(strings.TrimSpace(filters.Group), 10, 64) - if err != nil { - return nil, fmt.Errorf("invalid group filter: %w", err) - } - groupID = parsedGroupID - } - - const pageSize = 500 - page := 1 - accountIDs := make([]int64, 0, pageSize) - - for { - accounts, total, err := s.ListAccounts( - ctx, - page, - pageSize, - filters.Platform, - filters.Type, - filters.Status, - filters.Search, - groupID, - filters.PrivacyMode, - "", - "", - ) - if err != nil { - return nil, err - } - for _, account := range accounts { - accountIDs = append(accountIDs, account.ID) - } - if int64(len(accountIDs)) >= total || len(accounts) == 0 { - return accountIDs, nil - } - page++ - } -} - -func (s *adminServiceImpl) DeleteAccount(ctx context.Context, id int64) error { - // 级联删除 spark 影子账号(先删影子,再删母账号) - shadows, err := s.accountRepo.ListShadowsByParent(ctx, id) - if err != nil { - return fmt.Errorf("list spark shadows for cascade delete: %w", err) - } - for _, shadow := range shadows { - if err := s.accountRepo.Delete(ctx, shadow.ID); err != nil { - return fmt.Errorf("cascade delete spark shadow %d: %w", shadow.ID, err) - } - } - if err := s.accountRepo.Delete(ctx, id); err != nil { - return err - } - return nil -} - -func (s *adminServiceImpl) RefreshAccountCredentials(ctx context.Context, id int64) (*Account, error) { - account, err := s.accountRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - // TODO: Implement refresh logic - return account, nil -} - -func (s *adminServiceImpl) ClearAccountError(ctx context.Context, id int64) (*Account, error) { - if err := s.accountRepo.ClearError(ctx, id); err != nil { - return nil, err - } - if err := s.accountRepo.ClearRateLimit(ctx, id); err != nil { - return nil, err - } - if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, id); err != nil { - return nil, err - } - if err := s.accountRepo.ClearModelRateLimits(ctx, id); err != nil { - return nil, err - } - if err := s.accountRepo.ClearTempUnschedulable(ctx, id); err != nil { - return nil, err - } - if s.runtimeBlocker != nil { - s.runtimeBlocker.ClearAccountSchedulingBlock(id) - } - return s.accountRepo.GetByID(ctx, id) -} - -func (s *adminServiceImpl) SetAccountError(ctx context.Context, id int64, errorMsg string) error { - return s.accountRepo.SetError(ctx, id, errorMsg) -} - -func (s *adminServiceImpl) SetAccountSchedulable(ctx context.Context, id int64, schedulable bool) (*Account, error) { - if err := s.accountRepo.SetSchedulable(ctx, id, schedulable); err != nil { - return nil, err - } - updated, err := s.accountRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - return updated, nil -} - -func (s *adminServiceImpl) RevertAccountProxyFallback(ctx context.Context, id int64) error { - if err := s.accountRepo.RevertProxyFallback(ctx, id); err != nil { - return err - } - // 加载回退后的账号以获取实际 ProxyID,再传播到影子账号 - account, err := s.accountRepo.GetByID(ctx, id) - if err != nil { - return fmt.Errorf("get account after proxy revert: %w", err) - } - return s.propagateProxyToShadows(ctx, id, account.ProxyID) -} - -// CreateShadow 为指定 OpenAI OAuth 母账号创建 spark 维度影子账号(一母一影)。 -// 安全不变量:Credentials 恒不含 auth token(仅 model_mapping,守卫 isAllowedSparkShadowCredentialsUpdate 放行)。 -func (s *adminServiceImpl) CreateShadow(ctx context.Context, parentID int64, opts ShadowOptions) (*Account, error) { - // 1. 加载母账号并校验平台/类型 - parent, err := s.accountRepo.GetByID(ctx, parentID) - if err != nil { - return nil, fmt.Errorf("get parent account: %w", err) - } - if !parent.IsOpenAIOAuth() { - return nil, infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_INVALID_PARENT", - "spark shadow requires an OpenAI OAuth parent account") - } - // G6:母账号本身不能是影子,否则会建出二级影子——resolveCredentialAccount 只解一层, - // 会解析到无凭据的一级影子,进入坏调度/上游失败。 - if parent.IsCredentialShadow() { - return nil, infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_PARENT_IS_SHADOW", - "spark shadow parent must be a real account, not another spark shadow") - } - - // 2. 一母一影校验 - shadows, err := s.accountRepo.ListShadowsByParent(ctx, parentID) - if err != nil { - return nil, fmt.Errorf("check existing spark shadows: %w", err) - } - if len(shadows) > 0 { - return nil, infraerrors.New(http.StatusConflict, "SPARK_SHADOW_ALREADY_EXISTS", - "parent account already has a spark shadow account") - } - - // 3. 解析分组。未指定 GroupIDs 时:优先**继承母账号当前分组**(影子与母同路由域,母在自定义 - // 组时该组的 spark 请求也能选到影子;G1 决策);母无分组再回落 openai-default(F4)。 - // 显式指定 GroupIDs 时,与 UpdateAccount 对齐先校验存在性(创建前),避免建出影子后再因无效组 - // 失败而留下孤儿影子(一母一影唯一索引会挡住重试)——外审 C/P1。 - groupIDs := opts.GroupIDs - if len(groupIDs) > 0 { - if s.groupRepo != nil { - if err := s.validateGroupIDsExist(ctx, groupIDs); err != nil { - return nil, err - } - } - } else if len(parent.GroupIDs) > 0 { - groupIDs = append([]int64(nil), parent.GroupIDs...) - } else if s.groupRepo != nil { - defaultGroupName := PlatformOpenAI + "-default" - if groups, gerr := s.groupRepo.ListActiveByPlatform(ctx, PlatformOpenAI); gerr == nil { - for _, g := range groups { - if g.Name == defaultGroupName { - groupIDs = []int64{g.ID} - break - } - } - } - } - - // 4. 构造影子账号(安全不变量:Credentials 恒不含 auth token,仅含 model_mapping)。 - // name 为空时默认 "<母账号名> (Spark)"——否则空 name 会在 ent(name NotEmpty)处变成裸 500 - // (外审 E/P2);并 rune 安全截断到 ent MaxLen(100)。 - name := strings.TrimSpace(opts.Name) - if name == "" { - name = parent.Name + " (Spark)" - } - if runes := []rune(name); len(runes) > 100 { - name = string(runes[:100]) - } - // 并发未指定(<=0)时继承母账号,避免 0 被限流器解读为"无限并发"(外审 F3)。 - concurrency := opts.Concurrency - if concurrency <= 0 { - concurrency = parent.Concurrency - } - // 优先级未指定(<=0)时继承母账号——前端一键创建只传 name,opts.Priority 省略即 0,而调度 - // 比较是「数值越小越优先」(openai_account_scheduler.isOpenAIAccountCandidateBetter),且 repo - // 显式 SetPriority 会绕过 ent 默认 50,直写 0 会让影子意外抢到最高优先级(外审第5轮 P1)。 - // 与上方 Concurrency 一致采用「省略继承母账号」语义(影子的 proxy/分组/并发亦全部继承母账号)。 - priority := opts.Priority - if priority <= 0 { - priority = parent.Priority - } - shadow := &Account{ - Name: name, - Platform: PlatformOpenAI, - Type: AccountTypeOAuth, - Status: StatusActive, - Credentials: map[string]any{"model_mapping": defaultSparkShadowModelMapping()}, - ParentAccountID: &parentID, - QuotaDimension: QuotaDimensionSpark, - ProxyID: parent.ProxyID, - Priority: priority, - Concurrency: concurrency, - Schedulable: true, - } - - // 5. 持久化(Create 填充 shadow.ID)。并发竞态:预查(步骤2)放行后另一请求抢先建成,本次会撞 - // 一母一影唯一索引。复查确认确为"已存在"竞态时返回结构化 409 而非裸 500——外审 A/P1。 - if err := s.accountRepo.Create(ctx, shadow); err != nil { - if existing, qerr := s.accountRepo.ListShadowsByParent(ctx, parentID); qerr == nil && len(existing) > 0 { - return nil, infraerrors.New(http.StatusConflict, "SPARK_SHADOW_ALREADY_EXISTS", - "parent account already has a spark shadow account") - } - return nil, fmt.Errorf("create spark shadow: %w", err) - } - - // 6. 绑定分组。注意:create+bind 非单一 DB 事务(通用 Create 走 r.client、outbox 走 r.sql, - // 无现成共享事务路径),故绑组失败时做 best-effort 补偿删除刚建的影子,避免半成品影子(否则 - // 一母一影唯一索引会挡住重试)——外审 C/P1。补偿删除用 detached ctx,即便请求 ctx 已取消/超时 - // 仍能完成清理(外审第4轮);进程崩溃这种极端仍可能残留,属已知权衡。 - if len(groupIDs) > 0 { - if err := s.accountRepo.BindGroups(ctx, shadow.ID, groupIDs); err != nil { - if delErr := s.accountRepo.Delete(context.WithoutCancel(ctx), shadow.ID); delErr != nil { - slog.Error("spark_shadow_bind_groups_rollback_failed", - "shadow_id", shadow.ID, "parent_id", parentID, "delete_err", delErr) - } - return nil, fmt.Errorf("bind groups for spark shadow: %w", err) - } - shadow.GroupIDs = groupIDs - } - - return shadow, nil -} - -// propagateProxyToShadows syncs proxyID to all spark shadow accounts of parentID. -// It is called synchronously so that proxy changes are immediately consistent; -// accountRepo.Update triggers the scheduler outbox + cache propagation internally. -// Calling this for a non-parent account is a harmless no-op. -func (s *adminServiceImpl) propagateProxyToShadows(ctx context.Context, parentID int64, proxyID *int64) error { - return propagateAccountProxyToShadows(ctx, s.accountRepo, parentID, proxyID) -} - -// propagateAccountProxyToShadows 把母账号的 proxy 同步到其所有 spark 影子(影子 proxy 恒继承母账号)。 -// 供 AdminService 编辑路径与 CRS 同步路径共用——后者改动母账号 proxy 后必须同样传播,否则影子保留 -// 旧 proxy 出现出站漂移(外审第8轮)。 -func propagateAccountProxyToShadows(ctx context.Context, repo AccountRepository, parentID int64, proxyID *int64) error { - shadows, err := repo.ListShadowsByParent(ctx, parentID) - if err != nil { - return fmt.Errorf("list spark shadows for proxy propagation: %w", err) - } - for _, shadow := range shadows { - shadow.ProxyID = proxyID - if err := repo.Update(ctx, shadow); err != nil { - return fmt.Errorf("update spark shadow %d proxy: %w", shadow.ID, err) - } - } - return nil -} - -// Proxy management implementations -func (s *adminServiceImpl) ListProxies(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]Proxy, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - proxies, result, err := s.proxyRepo.ListWithFilters(ctx, params, protocol, status, search) - if err != nil { - return nil, 0, err - } - return proxies, result.Total, nil -} - -func (s *adminServiceImpl) ListProxiesWithAccountCount(ctx context.Context, page, pageSize int, protocol, status, search string, sortBy, sortOrder string) ([]ProxyWithAccountCount, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - proxies, result, err := s.proxyRepo.ListWithFiltersAndAccountCount(ctx, params, protocol, status, search) - if err != nil { - return nil, 0, err - } - s.attachProxyLatency(ctx, proxies) - return proxies, result.Total, nil -} - -func (s *adminServiceImpl) GetAllProxies(ctx context.Context) ([]Proxy, error) { - return s.proxyRepo.ListActive(ctx) -} - -func (s *adminServiceImpl) GetAllProxiesWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error) { - proxies, err := s.proxyRepo.ListActiveWithAccountCount(ctx) - if err != nil { - return nil, err - } - s.attachProxyLatency(ctx, proxies) - return proxies, nil -} - -func (s *adminServiceImpl) GetProxy(ctx context.Context, id int64) (*Proxy, error) { - return s.proxyRepo.GetByID(ctx, id) -} - -func (s *adminServiceImpl) GetProxiesByIDs(ctx context.Context, ids []int64) ([]Proxy, error) { - return s.proxyRepo.ListByIDs(ctx, ids) -} - -func (s *adminServiceImpl) CreateProxy(ctx context.Context, input *CreateProxyInput) (*Proxy, error) { - // 规范化 fallback_mode - mode := input.FallbackMode - if mode == "" { - mode = FallbackModeNone - } - // 校验:mode=proxy 必须有 backup - if mode == FallbackModeProxy && input.BackupProxyID == nil { - return nil, infraerrors.BadRequest("PROXY_BACKUP_REQUIRED", "backup proxy required when fallback_mode=proxy") - } - if input.ExpiryWarnDays < 0 { - return nil, infraerrors.BadRequest("PROXY_WARN_DAYS_INVALID", "expiry_warn_days must be >= 0") - } - - proxy := &Proxy{ - Name: input.Name, - Protocol: input.Protocol, - Host: input.Host, - Port: input.Port, - Username: input.Username, - Password: input.Password, - Status: StatusActive, - ExpiresAt: input.ExpiresAt, - FallbackMode: mode, - BackupProxyID: input.BackupProxyID, - ExpiryWarnDays: input.ExpiryWarnDays, - } - if err := s.proxyRepo.Create(ctx, proxy); err != nil { - return nil, err - } - // Probe latency asynchronously so creation isn't blocked by network timeout. - go s.probeProxyLatency(context.Background(), proxy) - return proxy, nil -} - -func (s *adminServiceImpl) UpdateProxy(ctx context.Context, id int64, input *UpdateProxyInput) (*Proxy, error) { - // 校验:backup_proxy_id 不能是自身 - if input.BackupProxyID != nil && *input.BackupProxyID == id { - return nil, infraerrors.BadRequest("PROXY_BACKUP_SELF", "backup proxy cannot be itself") - } - // 规范化 fallback_mode - mode := input.FallbackMode - if mode == "" { - mode = FallbackModeNone - } - // 校验:mode=proxy 必须有 backup - if mode == FallbackModeProxy && input.BackupProxyID == nil { - return nil, infraerrors.BadRequest("PROXY_BACKUP_REQUIRED", "backup proxy required when fallback_mode=proxy") - } - if input.ExpiryWarnDays < 0 { - return nil, infraerrors.BadRequest("PROXY_WARN_DAYS_INVALID", "expiry_warn_days must be >= 0") - } - - proxy, err := s.proxyRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - - if input.Name != "" { - proxy.Name = input.Name - } - if input.Protocol != "" { - proxy.Protocol = input.Protocol - } - if input.Host != "" { - proxy.Host = input.Host - } - if input.Port != 0 { - proxy.Port = input.Port - } - if input.Username != "" { - proxy.Username = input.Username - } - if input.Password != "" { - proxy.Password = input.Password - } - if input.Status != "" { - proxy.Status = input.Status - } - // 透传有效期与回退字段 - proxy.ExpiresAt = input.ExpiresAt - proxy.FallbackMode = mode - proxy.BackupProxyID = input.BackupProxyID - proxy.ExpiryWarnDays = input.ExpiryWarnDays - - if err := s.proxyRepo.Update(ctx, proxy); err != nil { - return nil, err - } - return proxy, nil -} - -func (s *adminServiceImpl) DeleteProxy(ctx context.Context, id int64) error { - count, err := s.proxyRepo.CountAccountsByProxyID(ctx, id) - if err != nil { - return err - } - if count > 0 { - return ErrProxyInUse - } - return s.proxyRepo.Delete(ctx, id) -} - -func (s *adminServiceImpl) BatchDeleteProxies(ctx context.Context, ids []int64) (*ProxyBatchDeleteResult, error) { - result := &ProxyBatchDeleteResult{} - if len(ids) == 0 { - return result, nil - } - - for _, id := range ids { - count, err := s.proxyRepo.CountAccountsByProxyID(ctx, id) - if err != nil { - result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{ - ID: id, - Reason: err.Error(), - }) - continue - } - if count > 0 { - result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{ - ID: id, - Reason: ErrProxyInUse.Error(), - }) - continue - } - if err := s.proxyRepo.Delete(ctx, id); err != nil { - result.Skipped = append(result.Skipped, ProxyBatchDeleteSkipped{ - ID: id, - Reason: err.Error(), - }) - continue - } - result.DeletedIDs = append(result.DeletedIDs, id) - } - - return result, nil -} - -func (s *adminServiceImpl) GetProxyAccounts(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error) { - return s.proxyRepo.ListAccountSummariesByProxyID(ctx, proxyID) -} - -func (s *adminServiceImpl) CheckProxyExists(ctx context.Context, host string, port int, username, password string) (bool, error) { - return s.proxyRepo.ExistsByHostPortAuth(ctx, host, port, username, password) -} - -// Redeem code management implementations -func (s *adminServiceImpl) ListRedeemCodes(ctx context.Context, page, pageSize int, codeType, status, search string, sortBy, sortOrder string) ([]RedeemCode, int64, error) { - params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} - codes, result, err := s.redeemCodeRepo.ListWithFilters(ctx, params, codeType, status, search) - if err != nil { - return nil, 0, err - } - return codes, result.Total, nil -} - -func (s *adminServiceImpl) GetRedeemCode(ctx context.Context, id int64) (*RedeemCode, error) { - return s.redeemCodeRepo.GetByID(ctx, id) -} - -func (s *adminServiceImpl) GenerateRedeemCodes(ctx context.Context, input *GenerateRedeemCodesInput) ([]RedeemCode, error) { - if input.ExpiresAt != nil && !input.ExpiresAt.After(time.Now()) { - return nil, ErrRedeemCodeExpired - } - - // 如果是订阅类型,验证必须有 GroupID - if input.Type == RedeemTypeSubscription { - if input.GroupID == nil { - return nil, errors.New("group_id is required for subscription type") - } - // 验证分组存在且为订阅类型 - group, err := s.groupRepo.GetByID(ctx, *input.GroupID) - if err != nil { - return nil, fmt.Errorf("group not found: %w", err) - } - if !group.IsSubscriptionType() { - return nil, errors.New("group must be subscription type") - } - } - - codes := make([]RedeemCode, 0, input.Count) - for i := 0; i < input.Count; i++ { - codeValue, err := GenerateRedeemCode() - if err != nil { - return nil, err - } - code := RedeemCode{ - Code: codeValue, - Type: input.Type, - Value: input.Value, - Status: StatusUnused, - ExpiresAt: input.ExpiresAt, - } - // 订阅类型专用字段 - if input.Type == RedeemTypeSubscription { - code.GroupID = input.GroupID - code.ValidityDays = input.ValidityDays - if code.ValidityDays <= 0 { - code.ValidityDays = 30 // 默认30天 - } - } - if err := s.redeemCodeRepo.Create(ctx, &code); err != nil { - return nil, err - } - codes = append(codes, code) - } - return codes, nil -} - -func (s *adminServiceImpl) DeleteRedeemCode(ctx context.Context, id int64) error { - return s.redeemCodeRepo.Delete(ctx, id) -} - -func (s *adminServiceImpl) BatchDeleteRedeemCodes(ctx context.Context, ids []int64) (int64, error) { - var deleted int64 - for _, id := range ids { - if err := s.redeemCodeRepo.Delete(ctx, id); err == nil { - deleted++ - } - } - return deleted, nil -} - -func (s *adminServiceImpl) ExpireRedeemCode(ctx context.Context, id int64) (*RedeemCode, error) { - code, err := s.redeemCodeRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - code.Status = StatusExpired - if err := s.redeemCodeRepo.Update(ctx, code); err != nil { - return nil, err - } - return code, nil -} - -func (s *adminServiceImpl) TestProxy(ctx context.Context, id int64) (*ProxyTestResult, error) { - proxy, err := s.proxyRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - - proxyURL := proxy.URL() - exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxyURL) - if err != nil { - s.saveProxyLatency(ctx, id, &ProxyLatencyInfo{ - Success: false, - Message: err.Error(), - UpdatedAt: time.Now(), - }) - return &ProxyTestResult{ - Success: false, - Message: err.Error(), - }, nil - } - - latency := latencyMs - s.saveProxyLatency(ctx, id, &ProxyLatencyInfo{ - Success: true, - LatencyMs: &latency, - Message: "Proxy is accessible", - IPAddress: exitInfo.IP, - Country: exitInfo.Country, - CountryCode: exitInfo.CountryCode, - Region: exitInfo.Region, - City: exitInfo.City, - UpdatedAt: time.Now(), - }) - return &ProxyTestResult{ - Success: true, - Message: "Proxy is accessible", - LatencyMs: latencyMs, - IPAddress: exitInfo.IP, - City: exitInfo.City, - Region: exitInfo.Region, - Country: exitInfo.Country, - CountryCode: exitInfo.CountryCode, - }, nil -} - -func (s *adminServiceImpl) CheckProxyQuality(ctx context.Context, id int64) (*ProxyQualityCheckResult, error) { - proxy, err := s.proxyRepo.GetByID(ctx, id) - if err != nil { - return nil, err - } - - result := &ProxyQualityCheckResult{ - ProxyID: id, - Score: 100, - Grade: "A", - CheckedAt: time.Now().Unix(), - Items: make([]ProxyQualityCheckItem, 0, len(proxyQualityTargets)+1), - } - - proxyURL := proxy.URL() - if s.proxyProber == nil { - result.Items = append(result.Items, ProxyQualityCheckItem{ - Target: "base_connectivity", - Status: "fail", - Message: "代理探测服务未配置", - }) - result.FailedCount++ - finalizeProxyQualityResult(result) - s.saveProxyQualitySnapshot(ctx, id, result, nil) - return result, nil - } - - exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxyURL) - if err != nil { - result.Items = append(result.Items, ProxyQualityCheckItem{ - Target: "base_connectivity", - Status: "fail", - LatencyMs: latencyMs, - Message: err.Error(), - }) - result.FailedCount++ - finalizeProxyQualityResult(result) - s.saveProxyQualitySnapshot(ctx, id, result, nil) - return result, nil - } - - result.ExitIP = exitInfo.IP - result.Country = exitInfo.Country - result.CountryCode = exitInfo.CountryCode - result.BaseLatencyMs = latencyMs - result.Items = append(result.Items, ProxyQualityCheckItem{ - Target: "base_connectivity", - Status: "pass", - LatencyMs: latencyMs, - Message: "代理出口连通正常", - }) - result.PassedCount++ - - client, err := httpclient.GetClient(httpclient.Options{ - ProxyURL: proxyURL, - Timeout: proxyQualityRequestTimeout, - ResponseHeaderTimeout: proxyQualityResponseHeaderTimeout, - }) - if err != nil { - result.Items = append(result.Items, ProxyQualityCheckItem{ - Target: "http_client", - Status: "fail", - Message: fmt.Sprintf("创建检测客户端失败: %v", err), - }) - result.FailedCount++ - finalizeProxyQualityResult(result) - s.saveProxyQualitySnapshot(ctx, id, result, exitInfo) - return result, nil - } - - for _, target := range proxyQualityTargets { - item := runProxyQualityTarget(ctx, client, target) - result.Items = append(result.Items, item) - switch item.Status { - case "pass": - result.PassedCount++ - case "warn": - result.WarnCount++ - case "challenge": - result.ChallengeCount++ - default: - result.FailedCount++ - } - } - - finalizeProxyQualityResult(result) - s.saveProxyQualitySnapshot(ctx, id, result, exitInfo) - return result, nil -} - -func runProxyQualityTarget(ctx context.Context, client *http.Client, target proxyQualityTarget) ProxyQualityCheckItem { - item := ProxyQualityCheckItem{ - Target: target.Target, - } - - req, err := http.NewRequestWithContext(ctx, target.Method, target.URL, nil) - if err != nil { - item.Status = "fail" - item.Message = fmt.Sprintf("构建请求失败: %v", err) - return item - } - req.Header.Set("Accept", "application/json,text/html,*/*") - req.Header.Set("User-Agent", proxyQualityClientUserAgent) - - start := time.Now() - resp, err := client.Do(req) - if err != nil { - item.Status = "fail" - item.LatencyMs = time.Since(start).Milliseconds() - item.Message = fmt.Sprintf("请求失败: %v", err) - return item - } - defer func() { _ = resp.Body.Close() }() - item.LatencyMs = time.Since(start).Milliseconds() - item.HTTPStatus = resp.StatusCode - - body, readErr := io.ReadAll(io.LimitReader(resp.Body, proxyQualityMaxBodyBytes+1)) - if readErr != nil { - item.Status = "fail" - item.Message = fmt.Sprintf("读取响应失败: %v", readErr) - return item - } - if int64(len(body)) > proxyQualityMaxBodyBytes { - body = body[:proxyQualityMaxBodyBytes] - } - - // Cloudflare challenge 检测 - if httputil.IsCloudflareChallengeResponse(resp.StatusCode, resp.Header, body) { - item.Status = "challenge" - item.CFRay = httputil.ExtractCloudflareRayID(resp.Header, body) - item.Message = "命中 Cloudflare challenge" - return item - } - - if _, ok := target.AllowedStatuses[resp.StatusCode]; ok { - // 白名单内的状态码均代表目标可达:2xx 表示接口直接可用, - // 401/405 等是无鉴权探测的预期结果,同样视为连通正常,不再扣分。 - item.Status = "pass" - if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices { - item.Message = fmt.Sprintf("HTTP %d", resp.StatusCode) - } else { - item.Message = fmt.Sprintf("HTTP %d(目标可达)", resp.StatusCode) - } - return item - } - - if resp.StatusCode == http.StatusTooManyRequests { - item.Status = "warn" - item.Message = "目标返回 429,可能存在频控" - return item - } - - item.Status = "fail" - item.Message = fmt.Sprintf("非预期状态码: %d", resp.StatusCode) - return item -} - -func finalizeProxyQualityResult(result *ProxyQualityCheckResult) { - if result == nil { - return - } - score := 100 - result.WarnCount*10 - result.FailedCount*22 - result.ChallengeCount*30 - if score < 0 { - score = 0 - } - result.Score = score - result.Grade = proxyQualityGrade(score) - result.Summary = fmt.Sprintf( - "通过 %d 项,告警 %d 项,失败 %d 项,挑战 %d 项", - result.PassedCount, - result.WarnCount, - result.FailedCount, - result.ChallengeCount, - ) -} - -func proxyQualityGrade(score int) string { - switch { - case score >= 90: - return "A" - case score >= 75: - return "B" - case score >= 60: - return "C" - case score >= 40: - return "D" - default: - return "F" - } -} - -func proxyQualityOverallStatus(result *ProxyQualityCheckResult) string { - if result == nil { - return "" - } - if result.ChallengeCount > 0 { - return "challenge" - } - if result.FailedCount > 0 { - return "failed" - } - if result.WarnCount > 0 { - return "warn" - } - if result.PassedCount > 0 { - return "healthy" - } - return "failed" -} - -func proxyQualityFirstCFRay(result *ProxyQualityCheckResult) string { - if result == nil { - return "" - } - for _, item := range result.Items { - if item.CFRay != "" { - return item.CFRay - } - } - return "" -} - -func proxyQualityBaseConnectivityPass(result *ProxyQualityCheckResult) bool { - if result == nil { - return false - } - for _, item := range result.Items { - if item.Target == "base_connectivity" { - return item.Status == "pass" - } - } - return false -} - -func (s *adminServiceImpl) saveProxyQualitySnapshot(ctx context.Context, proxyID int64, result *ProxyQualityCheckResult, exitInfo *ProxyExitInfo) { - if result == nil { - return - } - score := result.Score - checkedAt := result.CheckedAt - info := &ProxyLatencyInfo{ - Success: proxyQualityBaseConnectivityPass(result), - Message: result.Summary, - QualityStatus: proxyQualityOverallStatus(result), - QualityScore: &score, - QualityGrade: result.Grade, - QualitySummary: result.Summary, - QualityCheckedAt: &checkedAt, - QualityCFRay: proxyQualityFirstCFRay(result), - UpdatedAt: time.Now(), - } - if result.BaseLatencyMs > 0 { - latency := result.BaseLatencyMs - info.LatencyMs = &latency - } - if exitInfo != nil { - info.IPAddress = exitInfo.IP - info.Country = exitInfo.Country - info.CountryCode = exitInfo.CountryCode - info.Region = exitInfo.Region - info.City = exitInfo.City - } - s.saveProxyLatency(ctx, proxyID, info) -} - -func (s *adminServiceImpl) probeProxyLatency(ctx context.Context, proxy *Proxy) { - if s.proxyProber == nil || proxy == nil { - return - } - exitInfo, latencyMs, err := s.proxyProber.ProbeProxy(ctx, proxy.URL()) - if err != nil { - s.saveProxyLatency(ctx, proxy.ID, &ProxyLatencyInfo{ - Success: false, - Message: err.Error(), - UpdatedAt: time.Now(), - }) - return - } - - latency := latencyMs - s.saveProxyLatency(ctx, proxy.ID, &ProxyLatencyInfo{ - Success: true, - LatencyMs: &latency, - Message: "Proxy is accessible", - IPAddress: exitInfo.IP, - Country: exitInfo.Country, - CountryCode: exitInfo.CountryCode, - Region: exitInfo.Region, - City: exitInfo.City, - UpdatedAt: time.Now(), - }) -} - -// checkMixedChannelRisk 检查分组中是否存在混合渠道(Antigravity + Anthropic) -// 如果存在混合,返回错误提示用户确认 -func (s *adminServiceImpl) checkMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error { - // 判断当前账号的渠道类型(基于 platform 字段,而不是 type 字段) - currentPlatform := getAccountPlatform(currentAccountPlatform) - if currentPlatform == "" { - // 不是 Antigravity 或 Anthropic,无需检查 - return nil - } - - // 检查每个分组中的其他账号 - for _, groupID := range groupIDs { - accounts, err := s.accountRepo.ListByGroup(ctx, groupID) - if err != nil { - return fmt.Errorf("get accounts in group %d: %w", groupID, err) - } - - // 检查是否存在不同渠道的账号 - for _, account := range accounts { - if currentAccountID > 0 && account.ID == currentAccountID { - continue // 跳过当前账号 - } - - otherPlatform := getAccountPlatform(account.Platform) - if otherPlatform == "" { - continue // 不是 Antigravity 或 Anthropic,跳过 - } - - // 检测混合渠道 - if currentPlatform != otherPlatform { - group, _ := s.groupRepo.GetByID(ctx, groupID) - groupName := fmt.Sprintf("Group %d", groupID) - if group != nil { - groupName = group.Name - } - - return &MixedChannelError{ - GroupID: groupID, - GroupName: groupName, - CurrentPlatform: currentPlatform, - OtherPlatform: otherPlatform, - } - } - } - } - - return nil -} - -func (s *adminServiceImpl) validateGroupIDsExist(ctx context.Context, groupIDs []int64) error { - if len(groupIDs) == 0 { - return nil - } - if s.groupRepo == nil { - return errors.New("group repository not configured") - } - - if batchReader, ok := s.groupRepo.(groupExistenceBatchReader); ok { - existsByID, err := batchReader.ExistsByIDs(ctx, groupIDs) - if err != nil { - return fmt.Errorf("check groups exists: %w", err) - } - for _, groupID := range groupIDs { - if groupID <= 0 || !existsByID[groupID] { - return fmt.Errorf("get group: %w", ErrGroupNotFound) - } - } - return nil - } - - for _, groupID := range groupIDs { - if _, err := s.groupRepo.GetByID(ctx, groupID); err != nil { - return fmt.Errorf("get group: %w", err) - } - } - return nil -} - -// CheckMixedChannelRisk checks whether target groups contain mixed channels for the current account platform. -func (s *adminServiceImpl) CheckMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error { - return s.checkMixedChannelRisk(ctx, currentAccountID, currentAccountPlatform, groupIDs) -} - -func (s *adminServiceImpl) attachProxyLatency(ctx context.Context, proxies []ProxyWithAccountCount) { - if s.proxyLatencyCache == nil || len(proxies) == 0 { - return - } - - ids := make([]int64, 0, len(proxies)) - for i := range proxies { - ids = append(ids, proxies[i].ID) - } - - latencies, err := s.proxyLatencyCache.GetProxyLatencies(ctx, ids) - if err != nil { - logger.LegacyPrintf("service.admin", "Warning: load proxy latency cache failed: %v", err) - return - } - - for i := range proxies { - info := latencies[proxies[i].ID] - if info == nil { - continue - } - if info.Success { - proxies[i].LatencyStatus = "success" - proxies[i].LatencyMs = info.LatencyMs - } else { - proxies[i].LatencyStatus = "failed" - } - proxies[i].LatencyMessage = info.Message - proxies[i].IPAddress = info.IPAddress - proxies[i].Country = info.Country - proxies[i].CountryCode = info.CountryCode - proxies[i].Region = info.Region - proxies[i].City = info.City - proxies[i].QualityStatus = info.QualityStatus - proxies[i].QualityScore = info.QualityScore - proxies[i].QualityGrade = info.QualityGrade - proxies[i].QualitySummary = info.QualitySummary - proxies[i].QualityChecked = info.QualityCheckedAt - } -} - -func (s *adminServiceImpl) saveProxyLatency(ctx context.Context, proxyID int64, info *ProxyLatencyInfo) { - if s.proxyLatencyCache == nil || info == nil { - return - } - - merged := *info - if latencies, err := s.proxyLatencyCache.GetProxyLatencies(ctx, []int64{proxyID}); err == nil { - if existing := latencies[proxyID]; existing != nil { - if merged.QualityCheckedAt == nil && - merged.QualityScore == nil && - merged.QualityGrade == "" && - merged.QualityStatus == "" && - merged.QualitySummary == "" && - merged.QualityCFRay == "" { - merged.QualityStatus = existing.QualityStatus - merged.QualityScore = existing.QualityScore - merged.QualityGrade = existing.QualityGrade - merged.QualitySummary = existing.QualitySummary - merged.QualityCheckedAt = existing.QualityCheckedAt - merged.QualityCFRay = existing.QualityCFRay - } - } - } - - if err := s.proxyLatencyCache.SetProxyLatency(ctx, proxyID, &merged); err != nil { - logger.LegacyPrintf("service.admin", "Warning: store proxy latency cache failed: %v", err) - } -} - -// getAccountPlatform 根据账号 platform 判断混合渠道检查用的平台标识 -func getAccountPlatform(accountPlatform string) string { - switch strings.ToLower(strings.TrimSpace(accountPlatform)) { - case PlatformAntigravity: - return "Antigravity" - case PlatformAnthropic, "claude": - return "Anthropic" - default: - return "" - } -} - -// MixedChannelError 混合渠道错误 -type MixedChannelError struct { - GroupID int64 - GroupName string - CurrentPlatform string - OtherPlatform string -} - -func (e *MixedChannelError) Error() string { - return fmt.Sprintf("mixed_channel_warning: Group '%s' contains both %s and %s accounts. Using mixed channels in the same context may cause thinking block signature validation issues, which will fallback to non-thinking mode for historical messages.", - e.GroupName, e.CurrentPlatform, e.OtherPlatform) -} - -func (s *adminServiceImpl) ResetAccountQuota(ctx context.Context, id int64) error { - account, err := s.accountRepo.GetByID(ctx, id) - if err != nil { - return err - } - // spark 影子账号不持自有配额(凭据透传母账号、spark 用量走独立 codex_* 维度由 QueryUsage 维护), - // 通用 quota 重置对其无意义且语义不一致——明确 400 拒绝(与 OpenAI reset-credit 对影子一致)(外审第7轮 P2)。 - if account.IsCredentialShadow() { - return infraerrors.New(http.StatusBadRequest, "SPARK_SHADOW_NO_QUOTA_RESET", - "cannot reset quota for a spark shadow account; manage it on the parent account") - } - return s.accountRepo.ResetQuotaUsed(ctx, id) -} - -// EnsureOpenAIPrivacy 检查 OpenAI OAuth 账号是否已设置 privacy_mode, -// 未设置则调用 disableOpenAITraining 并持久化到 Extra,返回设置的 mode 值。 -func (s *adminServiceImpl) EnsureOpenAIPrivacy(ctx context.Context, account *Account) string { - // 影子账号不持凭据,隐私设置由母账号管理,直接跳过。 - if account.IsCredentialShadow() { - return "" - } - if account.Platform != PlatformOpenAI || account.Type != AccountTypeOAuth { - return "" - } - if s.privacyClientFactory == nil { - return "" - } - if shouldSkipOpenAIPrivacyEnsure(account.Extra) { - return "" - } - - token, _ := account.Credentials["access_token"].(string) - if token == "" { - return "" - } - - var proxyURL string - if account.ProxyID != nil { - if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { - proxyURL = p.URL() - } - } - - mode := disableOpenAITraining(ctx, s.privacyClientFactory, token, proxyURL) - if mode == "" { - return "" - } - - _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}) - return mode -} - -// ForceOpenAIPrivacy 强制重新设置 OpenAI OAuth 账号隐私,无论当前状态。 -func (s *adminServiceImpl) ForceOpenAIPrivacy(ctx context.Context, account *Account) string { - // 影子账号不持凭据,隐私由母账号管理,直接跳过(与 EnsureOpenAIPrivacy 一致——外审第4轮)。 - if account.IsCredentialShadow() { - return "" - } - if account.Platform != PlatformOpenAI || account.Type != AccountTypeOAuth { - return "" - } - if s.privacyClientFactory == nil { - return "" - } - - token, _ := account.Credentials["access_token"].(string) - if token == "" { - return "" - } - - var proxyURL string - if account.ProxyID != nil { - if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { - proxyURL = p.URL() - } - } - - mode := disableOpenAITraining(ctx, s.privacyClientFactory, token, proxyURL) - if mode == "" { - return "" - } - - if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil { - logger.LegacyPrintf("service.admin", "force_update_openai_privacy_mode_failed: account_id=%d err=%v", account.ID, err) - return mode - } - if account.Extra == nil { - account.Extra = make(map[string]any) - } - account.Extra["privacy_mode"] = mode - return mode -} - -// EnsureAntigravityPrivacy 检查 Antigravity OAuth 账号隐私状态。 -// 仅当 privacy_mode 已成功设置("privacy_set")时跳过; -// 未设置或之前失败("privacy_set_failed")均会重试。 -func (s *adminServiceImpl) EnsureAntigravityPrivacy(ctx context.Context, account *Account) string { - if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth { - return "" - } - if account.Extra != nil { - if existing, ok := account.Extra["privacy_mode"].(string); ok && existing == AntigravityPrivacySet { - return existing - } - } - - token, _ := account.Credentials["access_token"].(string) - if token == "" { - return "" - } - - projectID, _ := account.Credentials["project_id"].(string) - - var proxyURL string - if account.ProxyID != nil { - if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { - proxyURL = p.URL() - } - } - - mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL) - if mode == "" { - return "" - } - - if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil { - logger.LegacyPrintf("service.admin", "update_antigravity_privacy_mode_failed: account_id=%d err=%v", account.ID, err) - return mode - } - applyAntigravityPrivacyMode(account, mode) - return mode -} - -// ForceAntigravityPrivacy 强制重新设置 Antigravity OAuth 账号隐私,无论当前状态。 -func (s *adminServiceImpl) ForceAntigravityPrivacy(ctx context.Context, account *Account) string { - if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth { - return "" - } - - token, _ := account.Credentials["access_token"].(string) - if token == "" { - return "" - } - - projectID, _ := account.Credentials["project_id"].(string) - - var proxyURL string - if account.ProxyID != nil { - if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { - proxyURL = p.URL() - } - } - - mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL) - if mode == "" { - return "" - } - - if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil { - logger.LegacyPrintf("service.admin", "force_update_antigravity_privacy_mode_failed: account_id=%d err=%v", account.ID, err) - return mode - } - applyAntigravityPrivacyMode(account, mode) - return mode -} diff --git a/backend/internal/service/admin_service_group_test.go b/backend/internal/service/admin_service_group_test.go index 0b360c61c7..7731b33bb6 100644 --- a/backend/internal/service/admin_service_group_test.go +++ b/backend/internal/service/admin_service_group_test.go @@ -174,6 +174,42 @@ func TestAdminService_CreateGroup_WithImagePricing(t *testing.T) { require.InDelta(t, 0.30, *repo.created.ImagePrice4K, 0.0001) } +func TestAdminService_CreateGroup_WithVideoPricing(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} + + price480P := 0.08 + price720P := 0.12 + price1080P := 0.18 + videoMultiplier := 0.75 + + input := &CreateGroupInput{ + Name: "grok-video", + Description: "Grok video group", + Platform: PlatformGrok, + RateMultiplier: 1.0, + VideoRateIndependent: true, + VideoRateMultiplier: &videoMultiplier, + VideoPrice480P: &price480P, + VideoPrice720P: &price720P, + VideoPrice1080P: &price1080P, + } + + group, err := svc.CreateGroup(context.Background(), input) + require.NoError(t, err) + require.NotNil(t, group) + + require.NotNil(t, repo.created) + require.True(t, repo.created.VideoRateIndependent) + require.InDelta(t, 0.75, repo.created.VideoRateMultiplier, 1e-12) + require.NotNil(t, repo.created.VideoPrice480P) + require.NotNil(t, repo.created.VideoPrice720P) + require.NotNil(t, repo.created.VideoPrice1080P) + require.InDelta(t, 0.08, *repo.created.VideoPrice480P, 0.0001) + require.InDelta(t, 0.12, *repo.created.VideoPrice720P, 0.0001) + require.InDelta(t, 0.18, *repo.created.VideoPrice1080P, 0.0001) +} + // TestAdminService_CreateGroup_NilImagePricing 测试 ImagePrice 为 nil 时正常创建 func TestAdminService_CreateGroup_NilImagePricing(t *testing.T) { repo := &groupRepoStubForAdmin{} @@ -232,6 +268,46 @@ func TestAdminService_CreateGroup_PreservesNonGrokImageGenerationDisabled(t *tes require.False(t, group.AllowImageGeneration) } +func TestAdminService_CreateGroup_DisablesBatchImageWhenImageGenerationDisabled(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} + + group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ + Name: "gemini-no-image", + Description: "Gemini group without image generation", + Platform: PlatformGemini, + RateMultiplier: 1.0, + AllowImageGeneration: false, + AllowBatchImageGeneration: true, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.created) + require.False(t, repo.created.AllowImageGeneration) + require.False(t, repo.created.AllowBatchImageGeneration) + require.False(t, group.AllowBatchImageGeneration) +} + +func TestAdminService_CreateGroup_DisablesBatchImageForNonGeminiPlatform(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} + + group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ + Name: "openai-image", + Description: "OpenAI image group", + Platform: PlatformOpenAI, + RateMultiplier: 1.0, + AllowImageGeneration: true, + AllowBatchImageGeneration: true, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.created) + require.True(t, repo.created.AllowImageGeneration) + require.False(t, repo.created.AllowBatchImageGeneration) + require.False(t, group.AllowBatchImageGeneration) +} + // TestAdminService_UpdateGroup_WithImagePricing 测试更新分组时 ImagePrice 字段正确更新 func TestAdminService_UpdateGroup_WithImagePricing(t *testing.T) { existingGroup := &Group{ @@ -267,6 +343,42 @@ func TestAdminService_UpdateGroup_WithImagePricing(t *testing.T) { require.InDelta(t, 0.36, *repo.updated.ImagePrice4K, 0.0001) } +func TestAdminService_UpdateGroup_WithVideoPricing(t *testing.T) { + existingGroup := &Group{ + ID: 1, + Name: "existing-grok", + Platform: PlatformGrok, + Status: StatusActive, + } + repo := &groupRepoStubForAdmin{getByID: existingGroup} + svc := &adminServiceImpl{groupRepo: repo} + + price480P := 0.09 + price720P := 0.13 + price1080P := 0.19 + videoMultiplier := 0.6 + independent := true + + input := &UpdateGroupInput{ + VideoRateIndependent: &independent, + VideoRateMultiplier: &videoMultiplier, + VideoPrice480P: &price480P, + VideoPrice720P: &price720P, + VideoPrice1080P: &price1080P, + } + + group, err := svc.UpdateGroup(context.Background(), 1, input) + require.NoError(t, err) + require.NotNil(t, group) + + require.NotNil(t, repo.updated) + require.True(t, repo.updated.VideoRateIndependent) + require.InDelta(t, 0.6, repo.updated.VideoRateMultiplier, 1e-12) + require.InDelta(t, 0.09, *repo.updated.VideoPrice480P, 0.0001) + require.InDelta(t, 0.13, *repo.updated.VideoPrice720P, 0.0001) + require.InDelta(t, 0.19, *repo.updated.VideoPrice1080P, 0.0001) +} + // TestAdminService_UpdateGroup_PartialImagePricing 测试仅更新部分 ImagePrice 字段 func TestAdminService_UpdateGroup_PartialImagePricing(t *testing.T) { oldPrice2K := 0.15 @@ -326,6 +438,53 @@ func TestAdminService_UpdateGroup_PreservesImageGenerationControlsWhenOmitted(t require.InDelta(t, 0.5, repo.updated.ImageRateMultiplier, 1e-12) } +func TestAdminService_UpdateGroup_DisablesBatchImageWhenImageGenerationDisabled(t *testing.T) { + existingGroup := &Group{ + ID: 1, + Name: "existing-gemini", + Platform: PlatformGemini, + Status: StatusActive, + AllowImageGeneration: true, + AllowBatchImageGeneration: true, + } + repo := &groupRepoStubForAdmin{getByID: existingGroup} + svc := &adminServiceImpl{groupRepo: repo} + disabled := false + + group, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{ + AllowImageGeneration: &disabled, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.updated) + require.False(t, repo.updated.AllowImageGeneration) + require.False(t, repo.updated.AllowBatchImageGeneration) + require.False(t, group.AllowBatchImageGeneration) +} + +func TestAdminService_UpdateGroup_DisablesBatchImageWhenPlatformChangesFromGemini(t *testing.T) { + existingGroup := &Group{ + ID: 1, + Name: "existing-gemini", + Platform: PlatformGemini, + Status: StatusActive, + AllowImageGeneration: true, + AllowBatchImageGeneration: true, + } + repo := &groupRepoStubForAdmin{getByID: existingGroup} + svc := &adminServiceImpl{groupRepo: repo} + + group, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{ + Platform: PlatformOpenAI, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.updated) + require.Equal(t, PlatformOpenAI, repo.updated.Platform) + require.False(t, repo.updated.AllowBatchImageGeneration) + require.False(t, group.AllowBatchImageGeneration) +} + func TestAdminService_UpdateGroup_ClearsDescriptionWhenEmptyString(t *testing.T) { existingGroup := &Group{ ID: 1, @@ -384,6 +543,96 @@ func TestAdminService_UpdateGroup_RejectsNegativeImageRateMultiplier(t *testing. require.Nil(t, repo.updated) } +func TestAdminService_CreateGroup_BatchImagePricingSettings(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} + discount := 0.8 + hold := 0.9 + + group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ + Name: "batch-image-pricing", + Platform: PlatformGemini, + RateMultiplier: 1, + BatchImageDiscountMultiplier: &discount, + BatchImageHoldMultiplier: &hold, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.created) + require.InDelta(t, 0.8, repo.created.BatchImageDiscountMultiplier, 1e-12) + require.InDelta(t, 0.9, repo.created.BatchImageHoldMultiplier, 1e-12) +} + +func TestAdminService_CreateGroup_RejectsHoldBelowDiscount(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} + discount := 0.8 + hold := 0.6 + + // hold < discount 时,成功率足够高的批量任务实际成本会超过冻结额, + // 结算永远失败,必须在配置入口拒绝。 + _, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ + Name: "batch-image-pricing-invalid", + Platform: PlatformGemini, + RateMultiplier: 1, + BatchImageDiscountMultiplier: &discount, + BatchImageHoldMultiplier: &hold, + }) + require.Error(t, err) + require.Nil(t, repo.created) +} + +func TestAdminService_GroupBatchImagePricingValidation(t *testing.T) { + tests := []struct { + name string + input *CreateGroupInput + }{ + { + name: "negative_discount", + input: func() *CreateGroupInput { + v := -0.1 + return &CreateGroupInput{Name: "bad-discount", RateMultiplier: 1, BatchImageDiscountMultiplier: &v} + }(), + }, + { + name: "negative_hold", + input: func() *CreateGroupInput { + v := -0.1 + return &CreateGroupInput{Name: "bad-hold", RateMultiplier: 1, BatchImageHoldMultiplier: &v} + }(), + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} + + _, err := svc.CreateGroup(context.Background(), tt.input) + require.Error(t, err) + require.Nil(t, repo.created) + }) + } +} + +func TestAdminService_UpdateGroup_RejectsNegativeVideoRateMultiplier(t *testing.T) { + existingGroup := &Group{ + ID: 1, + Name: "existing-group", + Platform: PlatformGrok, + Status: StatusActive, + VideoRateMultiplier: 1, + } + repo := &groupRepoStubForAdmin{getByID: existingGroup} + svc := &adminServiceImpl{groupRepo: repo} + negative := -0.1 + + _, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{ + VideoRateMultiplier: &negative, + }) + require.Error(t, err) + require.Nil(t, repo.updated) +} + func TestAdminService_UpdateGroup_InvalidatesAuthCacheOnRPMLimitChange(t *testing.T) { existingGroup := &Group{ ID: 1, diff --git a/backend/internal/service/admin_service_role_test.go b/backend/internal/service/admin_service_role_test.go new file mode 100644 index 0000000000..956b7fae17 --- /dev/null +++ b/backend/internal/service/admin_service_role_test.go @@ -0,0 +1,143 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/stretchr/testify/require" +) + +func TestAdminService_CreateUser_WithAdminRole(t *testing.T) { + repo := &userRepoStub{nextID: 30} + svc := &adminServiceImpl{userRepo: repo} + + user, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "admin@test.com", + Password: "strong-pass", + Role: RoleAdmin, + }) + require.NoError(t, err) + require.Equal(t, RoleAdmin, user.Role) +} + +func TestAdminService_CreateUser_DefaultsToUserRole(t *testing.T) { + repo := &userRepoStub{nextID: 31} + svc := &adminServiceImpl{userRepo: repo} + + user, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "plain@test.com", + Password: "strong-pass", + }) + require.NoError(t, err) + require.Equal(t, RoleUser, user.Role) +} + +func TestAdminService_CreateUser_InvalidRoleRejected(t *testing.T) { + repo := &userRepoStub{nextID: 32} + svc := &adminServiceImpl{userRepo: repo} + + _, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "bad@test.com", + Password: "strong-pass", + Role: "superuser", + }) + require.Error(t, err) + require.Empty(t, repo.created, "非法角色不应写入用户") +} + +func TestAdminService_UpdateUser_PromoteToAdmin(t *testing.T) { + base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleUser}} + repo := &rpmUserRepoStub{userRepoStub: base} + invalidator := &authCacheInvalidatorStub{} + svc := &adminServiceImpl{ + userRepo: repo, + redeemCodeRepo: &redeemRepoStub{}, + authCacheInvalidator: invalidator, + } + + updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleAdmin}) + require.NoError(t, err) + require.Equal(t, RoleAdmin, updated.Role) + require.Equal(t, []int64{42}, invalidator.userIDs, "角色变更应失效认证缓存") +} + +func TestAdminService_UpdateUser_RoleOmittedKeepsExisting(t *testing.T) { + base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleAdmin}} + repo := &rpmUserRepoStub{userRepoStub: base} + svc := &adminServiceImpl{userRepo: repo, redeemCodeRepo: &redeemRepoStub{}} + + newName := "renamed" + updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Username: &newName}) + require.NoError(t, err) + require.Equal(t, RoleAdmin, updated.Role, "未提供 role 时不应改变现有角色") +} + +func TestAdminService_UpdateUser_InvalidRoleRejected(t *testing.T) { + base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleUser}} + repo := &rpmUserRepoStub{userRepoStub: base} + svc := &adminServiceImpl{userRepo: repo, redeemCodeRepo: &redeemRepoStub{}} + + _, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: "root"}) + require.Error(t, err) + require.Nil(t, repo.lastUpdated, "非法角色不应触发持久化") +} + +// roleGuardUserRepoStub 在 rpmUserRepoStub 之上提供可控的管理员计数, +// 用于测试"最后一个管理员不可降级"守卫。 +type roleGuardUserRepoStub struct { + *rpmUserRepoStub + adminTotal int64 + listCalls int +} + +func (s *roleGuardUserRepoStub) ListWithFilters(_ context.Context, _ pagination.PaginationParams, _ UserListFilters) ([]User, *pagination.PaginationResult, error) { + s.listCalls++ + return nil, &pagination.PaginationResult{Total: s.adminTotal}, nil +} + +func TestAdminService_UpdateUser_DemoteLastAdminRejected(t *testing.T) { + base := &userRepoStub{user: &User{ID: 42, Email: "a@example.com", Role: RoleAdmin}} + repo := &roleGuardUserRepoStub{rpmUserRepoStub: &rpmUserRepoStub{userRepoStub: base}, adminTotal: 1} + svc := &adminServiceImpl{userRepo: repo, redeemCodeRepo: &redeemRepoStub{}} + + _, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleUser}) + require.Error(t, err) + require.Contains(t, err.Error(), "last admin") + require.Nil(t, repo.lastUpdated, "最后一个管理员不应被降级持久化") + require.Equal(t, 1, repo.listCalls, "降级路径应触发管理员计数") +} + +func TestAdminService_UpdateUser_DemoteAdminAllowedWhenOthersExist(t *testing.T) { + base := &userRepoStub{user: &User{ID: 42, Email: "a@example.com", Role: RoleAdmin}} + repo := &roleGuardUserRepoStub{rpmUserRepoStub: &rpmUserRepoStub{userRepoStub: base}, adminTotal: 2} + invalidator := &authCacheInvalidatorStub{} + svc := &adminServiceImpl{ + userRepo: repo, + redeemCodeRepo: &redeemRepoStub{}, + authCacheInvalidator: invalidator, + } + + updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleUser}) + require.NoError(t, err) + require.Equal(t, RoleUser, updated.Role) + require.NotNil(t, repo.lastUpdated) + require.Equal(t, RoleUser, repo.lastUpdated.Role, "存在其他管理员时允许降级") +} + +func TestAdminService_UpdateUser_PromoteDoesNotCountAdmins(t *testing.T) { + base := &userRepoStub{user: &User{ID: 42, Email: "u@example.com", Role: RoleUser}} + repo := &roleGuardUserRepoStub{rpmUserRepoStub: &rpmUserRepoStub{userRepoStub: base}, adminTotal: 1} + svc := &adminServiceImpl{ + userRepo: repo, + redeemCodeRepo: &redeemRepoStub{}, + authCacheInvalidator: &authCacheInvalidatorStub{}, + } + + updated, err := svc.UpdateUser(context.Background(), 42, &UpdateUserInput{Role: RoleAdmin}) + require.NoError(t, err) + require.Equal(t, RoleAdmin, updated.Role) + require.Equal(t, 0, repo.listCalls, "升级路径不应触发管理员计数") +} diff --git a/backend/internal/service/admin_user.go b/backend/internal/service/admin_user.go new file mode 100644 index 0000000000..862a5ab5c7 --- /dev/null +++ b/backend/internal/service/admin_user.go @@ -0,0 +1,1251 @@ +package service + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + "time" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/ent/authidentity" + "github.com/Wei-Shaw/sub2api/ent/authidentitychannel" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" +) + +// User management implementations +func (s *adminServiceImpl) ListUsers(ctx context.Context, page, pageSize int, filters UserListFilters, sortBy, sortOrder string) ([]User, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + users, result, err := s.userRepo.ListWithFilters(ctx, params, filters) + if err != nil { + return nil, 0, err + } + if len(users) > 0 { + userIDs := make([]int64, 0, len(users)) + for i := range users { + userIDs = append(userIDs, users[i].ID) + } + lastUsedByUserID, latestErr := s.userRepo.GetLatestUsedAtByUserIDs(ctx, userIDs) + if latestErr != nil { + logger.LegacyPrintf("service.admin", "failed to load user last_used_at in batch: err=%v", latestErr) + } else { + for i := range users { + users[i].LastUsedAt = lastUsedByUserID[users[i].ID] + } + } + } + // 批量加载用户专属分组倍率 + if s.userGroupRateRepo != nil && len(users) > 0 { + if batchRepo, ok := s.userGroupRateRepo.(userGroupRateBatchReader); ok { + userIDs := make([]int64, 0, len(users)) + for i := range users { + userIDs = append(userIDs, users[i].ID) + } + ratesByUser, err := batchRepo.GetByUserIDs(ctx, userIDs) + if err != nil { + logger.LegacyPrintf("service.admin", "failed to load user group rates in batch: err=%v", err) + s.loadUserGroupRatesOneByOne(ctx, users) + } else { + for i := range users { + if rates, ok := ratesByUser[users[i].ID]; ok { + users[i].GroupRates = rates + } + } + } + } else { + s.loadUserGroupRatesOneByOne(ctx, users) + } + } + return users, result.Total, nil +} + +func (s *adminServiceImpl) loadUserGroupRatesOneByOne(ctx context.Context, users []User) { + if s.userGroupRateRepo == nil { + return + } + for i := range users { + rates, err := s.userGroupRateRepo.GetByUserID(ctx, users[i].ID) + if err != nil { + logger.LegacyPrintf("service.admin", "failed to load user group rates: user_id=%d err=%v", users[i].ID, err) + continue + } + users[i].GroupRates = rates + } +} + +func (s *adminServiceImpl) GetUser(ctx context.Context, id int64) (*User, error) { + user, err := s.userRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + lastUsedAt, latestErr := s.userRepo.GetLatestUsedAtByUserID(ctx, id) + if latestErr != nil { + logger.LegacyPrintf("service.admin", "failed to load user last_used_at: user_id=%d err=%v", id, latestErr) + } else { + user.LastUsedAt = lastUsedAt + } + // 加载用户专属分组倍率 + if s.userGroupRateRepo != nil { + rates, err := s.userGroupRateRepo.GetByUserID(ctx, id) + if err != nil { + logger.LegacyPrintf("service.admin", "failed to load user group rates: user_id=%d err=%v", id, err) + } else { + user.GroupRates = rates + } + } + return user, nil +} + +func (s *adminServiceImpl) GetUserIncludeDeleted(ctx context.Context, id int64) (*User, error) { + return s.userRepo.GetByIDIncludeDeleted(ctx, id) +} + +// normalizeUserRole 校验并归一化角色输入。 +// 空字符串返回 fallback(未提供时的默认角色);非法值返回错误。 +func normalizeUserRole(role, fallback string) (string, error) { + if role == "" { + return fallback, nil + } + if role != RoleAdmin && role != RoleUser { + return "", fmt.Errorf("invalid role: %q (must be %s or %s)", role, RoleAdmin, RoleUser) + } + return role, nil +} + +func (s *adminServiceImpl) CreateUser(ctx context.Context, input *CreateUserInput) (*User, error) { + balance := 0.0 + if input.Balance != nil { + balance = *input.Balance + } else if s.settingService != nil { + balance = s.settingService.GetDefaultBalance(ctx) + } + + // 角色可由管理员在创建时指定(admin/user);未提供时默认 user。 + role, err := normalizeUserRole(input.Role, RoleUser) + if err != nil { + return nil, err + } + + user := &User{ + Email: input.Email, + Username: input.Username, + Notes: input.Notes, + Role: role, + Balance: balance, + Concurrency: input.Concurrency, + RPMLimit: input.RPMLimit, + Status: StatusActive, + AllowedGroups: input.AllowedGroups, + } + if err := user.SetPassword(input.Password); err != nil { + return nil, err + } + if err := s.userRepo.Create(ctx, user); err != nil { + return nil, err + } + // 创建管理员属权限敏感操作,落审计日志(含操作者),便于事后追溯。 + if user.Role == RoleAdmin { + logger.LegacyPrintf("service.admin", "audit: admin user created actor_admin_id=%d target_user_id=%d", + input.ActorAdminID, user.ID) + } + s.assignDefaultSubscriptions(ctx, user.ID) + return user, nil +} + +// ensureNotLastAdmin 降级管理员前确认系统中仍存在其他管理员,防止零 admin 锁死。 +// 注:读取与写入之间存在竞态窗口,极端并发下仍可能双双降级;作为后台低频操作 +// 的兜底保护足够,彻底防护需依赖数据库层约束。 +func (s *adminServiceImpl) ensureNotLastAdmin(ctx context.Context) error { + noSubs := false + _, result, err := s.userRepo.ListWithFilters(ctx, + pagination.PaginationParams{Page: 1, PageSize: 1}, + UserListFilters{Role: RoleAdmin, IncludeSubscriptions: &noSubs}, + ) + if err != nil { + return fmt.Errorf("count admin users: %w", err) + } + if result == nil || result.Total <= 1 { + return errors.New("cannot demote the last admin user") + } + return nil +} + +func (s *adminServiceImpl) assignDefaultSubscriptions(ctx context.Context, userID int64) { + if s.settingService == nil || s.defaultSubAssigner == nil || userID <= 0 { + return + } + items := s.settingService.GetDefaultSubscriptions(ctx) + for _, item := range items { + if _, _, err := s.defaultSubAssigner.AssignOrExtendSubscription(ctx, &AssignSubscriptionInput{ + UserID: userID, + GroupID: item.GroupID, + ValidityDays: item.ValidityDays, + Notes: "auto assigned by default user subscriptions setting", + }); err != nil { + logger.LegacyPrintf("service.admin", "failed to assign default subscription: user_id=%d group_id=%d err=%v", userID, item.GroupID, err) + } + } +} + +func (s *adminServiceImpl) UpdateUser(ctx context.Context, id int64, input *UpdateUserInput) (*User, error) { + // 校验用户专属分组倍率:必须 > 0(nil 合法,表示清除专属倍率) + if input.GroupRates != nil { + for groupID, rate := range input.GroupRates { + if rate != nil && *rate <= 0 { + return nil, fmt.Errorf("rate_multiplier must be > 0 (group_id=%d)", groupID) + } + } + } + + user, err := s.userRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + + // Protect admin users: cannot disable admin accounts + if user.Role == "admin" && input.Status == "disabled" { + return nil, errors.New("cannot disable admin user") + } + + oldConcurrency := user.Concurrency + oldStatus := user.Status + oldRole := user.Role + oldRPMLimit := user.RPMLimit + oldAllowedGroups := append([]int64(nil), user.AllowedGroups...) + + if input.Email != "" { + user.Email = input.Email + } + if input.Password != "" { + if err := user.SetPassword(input.Password); err != nil { + return nil, err + } + } + + if input.Username != nil { + user.Username = *input.Username + } + if input.Notes != nil { + user.Notes = *input.Notes + } + + if input.Status != "" { + user.Status = input.Status + } + + // 角色变更(admin/user);空字符串表示不修改。 + if input.Role != "" { + role, err := normalizeUserRole(input.Role, user.Role) + if err != nil { + return nil, err + } + // 防锁死保护:不允许降级系统中最后一个管理员(自我降级已在 handler 层拦截, + // 此处兜底覆盖跨管理员互降导致零 admin 的场景)。 + if user.Role == RoleAdmin && role == RoleUser { + if err := s.ensureNotLastAdmin(ctx); err != nil { + return nil, err + } + } + user.Role = role + } + + if input.Concurrency != nil { + user.Concurrency = *input.Concurrency + } + + if input.RPMLimit != nil { + user.RPMLimit = *input.RPMLimit + } + + if input.AllowedGroups != nil { + user.AllowedGroups = *input.AllowedGroups + } + + if err := s.userRepo.Update(ctx, user); err != nil { + return nil, err + } + + // 角色变更属权限敏感操作,落审计日志(含操作者),便于事后追溯。 + if user.Role != oldRole { + logger.LegacyPrintf("service.admin", "audit: user role changed actor_admin_id=%d target_user_id=%d old_role=%s new_role=%s", + input.ActorAdminID, user.ID, oldRole, user.Role) + } + + // 同步用户专属分组倍率 + if input.GroupRates != nil && s.userGroupRateRepo != nil { + if err := s.userGroupRateRepo.SyncUserGroupRates(ctx, user.ID, input.GroupRates); err != nil { + logger.LegacyPrintf("service.admin", "failed to sync user group rates: user_id=%d err=%v", user.ID, err) + } + } + + if s.authCacheInvalidator != nil { + // RPMLimit 直接参与 billing_cache_service.checkRPM 的三级级联, + // allowed_groups 参与 API Key 专属分组授权判断;不失效缓存会让修改在一个 L2 TTL 内失去效果。 + if user.Concurrency != oldConcurrency || user.Status != oldStatus || user.Role != oldRole || user.RPMLimit != oldRPMLimit || !sameInt64Set(user.AllowedGroups, oldAllowedGroups) { + s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, user.ID) + } + } + + concurrencyDiff := user.Concurrency - oldConcurrency + if concurrencyDiff != 0 { + code, err := GenerateRedeemCode() + if err != nil { + logger.LegacyPrintf("service.admin", "failed to generate adjustment redeem code: %v", err) + return user, nil + } + adjustmentRecord := &RedeemCode{ + Code: code, + Type: AdjustmentTypeAdminConcurrency, + Value: float64(concurrencyDiff), + Status: StatusUsed, + UsedBy: &user.ID, + } + now := time.Now() + adjustmentRecord.UsedAt = &now + if err := s.redeemCodeRepo.Create(ctx, adjustmentRecord); err != nil { + logger.LegacyPrintf("service.admin", "failed to create concurrency adjustment redeem code: %v", err) + } + } + + return user, nil +} + +func sameInt64Set(a, b []int64) bool { + if len(a) != len(b) { + return false + } + if len(a) == 0 { + return true + } + counts := make(map[int64]int, len(a)) + for _, v := range a { + counts[v]++ + } + for _, v := range b { + if counts[v] == 0 { + return false + } + counts[v]-- + } + return true +} + +func (s *adminServiceImpl) DeleteUser(ctx context.Context, id int64) error { + // Protect admin users: cannot delete admin accounts + user, err := s.userRepo.GetByID(ctx, id) + if err != nil { + return err + } + if user.Role == "admin" { + return errors.New("cannot delete admin user") + } + + apiKeys, err := s.listUserAPIKeysForDeletion(ctx, id) + if err != nil { + return err + } + + if s.entClient != nil { + tx, err := s.entClient.Tx(ctx) + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + + opCtx := dbent.NewTxContext(ctx, tx) + if err := s.deleteUserWithAPIKeys(opCtx, id, apiKeys); err != nil { + return err + } + if err := tx.Commit(); err != nil { + return err + } + } else { + if err := s.deleteUserWithAPIKeys(ctx, id, apiKeys); err != nil { + return err + } + } + + if s.authCacheInvalidator != nil { + for _, key := range apiKeys { + if keyValue := strings.TrimSpace(key.Key); keyValue != "" { + s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, keyValue) + } + } + s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, id) + } + return nil +} + +func (s *adminServiceImpl) listUserAPIKeysForDeletion(ctx context.Context, userID int64) ([]APIKey, error) { + if s.apiKeyRepo == nil { + return nil, nil + } + + const pageSize = 1000 + keys := make([]APIKey, 0) + for page := 1; ; page++ { + batch, result, err := s.apiKeyRepo.ListByUserID(ctx, userID, pagination.PaginationParams{ + Page: page, + PageSize: pageSize, + SortBy: "id", + SortOrder: pagination.SortOrderAsc, + }, APIKeyListFilters{}) + if err != nil { + return nil, fmt.Errorf("list user api keys: %w", err) + } + keys = append(keys, batch...) + if len(batch) == 0 || len(batch) < pageSize || result == nil || int64(len(keys)) >= result.Total { + break + } + } + return keys, nil +} + +func (s *adminServiceImpl) deleteUserWithAPIKeys(ctx context.Context, userID int64, apiKeys []APIKey) error { + if s.apiKeyRepo != nil { + for _, key := range apiKeys { + if key.ID <= 0 { + continue + } + if err := s.apiKeyRepo.DeleteWithAudit(ctx, key.ID); err != nil { + logger.LegacyPrintf("service.admin", "delete user api key failed: user_id=%d api_key_id=%d err=%v", userID, key.ID, err) + return fmt.Errorf("delete user api key %d: %w", key.ID, err) + } + } + } + + if err := s.userRepo.Delete(ctx, userID); err != nil { + logger.LegacyPrintf("service.admin", "delete user failed: user_id=%d err=%v", userID, err) + return err + } + return nil +} + +func (s *adminServiceImpl) BatchUpdateConcurrency(ctx context.Context, userIDs []int64, value int, mode string) (int, error) { + cleaned := make([]int64, 0, len(userIDs)) + for _, uid := range userIDs { + if uid > 0 { + cleaned = append(cleaned, uid) + } + } + if len(cleaned) == 0 { + return 0, nil + } + + var affected int + var err error + switch mode { + case "set": + affected, err = s.userRepo.BatchSetConcurrency(ctx, cleaned, value) + case "add": + affected, err = s.userRepo.BatchAddConcurrency(ctx, cleaned, value) + default: + return 0, errors.New("invalid mode: must be 'set' or 'add'") + } + if err != nil { + return 0, err + } + + if s.authCacheInvalidator != nil { + for _, uid := range cleaned { + s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, uid) + } + } + return affected, nil +} + +func (s *adminServiceImpl) UpdateUserBalance(ctx context.Context, userID int64, balance float64, operation string, notes string) (*User, error) { + user, err := s.userRepo.GetByID(ctx, userID) + if err != nil { + return nil, err + } + + oldBalance := user.Balance + + switch operation { + case "set": + user.Balance = balance + case "add": + user.Balance += balance + case "subtract": + user.Balance -= balance + } + + if user.Balance < 0 { + return nil, fmt.Errorf("balance cannot be negative, current balance: %.2f, requested operation would result in: %.2f", oldBalance, user.Balance) + } + + if err := s.userRepo.Update(ctx, user); err != nil { + return nil, err + } + balanceDiff := user.Balance - oldBalance + if s.authCacheInvalidator != nil && balanceDiff != 0 { + s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID) + } + + if s.billingCacheService != nil { + go func() { + cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := s.billingCacheService.InvalidateUserBalance(cacheCtx, userID); err != nil { + logger.LegacyPrintf("service.admin", "invalidate user balance cache failed: user_id=%d err=%v", userID, err) + } + }() + } + + if balanceDiff != 0 { + code, err := GenerateRedeemCode() + if err != nil { + logger.LegacyPrintf("service.admin", "failed to generate adjustment redeem code: %v", err) + return user, nil + } + + adjustmentRecord := &RedeemCode{ + Code: code, + Type: AdjustmentTypeAdminBalance, + Value: balanceDiff, + Status: StatusUsed, + UsedBy: &user.ID, + Notes: notes, + } + now := time.Now() + adjustmentRecord.UsedAt = &now + + if err := s.redeemCodeRepo.Create(ctx, adjustmentRecord); err != nil { + logger.LegacyPrintf("service.admin", "failed to create balance adjustment redeem code: %v", err) + } + } + + return user, nil +} + +func (s *adminServiceImpl) GetUserAPIKeys(ctx context.Context, userID int64, page, pageSize int, sortBy, sortOrder string) ([]APIKey, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + keys, result, err := s.apiKeyRepo.ListByUserID(ctx, userID, params, APIKeyListFilters{}) + if err != nil { + return nil, 0, err + } + return keys, result.Total, nil +} + +func (s *adminServiceImpl) GetUserRPMStatus(ctx context.Context, userID int64) (*UserRPMStatus, error) { + if s.userRPMCache == nil { + return nil, ErrRPMStatusUnavailable + } + + user, err := s.userRepo.GetByID(ctx, userID) + if err != nil { + return nil, err + } + + userRPMUsed, err := s.userRPMCache.GetUserRPM(ctx, userID) + if err != nil { + logger.LegacyPrintf("service.admin", "failed to get user rpm: user_id=%d err=%v", userID, err) + } + + keys, _, err := s.GetUserAPIKeys(ctx, userID, 1, 1000, "", "") + if err != nil { + return nil, err + } + + groupIDSet := make(map[int64]struct{}) + for _, key := range keys { + if key.GroupID != nil && *key.GroupID > 0 { + groupIDSet[*key.GroupID] = struct{}{} + } + } + + groupIDs := make([]int64, 0, len(groupIDSet)) + for groupID := range groupIDSet { + groupIDs = append(groupIDs, groupID) + } + sort.Slice(groupIDs, func(i, j int) bool { return groupIDs[i] < groupIDs[j] }) + + var perGroup []UserGroupRPMStatus + for _, groupID := range groupIDs { + used, getErr := s.userRPMCache.GetUserGroupRPM(ctx, userID, groupID) + if getErr != nil { + logger.LegacyPrintf("service.admin", "failed to get user group rpm: user_id=%d group_id=%d err=%v", userID, groupID, getErr) + } + + entry := UserGroupRPMStatus{ + GroupID: groupID, + Used: used, + } + + if s.groupRepo != nil { + if group, groupErr := s.groupRepo.GetByIDLite(ctx, groupID); groupErr == nil && group != nil { + entry.GroupName = group.Name + entry.Limit = group.RPMLimit + entry.Source = "group" + } else if groupErr != nil { + logger.LegacyPrintf("service.admin", "failed to get group rpm status metadata: group_id=%d err=%v", groupID, groupErr) + } + } + + if s.userGroupRateRepo != nil { + override, overrideErr := s.userGroupRateRepo.GetRPMOverrideByUserAndGroup(ctx, userID, groupID) + if overrideErr != nil { + logger.LegacyPrintf("service.admin", "failed to get rpm override: user_id=%d group_id=%d err=%v", userID, groupID, overrideErr) + } else if override != nil { + entry.Limit = *override + entry.Source = "override" + } + } + + perGroup = append(perGroup, entry) + } + + return &UserRPMStatus{ + UserRPMUsed: userRPMUsed, + UserRPMLimit: user.RPMLimit, + PerGroup: perGroup, + }, nil +} + +func (s *adminServiceImpl) GetUserUsageStats(ctx context.Context, userID int64, period string) (any, error) { + // Return mock data for now + return map[string]any{ + "period": period, + "total_requests": 0, + "total_cost": 0.0, + "total_tokens": 0, + "avg_duration_ms": 0, + }, nil +} + +// GetUserBalanceHistory returns paginated balance/concurrency change records for a user. +func (s *adminServiceImpl) GetUserBalanceHistory(ctx context.Context, userID int64, page, pageSize int, codeType string) ([]RedeemCode, int64, float64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize} + if codeType == RedeemTypeAffiliateBalance { + codes, total, err := s.listAffiliateBalanceHistory(ctx, userID, params) + if err != nil { + return nil, 0, 0, err + } + totalRecharged, err := s.redeemCodeRepo.SumPositiveBalanceByUser(ctx, userID) + if err != nil { + return nil, 0, 0, err + } + return codes, total, totalRecharged, nil + } + + if codeType == "" { + return s.getAllUserBalanceHistory(ctx, userID, params) + } + + codes, result, err := s.redeemCodeRepo.ListByUserPaginated(ctx, userID, params, codeType) + if err != nil { + return nil, 0, 0, err + } + total := result.Total + // Aggregate total recharged amount (only once, regardless of type filter) + totalRecharged, err := s.redeemCodeRepo.SumPositiveBalanceByUser(ctx, userID) + if err != nil { + return nil, 0, 0, err + } + return codes, total, totalRecharged, nil +} + +func (s *adminServiceImpl) getAllUserBalanceHistory(ctx context.Context, userID int64, params pagination.PaginationParams) ([]RedeemCode, int64, float64, error) { + needed := params.Offset() + params.Limit() + if needed < params.Limit() { + needed = params.Limit() + } + + redeemCodes, redeemTotal, err := s.listRedeemBalanceHistoryForMerge(ctx, userID, needed) + if err != nil { + return nil, 0, 0, err + } + affiliateCodes, affiliateTotal, err := s.listAffiliateBalanceHistoryForMerge(ctx, userID, needed) + if err != nil { + return nil, 0, 0, err + } + codes := mergeBalanceHistoryCodes(redeemCodes, affiliateCodes, params) + + totalRecharged, err := s.redeemCodeRepo.SumPositiveBalanceByUser(ctx, userID) + if err != nil { + return nil, 0, 0, err + } + return codes, redeemTotal + affiliateTotal, totalRecharged, nil +} + +func (s *adminServiceImpl) listRedeemBalanceHistoryForMerge(ctx context.Context, userID int64, needed int) ([]RedeemCode, int64, error) { + if needed <= 0 { + return nil, 0, nil + } + + var ( + out []RedeemCode + total int64 + ) + for page := 1; len(out) < needed; page++ { + params := pagination.PaginationParams{Page: page, PageSize: 1000} + codes, result, err := s.redeemCodeRepo.ListByUserPaginated(ctx, userID, params, "") + if err != nil { + return nil, 0, err + } + if result != nil { + total = result.Total + } + out = append(out, codes...) + if len(codes) < params.Limit() || int64(len(out)) >= total { + break + } + } + if len(out) > needed { + out = out[:needed] + } + return out, total, nil +} + +func (s *adminServiceImpl) listAffiliateBalanceHistoryForMerge(ctx context.Context, userID int64, needed int) ([]RedeemCode, int64, error) { + if needed <= 0 { + return nil, 0, nil + } + + var ( + out []RedeemCode + total int64 + ) + for page := 1; len(out) < needed; page++ { + params := pagination.PaginationParams{Page: page, PageSize: 1000} + codes, currentTotal, err := s.listAffiliateBalanceHistory(ctx, userID, params) + if err != nil { + return nil, 0, err + } + total = currentTotal + out = append(out, codes...) + if len(codes) < params.Limit() || int64(len(out)) >= total { + break + } + } + if len(out) > needed { + out = out[:needed] + } + return out, total, nil +} + +func (s *adminServiceImpl) listAffiliateBalanceHistory(ctx context.Context, userID int64, params pagination.PaginationParams) ([]RedeemCode, int64, error) { + if s == nil || s.entClient == nil || userID <= 0 { + return nil, 0, nil + } + + rows, err := s.entClient.QueryContext(ctx, ` +SELECT id, + amount::double precision, + created_at +FROM user_affiliate_ledger +WHERE user_id = $1 + AND action = 'transfer' +ORDER BY created_at DESC, id DESC +OFFSET $2 +LIMIT $3`, userID, params.Offset(), params.Limit()) + if err != nil { + return nil, 0, err + } + defer func() { _ = rows.Close() }() + + codes := make([]RedeemCode, 0, params.Limit()) + for rows.Next() { + var id int64 + var amount float64 + var createdAt time.Time + if err := rows.Scan(&id, &amount, &createdAt); err != nil { + return nil, 0, err + } + usedBy := userID + usedAt := createdAt + codes = append(codes, RedeemCode{ + ID: -id, + Code: fmt.Sprintf("AFF-%d", id), + Type: RedeemTypeAffiliateBalance, + Value: amount, + Status: StatusUsed, + UsedBy: &usedBy, + UsedAt: &usedAt, + CreatedAt: createdAt, + }) + } + if err := rows.Err(); err != nil { + return nil, 0, err + } + + total, err := countAffiliateBalanceHistory(ctx, s.entClient, userID) + if err != nil { + return nil, 0, err + } + return codes, total, nil +} + +func countAffiliateBalanceHistory(ctx context.Context, client *dbent.Client, userID int64) (int64, error) { + rows, err := client.QueryContext(ctx, ` +SELECT COUNT(*) +FROM user_affiliate_ledger +WHERE user_id = $1 + AND action = 'transfer'`, userID) + if err != nil { + return 0, err + } + defer func() { _ = rows.Close() }() + + var total sql.NullInt64 + if rows.Next() { + if err := rows.Scan(&total); err != nil { + return 0, err + } + } + if err := rows.Err(); err != nil { + return 0, err + } + if !total.Valid { + return 0, nil + } + return total.Int64, nil +} + +func mergeBalanceHistoryCodes(redeemCodes, affiliateCodes []RedeemCode, params pagination.PaginationParams) []RedeemCode { + combined := append(append([]RedeemCode{}, redeemCodes...), affiliateCodes...) + sort.SliceStable(combined, func(i, j int) bool { + return redeemCodeHistoryTime(combined[i]).After(redeemCodeHistoryTime(combined[j])) + }) + offset := params.Offset() + if offset >= len(combined) { + return []RedeemCode{} + } + end := offset + params.Limit() + if end > len(combined) { + end = len(combined) + } + return combined[offset:end] +} + +func redeemCodeHistoryTime(code RedeemCode) time.Time { + if code.UsedAt != nil { + return *code.UsedAt + } + return code.CreatedAt +} + +func (s *adminServiceImpl) BindUserAuthIdentity(ctx context.Context, userID int64, input AdminBindAuthIdentityInput) (*AdminBoundAuthIdentity, error) { + if userID <= 0 { + return nil, infraerrors.BadRequest("INVALID_INPUT", "user_id must be greater than 0") + } + if s == nil || s.entClient == nil || s.userRepo == nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_UNAVAILABLE", "auth identity binding service is unavailable") + } + if _, err := s.userRepo.GetByID(ctx, userID); err != nil { + return nil, err + } + + providerType := normalizeAdminAuthIdentityProviderType(input.ProviderType) + providerKey := strings.TrimSpace(input.ProviderKey) + providerSubject := strings.TrimSpace(input.ProviderSubject) + if providerType == "" { + return nil, infraerrors.BadRequest("INVALID_INPUT", "provider_type must be one of email, linuxdo, oidc, wechat, or dingtalk") + } + if providerKey == "" || providerSubject == "" { + return nil, infraerrors.BadRequest("INVALID_INPUT", "provider_type, provider_key, and provider_subject are required") + } + canonicalProviderKey := canonicalAdminAuthIdentityProviderKey(providerType, "", providerKey) + compatibleProviderKeys := compatibleAdminAuthIdentityProviderKeys(providerType, providerKey) + + var issuer *string + if input.Issuer != nil { + trimmed := strings.TrimSpace(*input.Issuer) + if trimmed != "" { + issuer = &trimmed + } + } + + channelInput := normalizeAdminBindChannelInput(input.Channel) + if input.Channel != nil && channelInput == nil { + return nil, infraerrors.BadRequest("INVALID_INPUT", "channel, channel_app_id, and channel_subject are required when channel binding is provided") + } + + verifiedAt := time.Now().UTC() + tx, err := s.entClient.Tx(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_TX_FAILED", "failed to start auth identity bind transaction").WithCause(err) + } + defer func() { _ = tx.Rollback() }() + + identityRecords, err := tx.AuthIdentity.Query(). + Where( + authidentity.ProviderTypeEQ(providerType), + authidentity.ProviderKeyIn(compatibleProviderKeys...), + authidentity.ProviderSubjectEQ(providerSubject), + ). + All(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_LOOKUP_FAILED", "failed to inspect auth identity ownership").WithCause(err) + } + if hasAdminAuthIdentityOwnershipConflict(identityRecords, userID) { + return nil, infraerrors.Conflict("AUTH_IDENTITY_OWNERSHIP_CONFLICT", "auth identity already belongs to another user") + } + identity := selectOwnedAdminAuthIdentity(identityRecords, userID) + + if identity == nil { + create := tx.AuthIdentity.Create(). + SetUserID(userID). + SetProviderType(providerType). + SetProviderKey(canonicalProviderKey). + SetProviderSubject(providerSubject). + SetVerifiedAt(verifiedAt) + if issuer != nil { + create = create.SetIssuer(*issuer) + } + if input.Metadata != nil { + create = create.SetMetadata(cloneAdminAuthIdentityMetadata(input.Metadata)) + } + identity, err = create.Save(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_SAVE_FAILED", "failed to save auth identity").WithCause(err) + } + } else { + update := tx.AuthIdentity.UpdateOneID(identity.ID). + SetVerifiedAt(verifiedAt). + SetProviderKey(canonicalProviderKey) + if issuer != nil { + update = update.SetIssuer(*issuer) + } + if input.Metadata != nil { + update = update.SetMetadata(cloneAdminAuthIdentityMetadata(input.Metadata)) + } + identity, err = update.Save(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_SAVE_FAILED", "failed to save auth identity").WithCause(err) + } + } + + var channel *dbent.AuthIdentityChannel + if channelInput != nil { + channelRecords, err := tx.AuthIdentityChannel.Query(). + Where( + authidentitychannel.ProviderTypeEQ(providerType), + authidentitychannel.ProviderKeyIn(compatibleProviderKeys...), + authidentitychannel.ChannelEQ(channelInput.Channel), + authidentitychannel.ChannelAppIDEQ(channelInput.ChannelAppID), + authidentitychannel.ChannelSubjectEQ(channelInput.ChannelSubject), + ). + WithIdentity(). + All(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_CHANNEL_LOOKUP_FAILED", "failed to inspect auth identity channel ownership").WithCause(err) + } + if hasAdminAuthIdentityChannelOwnershipConflict(channelRecords, userID) { + return nil, infraerrors.Conflict("AUTH_IDENTITY_CHANNEL_OWNERSHIP_CONFLICT", "auth identity channel already belongs to another user") + } + channel = selectOwnedAdminAuthIdentityChannel(channelRecords, userID) + if channel == nil { + create := tx.AuthIdentityChannel.Create(). + SetIdentityID(identity.ID). + SetProviderType(providerType). + SetProviderKey(canonicalProviderKey). + SetChannel(channelInput.Channel). + SetChannelAppID(channelInput.ChannelAppID). + SetChannelSubject(channelInput.ChannelSubject) + if channelInput.Metadata != nil { + create = create.SetMetadata(cloneAdminAuthIdentityMetadata(channelInput.Metadata)) + } + channel, err = create.Save(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_CHANNEL_SAVE_FAILED", "failed to save auth identity channel").WithCause(err) + } + } else { + update := tx.AuthIdentityChannel.UpdateOneID(channel.ID). + SetIdentityID(identity.ID). + SetProviderKey(canonicalProviderKey) + if channelInput.Metadata != nil { + update = update.SetMetadata(cloneAdminAuthIdentityMetadata(channelInput.Metadata)) + } + channel, err = update.Save(ctx) + if err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_CHANNEL_SAVE_FAILED", "failed to save auth identity channel").WithCause(err) + } + } + } + + if err := tx.Commit(); err != nil { + return nil, infraerrors.InternalServer("ADMIN_AUTH_IDENTITY_BIND_COMMIT_FAILED", "failed to commit auth identity bind").WithCause(err) + } + return buildAdminBoundAuthIdentity(identity, channel), nil +} + +func compatibleAdminAuthIdentityProviderKeys(providerType, providerKey string) []string { + providerType = strings.TrimSpace(strings.ToLower(providerType)) + providerKey = strings.TrimSpace(providerKey) + if providerKey == "" { + return []string{providerKey} + } + if providerType != "wechat" { + return []string{providerKey} + } + + keys := []string{providerKey} + if !strings.EqualFold(providerKey, "wechat-main") { + keys = append(keys, "wechat-main") + } + if !strings.EqualFold(providerKey, "wechat") { + keys = append(keys, "wechat") + } + return keys +} + +func canonicalAdminAuthIdentityProviderKey(providerType, existingKey, requestedKey string) string { + providerType = strings.TrimSpace(strings.ToLower(providerType)) + existingKey = strings.TrimSpace(existingKey) + requestedKey = strings.TrimSpace(requestedKey) + if providerType != "wechat" { + if requestedKey != "" { + return requestedKey + } + return existingKey + } + if strings.EqualFold(existingKey, "wechat") || strings.EqualFold(existingKey, "wechat-main") || strings.EqualFold(requestedKey, "wechat-main") { + return "wechat-main" + } + if requestedKey != "" { + return requestedKey + } + return existingKey +} + +func adminAuthIdentityProviderKeyRank(providerType, providerKey string) int { + providerType = strings.TrimSpace(strings.ToLower(providerType)) + providerKey = strings.TrimSpace(providerKey) + if providerType != "wechat" { + return 0 + } + switch { + case strings.EqualFold(providerKey, "wechat-main"): + return 0 + case strings.EqualFold(providerKey, "wechat"): + return 2 + default: + return 1 + } +} + +func selectOwnedAdminAuthIdentity(records []*dbent.AuthIdentity, userID int64) *dbent.AuthIdentity { + var selected *dbent.AuthIdentity + for _, record := range records { + if record.UserID != userID { + continue + } + if selected == nil || adminAuthIdentityProviderKeyRank(record.ProviderType, record.ProviderKey) < adminAuthIdentityProviderKeyRank(selected.ProviderType, selected.ProviderKey) { + selected = record + } + } + return selected +} + +func hasAdminAuthIdentityOwnershipConflict(records []*dbent.AuthIdentity, userID int64) bool { + for _, record := range records { + if record.UserID != userID { + return true + } + } + return false +} + +func selectOwnedAdminAuthIdentityChannel(records []*dbent.AuthIdentityChannel, userID int64) *dbent.AuthIdentityChannel { + var selected *dbent.AuthIdentityChannel + for _, record := range records { + if record.Edges.Identity == nil || record.Edges.Identity.UserID != userID { + continue + } + if selected == nil || adminAuthIdentityProviderKeyRank(record.ProviderType, record.ProviderKey) < adminAuthIdentityProviderKeyRank(selected.ProviderType, selected.ProviderKey) { + selected = record + } + } + return selected +} + +func hasAdminAuthIdentityChannelOwnershipConflict(records []*dbent.AuthIdentityChannel, userID int64) bool { + for _, record := range records { + if record.Edges.Identity != nil && record.Edges.Identity.UserID != userID { + return true + } + } + return false +} + +func normalizeAdminBindChannelInput(input *AdminBindAuthIdentityChannelInput) *AdminBindAuthIdentityChannelInput { + if input == nil { + return nil + } + channel := &AdminBindAuthIdentityChannelInput{ + Channel: strings.TrimSpace(input.Channel), + ChannelAppID: strings.TrimSpace(input.ChannelAppID), + ChannelSubject: strings.TrimSpace(input.ChannelSubject), + Metadata: cloneAdminAuthIdentityMetadata(input.Metadata), + } + if channel.Channel == "" || channel.ChannelAppID == "" || channel.ChannelSubject == "" { + return nil + } + return channel +} + +func normalizeAdminAuthIdentityProviderType(input string) string { + switch strings.ToLower(strings.TrimSpace(input)) { + case "email": + return "email" + case "linuxdo": + return "linuxdo" + case "oidc": + return "oidc" + case "wechat": + return "wechat" + case "dingtalk": + return "dingtalk" + default: + return "" + } +} + +func buildAdminBoundAuthIdentity(identity *dbent.AuthIdentity, channel *dbent.AuthIdentityChannel) *AdminBoundAuthIdentity { + if identity == nil { + return nil + } + result := &AdminBoundAuthIdentity{ + UserID: identity.UserID, + ProviderType: strings.TrimSpace(identity.ProviderType), + ProviderKey: strings.TrimSpace(identity.ProviderKey), + ProviderSubject: strings.TrimSpace(identity.ProviderSubject), + VerifiedAt: identity.VerifiedAt, + Issuer: identity.Issuer, + Metadata: cloneAdminAuthIdentityMetadata(identity.Metadata), + CreatedAt: identity.CreatedAt, + UpdatedAt: identity.UpdatedAt, + } + if channel != nil { + result.Channel = &AdminBoundAuthIdentityChannel{ + Channel: strings.TrimSpace(channel.Channel), + ChannelAppID: strings.TrimSpace(channel.ChannelAppID), + ChannelSubject: strings.TrimSpace(channel.ChannelSubject), + Metadata: cloneAdminAuthIdentityMetadata(channel.Metadata), + CreatedAt: channel.CreatedAt, + UpdatedAt: channel.UpdatedAt, + } + } + return result +} + +func cloneAdminAuthIdentityMetadata(input map[string]any) map[string]any { + if input == nil { + return nil + } + if len(input) == 0 { + return map[string]any{} + } + data, err := json.Marshal(input) + if err != nil { + out := make(map[string]any, len(input)) + for key, value := range input { + out[key] = value + } + return out + } + var out map[string]any + if err := json.Unmarshal(data, &out); err != nil { + out = make(map[string]any, len(input)) + for key, value := range input { + out[key] = value + } + } + return out +} + +// Redeem code management implementations +func (s *adminServiceImpl) ListRedeemCodes(ctx context.Context, page, pageSize int, codeType, status, search string, sortBy, sortOrder string) ([]RedeemCode, int64, error) { + params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder} + codes, result, err := s.redeemCodeRepo.ListWithFilters(ctx, params, codeType, status, search) + if err != nil { + return nil, 0, err + } + return codes, result.Total, nil +} + +func (s *adminServiceImpl) GetRedeemCode(ctx context.Context, id int64) (*RedeemCode, error) { + return s.redeemCodeRepo.GetByID(ctx, id) +} + +func (s *adminServiceImpl) GenerateRedeemCodes(ctx context.Context, input *GenerateRedeemCodesInput) ([]RedeemCode, error) { + if input.ExpiresAt != nil && !input.ExpiresAt.After(time.Now()) { + return nil, ErrRedeemCodeExpired + } + + // 如果是订阅类型,验证必须有 GroupID + if input.Type == RedeemTypeSubscription { + if input.GroupID == nil { + return nil, errors.New("group_id is required for subscription type") + } + // 验证分组存在且为订阅类型 + group, err := s.groupRepo.GetByID(ctx, *input.GroupID) + if err != nil { + return nil, fmt.Errorf("group not found: %w", err) + } + if !group.IsSubscriptionType() { + return nil, errors.New("group must be subscription type") + } + } + + codes := make([]RedeemCode, 0, input.Count) + for i := 0; i < input.Count; i++ { + codeValue, err := GenerateRedeemCode() + if err != nil { + return nil, err + } + code := RedeemCode{ + Code: codeValue, + Type: input.Type, + Value: input.Value, + Status: StatusUnused, + ExpiresAt: input.ExpiresAt, + } + // 订阅类型专用字段 + if input.Type == RedeemTypeSubscription { + code.GroupID = input.GroupID + code.ValidityDays = input.ValidityDays + if code.ValidityDays <= 0 { + code.ValidityDays = 30 // 默认30天 + } + } + if err := s.redeemCodeRepo.Create(ctx, &code); err != nil { + return nil, err + } + codes = append(codes, code) + } + return codes, nil +} + +func (s *adminServiceImpl) DeleteRedeemCode(ctx context.Context, id int64) error { + return s.redeemCodeRepo.Delete(ctx, id) +} + +func (s *adminServiceImpl) BatchDeleteRedeemCodes(ctx context.Context, ids []int64) (int64, error) { + var deleted int64 + for _, id := range ids { + if err := s.redeemCodeRepo.Delete(ctx, id); err == nil { + deleted++ + } + } + return deleted, nil +} + +func (s *adminServiceImpl) ExpireRedeemCode(ctx context.Context, id int64) (*RedeemCode, error) { + code, err := s.redeemCodeRepo.GetByID(ctx, id) + if err != nil { + return nil, err + } + code.Status = StatusExpired + if err := s.redeemCodeRepo.Update(ctx, code); err != nil { + return nil, err + } + return code, nil +} diff --git a/backend/internal/service/antigravity_gateway_claude.go b/backend/internal/service/antigravity_gateway_claude.go new file mode 100644 index 0000000000..2011625984 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_claude.go @@ -0,0 +1,754 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// Forward 转发 Claude 协议请求(Claude → Gemini 转换) +// +// 限流处理流程: +// +// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 +// ├─ 成功 → 正常返回 +// └─ 429/503 → handleSmartRetry +// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 +// └─ retryDelay < 7s → 等待后重试 1 次 +// ├─ 成功 → 正常返回 +// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 +func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte, isStickySession bool) (*ForwardResult, error) { + // 上游透传账号直接转发,不走 OAuth token 刷新 + if account.Type == AccountTypeUpstream { + return s.ForwardUpstream(ctx, c, account, body) + } + + startTime := time.Now() + + sessionID := getSessionID(c) + prefix := logPrefix(sessionID, account.Name) + + // 解析 Claude 请求 + var claudeReq antigravity.ClaudeRequest + if err := json.Unmarshal(body, &claudeReq); err != nil { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") + } + if strings.TrimSpace(claudeReq.Model) == "" { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Missing model") + } + + originalModel := claudeReq.Model + mappedModel := s.getMappedModel(account, claudeReq.Model) + if mappedModel == "" { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + return nil, s.writeClaudeError(c, http.StatusForbidden, "permission_error", fmt.Sprintf("model %s not in whitelist", claudeReq.Model)) + } + // 应用 thinking 模式自动后缀:如果 thinking 开启且目标是 claude-sonnet-4-5,自动改为 thinking 版本 + thinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive") + mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled) + billingModel := mappedModel + + // 获取 access_token + if s.tokenProvider == nil { + return nil, s.writeClaudeError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`), + } + } + + projectID, err := resolveAntigravityProjectID(account) + if err != nil { + _ = s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) + return nil, err + } + + // 代理 URL + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + // 获取转换选项 + // Antigravity 上游要求必须包含身份提示词,否则会返回 429 + transformOpts := s.getClaudeTransformOptions(ctx) + transformOpts.EnableIdentityPatch = true // 强制启用,Antigravity 上游必需 + + // 转换 Claude 请求为 Gemini 格式 + geminiBody, err := antigravity.TransformClaudeToGeminiWithOptions(&claudeReq, projectID, mappedModel, transformOpts) + if err != nil { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request") + } + + // Antigravity 上游只支持流式请求,统一使用 streamGenerateContent + // 如果客户端请求非流式,在响应处理阶段会收集完整流式响应后转换返回 + action := "streamGenerateContent" + + // 执行带重试的请求 + result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: action, + body: geminiBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, // Forward 由上层判断粘性会话 + groupID: 0, // Forward 方法没有 groupID,由上层处理粘性会话清除 + sessionHash: "", // Forward 方法没有 sessionHash,由上层处理粘性会话清除 + }) + if err != nil { + // 检查是否是账号切换信号,转换为 UpstreamFailoverError 让 Handler 切换账号 + if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusServiceUnavailable, + ForceCacheBilling: switchErr.IsStickySession, + } + } + // 区分客户端取消和真正的上游失败,返回更准确的错误消息 + if c.Request.Context().Err() != nil { + return nil, s.writeClaudeError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response") + } + return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries") + } + resp := result.resp + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + + // 优先检测 thinking block 的 signature 相关错误(400)并重试一次: + // Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验, + // 当历史消息携带的 signature 不合法时会直接 400;去除 thinking 后可继续完成请求。 + if resp.StatusCode == http.StatusBadRequest && isSignatureRelatedError(respBody) && s.settingService.IsSignatureRectifierEnabled(ctx) { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + logBody, maxBytes := s.getLogConfig() + upstreamDetail := s.getUpstreamErrorDetail(respBody) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "signature_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + // Conservative two-stage fallback: + // 1) Disable top-level thinking + thinking->text + // 2) Only if still signature-related 400: also downgrade tool_use/tool_result to text. + + retryStages := []struct { + name string + strip func(*antigravity.ClaudeRequest) (bool, error) + }{ + {name: "thinking-only", strip: stripThinkingFromClaudeRequest}, + {name: "thinking+tools", strip: stripSignatureSensitiveBlocksFromClaudeRequest}, + } + + for _, stage := range retryStages { + retryClaudeReq := claudeReq + retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) + + stripped, stripErr := stage.strip(&retryClaudeReq) + if stripErr != nil || !stripped { + continue + } + + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected signature-related 400, retrying once (%s)", account.ID, stage.name) + + retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, s.getClaudeTransformOptions(ctx)) + if txErr != nil { + continue + } + retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: action, + body: retryGeminiBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, + groupID: 0, // Forward 方法没有 groupID,由上层处理粘性会话清除 + sessionHash: "", // Forward 方法没有 sessionHash,由上层处理粘性会话清除 + }) + if retryErr != nil { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + Kind: "signature_retry_request_error", + Message: sanitizeUpstreamErrorMessage(retryErr.Error()), + }) + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: signature retry request failed (%s): %v", account.ID, stage.name, retryErr) + continue + } + + retryResp := retryResult.resp + if retryResp.StatusCode < 400 { + _ = resp.Body.Close() + resp = retryResp + respBody = nil + break + } + + retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) + _ = retryResp.Body.Close() + if retryResp.StatusCode == http.StatusTooManyRequests { + retryBaseURL := "" + if retryResp.Request != nil && retryResp.Request.URL != nil { + retryBaseURL = retryResp.Request.URL.Scheme + "://" + retryResp.Request.URL.Host + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited base_url=%s retry_stage=%s body=%s", prefix, retryBaseURL, stage.name, truncateForLog(retryBody, 200)) + } + kind := "signature_retry" + if strings.TrimSpace(stage.name) != "" { + kind = "signature_retry_" + strings.ReplaceAll(stage.name, "+", "_") + } + retryUpstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(retryBody)) + retryUpstreamMsg = sanitizeUpstreamErrorMessage(retryUpstreamMsg) + retryUpstreamDetail := "" + if logBody { + retryUpstreamDetail = truncateString(string(retryBody), maxBytes) + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: retryResp.StatusCode, + UpstreamRequestID: retryResp.Header.Get("x-request-id"), + Kind: kind, + Message: retryUpstreamMsg, + Detail: retryUpstreamDetail, + }) + + // If this stage fixed the signature issue, we stop; otherwise we may try the next stage. + if retryResp.StatusCode != http.StatusBadRequest || !isSignatureRelatedError(retryBody) { + respBody = retryBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + } + break + } + + // Still signature-related; capture context and allow next stage. + respBody = retryBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + } + } + } + + // Budget 整流:检测 budget_tokens 约束错误并自动修正重试 + if resp.StatusCode == http.StatusBadRequest && respBody != nil && !isSignatureRelatedError(respBody) { + errMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + if isThinkingBudgetConstraintError(errMsg) && s.settingService.IsBudgetRectifierEnabled(ctx) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "budget_constraint_error", + Message: errMsg, + Detail: s.getUpstreamErrorDetail(respBody), + }) + + // 修正 claudeReq 的 thinking 参数(adaptive 模式不修正) + if claudeReq.Thinking == nil || claudeReq.Thinking.Type != "adaptive" { + retryClaudeReq := claudeReq + retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) + // 创建新的 ThinkingConfig 避免修改原始 claudeReq.Thinking 指针 + retryClaudeReq.Thinking = &antigravity.ThinkingConfig{ + Type: "enabled", + BudgetTokens: BudgetRectifyBudgetTokens, + } + if retryClaudeReq.MaxTokens < BudgetRectifyMinMaxTokens { + retryClaudeReq.MaxTokens = BudgetRectifyMaxTokens + } + + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) + + retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, transformOpts) + if txErr == nil { + retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: action, + body: retryGeminiBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, + groupID: 0, + sessionHash: "", + }) + if retryErr == nil { + retryResp := retryResult.resp + if retryResp.StatusCode < 400 { + _ = resp.Body.Close() + resp = retryResp + respBody = nil + } else { + retryBody := s.readUpstreamErrorBody(retryResp) + _ = retryResp.Body.Close() + respBody = retryBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + } + } + } else { + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: budget rectifier retry failed: %v", account.ID, retryErr) + } + } + } + } + } + + // 处理错误响应(重试后仍失败或不触发重试) + if resp.StatusCode >= 400 { + // 检测 prompt too long 错误,返回特殊错误类型供上层 fallback + if resp.StatusCode == http.StatusBadRequest && isPromptTooLongError(respBody) { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.getUpstreamErrorDetail(respBody) + logBody, maxBytes := s.getLogConfig() + if logBody { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=400 prompt_too_long=true upstream_message=%q request_id=%s body=%s", prefix, upstreamMsg, resp.Header.Get("x-request-id"), truncateForLog(respBody, maxBytes)) + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "prompt_too_long", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &PromptTooLongError{ + StatusCode: resp.StatusCode, + RequestID: resp.Header.Get("x-request-id"), + Body: respBody, + } + } + + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession) + + // 精确匹配服务端配置类 400 错误,触发同账号重试 + failover + if resp.StatusCode == http.StatusBadRequest { + msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + if isGoogleProjectConfigError(msg) { + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + upstreamDetail := s.getUpstreamErrorDetail(respBody) + log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: true} + } + } + + if s.shouldFailoverUpstreamError(resp.StatusCode) { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.getUpstreamErrorDetail(respBody) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} + } + + return nil, s.writeMappedClaudeError(c, account, resp.StatusCode, resp.Header.Get("x-request-id"), respBody) + } + } + + requestID := resp.Header.Get("x-request-id") + if requestID != "" { + c.Header("x-request-id", requestID) + } + + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + if claudeReq.Stream { + // 客户端要求流式,直接透传转换 + streamRes, err := s.handleClaudeStreamingResponse(c, resp, startTime, originalModel) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + } else { + // 客户端要求非流式,收集流式响应后转换返回 + streamRes, err := s.handleClaudeStreamToNonStreaming(c, resp, startTime, originalModel) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + } + + return &ForwardResult{ + RequestID: requestID, + Usage: *usage, + Model: originalModel, + UpstreamModel: billingModel, + Stream: claudeReq.Stream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + }, nil +} + +func isSignatureRelatedError(respBody []byte) bool { + msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + if msg == "" { + // Fallback: best-effort scan of the raw payload. + msg = strings.ToLower(string(respBody)) + } + + // Keep this intentionally broad: different upstreams may use "signature" or "thought_signature". + if strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature") { + return true + } + + // Also detect thinking block structural errors: + // "Expected `thinking` or `redacted_thinking`, but found `text`" + if strings.Contains(msg, "expected") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { + return true + } + + return false +} + +// isPromptTooLongError 检测是否为 prompt too long 错误 +func isPromptTooLongError(respBody []byte) bool { + msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + if msg == "" { + msg = strings.ToLower(string(respBody)) + } + return strings.Contains(msg, "prompt is too long") || + strings.Contains(msg, "request is too long") || + strings.Contains(msg, "context length exceeded") || + strings.Contains(msg, "max_tokens") +} + +// isPassthroughErrorMessage 检查错误消息是否在透传白名单中 +func isPassthroughErrorMessage(msg string) bool { + lower := strings.ToLower(msg) + for _, pattern := range antigravityPassthroughErrorMessages { + if strings.Contains(lower, pattern) { + return true + } + } + return false +} + +// getPassthroughOrDefault 若消息在白名单内则返回原始消息,否则返回默认消息 +func getPassthroughOrDefault(upstreamMsg, defaultMsg string) string { + if isPassthroughErrorMessage(upstreamMsg) { + return upstreamMsg + } + return defaultMsg +} + +func extractAntigravityErrorMessage(body []byte) string { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return "" + } + + // Google-style: {"error": {"message": "..."}} + if errObj, ok := payload["error"].(map[string]any); ok { + if msg, ok := errObj["message"].(string); ok && strings.TrimSpace(msg) != "" { + return msg + } + } + + // Fallback: top-level message + if msg, ok := payload["message"].(string); ok && strings.TrimSpace(msg) != "" { + return msg + } + + return "" +} + +// stripThinkingFromClaudeRequest converts thinking blocks to text blocks in a Claude Messages request. +// This preserves the thinking content while avoiding signature validation errors. +// Note: redacted_thinking blocks are removed because they cannot be converted to text. +// It also disables top-level `thinking` to avoid upstream structural constraints for thinking mode. +func stripThinkingFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { + if req == nil { + return false, nil + } + + changed := false + if req.Thinking != nil { + req.Thinking = nil + changed = true + } + + for i := range req.Messages { + raw := req.Messages[i].Content + if len(raw) == 0 { + continue + } + + // If content is a string, nothing to strip. + var str string + if json.Unmarshal(raw, &str) == nil { + continue + } + + // Otherwise treat as an array of blocks and convert thinking blocks to text. + var blocks []map[string]any + if err := json.Unmarshal(raw, &blocks); err != nil { + continue + } + + filtered := make([]map[string]any, 0, len(blocks)) + modifiedAny := false + for _, block := range blocks { + t, _ := block["type"].(string) + switch t { + case "thinking": + thinkingText, _ := block["thinking"].(string) + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + case "redacted_thinking": + modifiedAny = true + case "": + if thinkingText, hasThinking := block["thinking"].(string); hasThinking { + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + } else { + filtered = append(filtered, block) + } + default: + filtered = append(filtered, block) + } + } + + if !modifiedAny { + continue + } + + if len(filtered) == 0 { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": "(content removed)", + }) + } + + newRaw, err := json.Marshal(filtered) + if err != nil { + return changed, err + } + req.Messages[i].Content = newRaw + changed = true + } + + return changed, nil +} + +// stripSignatureSensitiveBlocksFromClaudeRequest is a stronger retry degradation that additionally converts +// tool blocks to plain text. Use this only after a thinking-only retry still fails with signature errors. +func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { + if req == nil { + return false, nil + } + + changed := false + if req.Thinking != nil { + req.Thinking = nil + changed = true + } + + for i := range req.Messages { + raw := req.Messages[i].Content + if len(raw) == 0 { + continue + } + + // If content is a string, nothing to strip. + var str string + if json.Unmarshal(raw, &str) == nil { + continue + } + + // Otherwise treat as an array of blocks and convert signature-sensitive blocks to text. + var blocks []map[string]any + if err := json.Unmarshal(raw, &blocks); err != nil { + continue + } + + filtered := make([]map[string]any, 0, len(blocks)) + modifiedAny := false + for _, block := range blocks { + t, _ := block["type"].(string) + switch t { + case "thinking": + // Convert thinking to text, skip if empty + thinkingText, _ := block["thinking"].(string) + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + case "redacted_thinking": + // Remove redacted_thinking (cannot convert encrypted content) + modifiedAny = true + case "tool_use": + // Convert tool_use to text to avoid upstream signature/thought_signature validation errors. + // This is a retry-only degradation path, so we prioritise request validity over tool semantics. + name, _ := block["name"].(string) + id, _ := block["id"].(string) + input := block["input"] + inputJSON, _ := json.Marshal(input) + text := "(tool_use)" + if name != "" { + text += " name=" + name + } + if id != "" { + text += " id=" + id + } + if len(inputJSON) > 0 && string(inputJSON) != "null" { + text += " input=" + string(inputJSON) + } + filtered = append(filtered, map[string]any{ + "type": "text", + "text": text, + }) + modifiedAny = true + case "tool_result": + // Convert tool_result to text so it stays consistent when tool_use is downgraded. + toolUseID, _ := block["tool_use_id"].(string) + isError, _ := block["is_error"].(bool) + content := block["content"] + contentJSON, _ := json.Marshal(content) + text := "(tool_result)" + if toolUseID != "" { + text += " tool_use_id=" + toolUseID + } + if isError { + text += " is_error=true" + } + if len(contentJSON) > 0 && string(contentJSON) != "null" { + text += "\n" + string(contentJSON) + } + filtered = append(filtered, map[string]any{ + "type": "text", + "text": text, + }) + modifiedAny = true + case "": + // Handle untyped block with "thinking" field + if thinkingText, hasThinking := block["thinking"].(string); hasThinking { + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + } else { + filtered = append(filtered, block) + } + default: + filtered = append(filtered, block) + } + } + + if !modifiedAny { + continue + } + + if len(filtered) == 0 { + // Keep request valid: upstream rejects empty content arrays. + filtered = append(filtered, map[string]any{ + "type": "text", + "text": "(content removed)", + }) + } + + newRaw, err := json.Marshal(filtered) + if err != nil { + return changed, err + } + req.Messages[i].Content = newRaw + changed = true + } + + return changed, nil +} diff --git a/backend/internal/service/antigravity_gateway_gemini.go b/backend/internal/service/antigravity_gateway_gemini.go new file mode 100644 index 0000000000..f0d20244a8 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_gemini.go @@ -0,0 +1,550 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// ForwardGemini 转发 Gemini 协议请求 +// +// 限流处理流程: +// +// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 +// ├─ 成功 → 正常返回 +// └─ 429/503 → handleSmartRetry +// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 +// └─ retryDelay < 7s → 等待后重试 1 次 +// ├─ 成功 → 正常返回 +// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 +type ForwardGeminiOption func(*forwardGeminiOptions) + +type forwardGeminiOptions struct { + groupID int64 + sessionHash string +} + +func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption { + return func(opts *forwardGeminiOptions) { + opts.groupID = groupID + opts.sessionHash = sessionHash + } +} + +func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) { + startTime := time.Now() + forwardOpts := forwardGeminiOptions{} + for _, apply := range options { + if apply != nil { + apply(&forwardOpts) + } + } + + sessionID := getSessionID(c) + prefix := logPrefix(sessionID, account.Name) + + if strings.TrimSpace(originalModel) == "" { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing model in URL") + } + if strings.TrimSpace(action) == "" { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing action in URL") + } + if len(body) == 0 { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty") + } + + // 解析请求以获取 image_size(用于图片计费) + imageInputSize := s.extractImageInputSize(body) + imageSize := normalizeOpenAIImageSizeTier(imageInputSize) + + switch action { + case "generateContent", "streamGenerateContent": + // ok + case "countTokens": + // 直接返回空值,不透传上游 + c.JSON(http.StatusOK, map[string]any{"totalTokens": 0}) + return &ForwardResult{ + RequestID: "", + Usage: ClaudeUsage{}, + Model: originalModel, + Stream: false, + Duration: time.Since(startTime), + FirstTokenMs: nil, + }, nil + default: + return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action) + } + + mappedModel := s.getMappedModel(account, originalModel) + if mappedModel == "" { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + return nil, s.writeGoogleError(c, http.StatusForbidden, fmt.Sprintf("model %s not in whitelist", originalModel)) + } + billingModel := mappedModel + + // 获取 access_token + if s.tokenProvider == nil { + return nil, s.writeGoogleError(c, http.StatusBadGateway, "Antigravity token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":{"message":"Failed to get upstream access token","status":"UNAVAILABLE"}}`), + } + } + + projectID, err := resolveAntigravityProjectID(account) + if err != nil { + _ = s.writeGoogleError(c, http.StatusBadRequest, err.Error()) + return nil, err + } + + // 代理 URL + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + // Antigravity 上游要求必须包含身份提示词,注入到请求中 + injectedBody, err := injectIdentityPatchToGeminiRequest(body) + if err != nil { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Invalid request body") + } + + // 清理 Schema + if cleanedBody, err := cleanGeminiRequest(injectedBody); err == nil { + injectedBody = cleanedBody + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Cleaned request schema in forwarded request for account %s", account.Name) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Failed to clean schema: %v", err) + } + + // 包装请求 + wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody) + if err != nil { + return nil, s.writeGoogleError(c, http.StatusInternalServerError, "Failed to build upstream request") + } + + // Antigravity 上游只支持流式请求,统一使用 streamGenerateContent + // 如果客户端请求非流式,在响应处理阶段会收集完整流式响应后返回 + upstreamAction := "streamGenerateContent" + + // 执行带重试的请求 + result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: upstreamAction, + body: wrappedBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话 + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, + }) + if err != nil { + // 检查是否是账号切换信号,转换为 UpstreamFailoverError 让 Handler 切换账号 + if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusServiceUnavailable, + ForceCacheBilling: switchErr.IsStickySession, + } + } + // 区分客户端取消和真正的上游失败,返回更准确的错误消息 + if c.Request.Context().Err() != nil { + return nil, s.writeGoogleError(c, http.StatusBadGateway, "Client disconnected before upstream response") + } + return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries") + } + resp := result.resp + defer func() { + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + }() + + // 处理错误响应 + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + contentType := resp.Header.Get("Content-Type") + // 尽早关闭原始响应体,释放连接;后续逻辑仍可能需要读取 body,因此用内存副本重新包装。 + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + // 模型兜底:模型不存在且开启 fallback 时,自动用 fallback 模型重试一次 + if s.settingService != nil && s.settingService.IsModelFallbackEnabled(ctx) && + isModelNotFoundError(resp.StatusCode, respBody) { + fallbackModel := s.settingService.GetFallbackModel(ctx, PlatformAntigravity) + if fallbackModel != "" && fallbackModel != mappedModel { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Model not found (%s), retrying with fallback model %s (account: %s)", mappedModel, fallbackModel, account.Name) + + fallbackWrapped, err := s.wrapV1InternalRequest(projectID, fallbackModel, injectedBody) + if err == nil { + fallbackReq, err := antigravity.NewAPIRequest(ctx, upstreamAction, accessToken, fallbackWrapped) + if err == nil { + fallbackResp, err := s.httpUpstream.Do(fallbackReq, proxyURL, account.ID, account.Concurrency) + if err == nil && fallbackResp.StatusCode < 400 { + _ = resp.Body.Close() + resp = fallbackResp + } else if fallbackResp != nil { + _ = fallbackResp.Body.Close() + } + } + } + } + } + + // Gemini 原生请求中的 thoughtSignature 可能来自旧上下文/旧账号,触发上游严格校验后返回 + // "Corrupted thought signature."。检测到此类 400 时,将 thoughtSignature 清理为 dummy 值后重试一次。 + signatureCheckBody := respBody + if unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody); unwrapErr == nil && len(unwrapped) > 0 { + signatureCheckBody = unwrapped + } + if resp.StatusCode == http.StatusBadRequest && + s.settingService != nil && + s.settingService.IsSignatureRectifierEnabled(ctx) && + isSignatureRelatedError(signatureCheckBody) && + bytes.Contains(injectedBody, []byte(`"thoughtSignature"`)) { + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(signatureCheckBody))) + upstreamDetail := s.getUpstreamErrorDetail(signatureCheckBody) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "signature_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: detected signature-related 400, retrying with cleaned thought signatures", account.ID) + + cleanedInjectedBody := CleanGeminiNativeThoughtSignatures(injectedBody) + retryWrappedBody, wrapErr := s.wrapV1InternalRequest(projectID, mappedModel, cleanedInjectedBody) + if wrapErr == nil { + retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: upstreamAction, + body: retryWrappedBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, + }) + if retryErr == nil { + retryResp := retryResult.resp + if retryResp.StatusCode < 400 { + resp = retryResp + } else { + retryRespBody := s.readUpstreamErrorBody(retryResp) + _ = retryResp.Body.Close() + retryOpsBody := retryRespBody + if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 { + retryOpsBody = retryUnwrapped + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: retryResp.StatusCode, + UpstreamRequestID: retryResp.Header.Get("x-request-id"), + Kind: "signature_retry", + Message: sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(retryOpsBody))), + Detail: s.getUpstreamErrorDetail(retryOpsBody), + }) + respBody = retryRespBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryRespBody)), + } + contentType = resp.Header.Get("Content-Type") + } + } else { + if switchErr, ok := IsAntigravityAccountSwitchError(retryErr); ok { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: http.StatusServiceUnavailable, + Kind: "failover", + Message: sanitizeUpstreamErrorMessage(retryErr.Error()), + }) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusServiceUnavailable, + ForceCacheBilling: switchErr.IsStickySession, + } + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + Kind: "signature_retry_request_error", + Message: sanitizeUpstreamErrorMessage(retryErr.Error()), + }) + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry request failed: %v", account.ID, retryErr) + } + } else { + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry wrap failed: %v", account.ID, wrapErr) + } + } + + // fallback 成功:继续按正常响应处理 + if resp.StatusCode < 400 { + goto handleSuccess + } + + requestID := resp.Header.Get("x-request-id") + if requestID != "" { + c.Header("x-request-id", requestID) + } + + unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody) + unwrappedForOps := unwrapped + if unwrapErr != nil || len(unwrappedForOps) == 0 { + unwrappedForOps = respBody + } + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession) + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps) + + // Always record upstream context for Ops error logs, even when we will failover. + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + + // 精确匹配服务端配置类 400 错误,触发同账号重试 + failover + if resp.StatusCode == http.StatusBadRequest && isGoogleProjectConfigError(strings.ToLower(upstreamMsg)) { + log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestID, + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps, RetryableOnSameAccount: true} + } + + if s.shouldFailoverUpstreamError(resp.StatusCode) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestID, + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps} + } + if contentType == "" { + contentType = "application/json" + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestID, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream error status=%d body=%s", resp.StatusCode, truncateForLog(unwrappedForOps, 500)) + MarkResponseCommitted(c) + c.Data(resp.StatusCode, contentType, unwrappedForOps) + return nil, fmt.Errorf("antigravity upstream error: %d", resp.StatusCode) + } + +handleSuccess: + requestID := resp.Header.Get("x-request-id") + if requestID != "" { + c.Header("x-request-id", requestID) + } + + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + + if stream { + // 客户端要求流式,直接透传 + streamRes, err := s.handleGeminiStreamingResponse(c, resp, startTime) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + } else { + // 客户端要求非流式,收集流式响应后返回 + streamRes, err := s.handleGeminiStreamToNonStreaming(c, resp, startTime) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + } + + if usage == nil { + usage = &ClaudeUsage{} + } + + // 判断是否为图片生成模型 + imageCount := 0 + if isImageGenerationModel(mappedModel) { + // Gemini 图片生成 API 每次请求只生成一张图片(API 限制) + imageCount = 1 + } + + return &ForwardResult{ + RequestID: requestID, + Usage: *usage, + Model: originalModel, + UpstreamModel: billingModel, + Stream: stream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + ImageCount: imageCount, + ImageSize: imageSize, + ImageInputSize: imageInputSize, + }, nil +} + +// cleanGeminiRequest 清理 Gemini 请求体中的 Schema +func cleanGeminiRequest(body []byte) ([]byte, error) { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return nil, err + } + + modified := false + + // 1. 清理 Tools + if tools, ok := payload["tools"].([]any); ok && len(tools) > 0 { + for _, t := range tools { + toolMap, ok := t.(map[string]any) + if !ok { + continue + } + + // function_declarations (snake_case) or functionDeclarations (camelCase) + var funcs []any + if f, ok := toolMap["functionDeclarations"].([]any); ok { + funcs = f + } else if f, ok := toolMap["function_declarations"].([]any); ok { + funcs = f + } + + if len(funcs) == 0 { + continue + } + + for _, f := range funcs { + funcMap, ok := f.(map[string]any) + if !ok { + continue + } + + if params, ok := funcMap["parameters"].(map[string]any); ok { + antigravity.DeepCleanUndefined(params) + cleaned := antigravity.CleanJSONSchema(params) + funcMap["parameters"] = cleaned + modified = true + } + } + } + } + + if !modified { + return body, nil + } + + return json.Marshal(payload) +} + +// filterEmptyPartsFromGeminiRequest 过滤掉 parts 为空的消息 +// Gemini API 不接受空 parts,需要在请求前过滤 +func filterEmptyPartsFromGeminiRequest(body []byte) ([]byte, error) { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return nil, err + } + + contents, ok := payload["contents"].([]any) + if !ok || len(contents) == 0 { + return body, nil + } + + filtered := make([]any, 0, len(contents)) + modified := false + + for _, c := range contents { + contentMap, ok := c.(map[string]any) + if !ok { + filtered = append(filtered, c) + continue + } + + parts, hasParts := contentMap["parts"] + if !hasParts { + filtered = append(filtered, c) + continue + } + + partsSlice, ok := parts.([]any) + if !ok { + filtered = append(filtered, c) + continue + } + + // 跳过 parts 为空数组的消息 + if len(partsSlice) == 0 { + modified = true + continue + } + + filtered = append(filtered, c) + } + + if !modified { + return body, nil + } + + payload["contents"] = filtered + return json.Marshal(payload) +} diff --git a/backend/internal/service/antigravity_gateway_retry.go b/backend/internal/service/antigravity_gateway_retry.go new file mode 100644 index 0000000000..f2f5d44f84 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_retry.go @@ -0,0 +1,1279 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + mathrand "math/rand" + "net" + "net/http" + "os" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// antigravityRetryLoopParams 重试循环的参数 +type antigravityRetryLoopParams struct { + ctx context.Context + prefix string + account *Account + proxyURL string + accessToken string + action string + body []byte + c *gin.Context + httpUpstream HTTPUpstream + settingService *SettingService + accountRepo AccountRepository // 用于智能重试的模型级别限流 + handleError func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult + requestedModel string // 用于限流检查的原始请求模型 + isStickySession bool // 是否为粘性会话(用于账号切换时的缓存计费判断) + groupID int64 // 用于模型级限流时清除粘性会话 + sessionHash string // 用于模型级限流时清除粘性会话 +} + +// antigravityRetryLoopResult 重试循环的结果 +type antigravityRetryLoopResult struct { + resp *http.Response +} + +// resolveAntigravityForwardBaseURL 解析转发用 base URL。 +// +// 默认使用生产端点 cloudcode-pa.googleapis.com(antigravity.BaseURLs 的首个地址, +// 与账号 OAuth 登录/测试连接所用的 antigravity.BaseURL 一致)。 +// +// 历史上这里改用 ForwardBaseURLs()(把 daily/sandbox 排到首位)并默认取首个地址, +// 导致网关把带生产 OAuth token 的请求发到 daily-cloudcode-pa.sandbox.googleapis.com, +// 上游拒绝 → 账号被 401「Invalid bearer token」/502 打入临时不可调度且无法恢复 +// (见 #3611 / #2962)。后台「测试连接」用的是生产端点,所以「测试成功但网关 401」。 +// +// daily/sandbox 端点仅供内部联调,需显式设置 +// GATEWAY_ANTIGRAVITY_FORWARD_BASE_URL=daily(或 sandbox)才启用。 +func resolveAntigravityForwardBaseURL() string { + baseURLs := antigravity.BaseURLs + if len(baseURLs) == 0 { + return "" + } + mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv))) + if (mode == "daily" || mode == "sandbox") && len(baseURLs) > 1 { + return baseURLs[1] + } + return baseURLs[0] +} + +// smartRetryAction 智能重试的处理结果 +type smartRetryAction int + +const ( + smartRetryActionContinue smartRetryAction = iota // 继续默认重试逻辑 + smartRetryActionBreakWithResp // 结束循环并返回 resp + smartRetryActionContinueURL // 继续 URL fallback 循环 +) + +// smartRetryResult 智能重试的结果 +type smartRetryResult struct { + action smartRetryAction + resp *http.Response + err error + switchError *AntigravityAccountSwitchError // 模型限流时返回账号切换信号 +} + +// handleSmartRetry 处理 OAuth 账号的智能重试逻辑 +// 将 429/503 限流处理逻辑抽取为独立函数,减少 antigravityRetryLoop 的复杂度 +func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParams, resp *http.Response, respBody []byte, baseURL string, urlIdx int, availableURLs []string) *smartRetryResult { + // "Resource has been exhausted" 是 URL 级别限流,切换 URL(仅 429) + if resp.StatusCode == http.StatusTooManyRequests && isURLLevelRateLimit(respBody) && urlIdx < len(availableURLs)-1 { + logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (429): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) + return &smartRetryResult{action: smartRetryActionContinueURL} + } + + category := antigravity429Unknown + if resp.StatusCode == http.StatusTooManyRequests { + category = classifyAntigravity429(respBody) + } + + // 判断是否触发智能重试 + shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) + + // AI Credits 超量请求: + // 仅在上游明确返回免费配额耗尽时才允许切换到 credits。 + if resp.StatusCode == http.StatusTooManyRequests && + category == antigravity429QuotaExhausted && + p.account.IsOveragesEnabled() && + !p.account.isCreditsExhausted() { + result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody) + if result.handled && result.resp != nil { + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: result.resp, + } + } + } + + // 情况1: retryDelay >= 阈值,限流模型并切换账号 + if shouldRateLimitModel { + // 单账号 503 退避重试模式:不设限流、不切换账号,改为原地等待+重试 + // 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。 + // 多账号场景下切换账号是最优选择,但单账号场景下设限流毫无意义(只会导致双重等待)。 + if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { + return s.handleSingleAccountRetryInPlace(p, resp, respBody, baseURL, waitDuration, modelName) + } + + rateLimitDuration := waitDuration + if rateLimitDuration <= 0 { + rateLimitDuration = antigravityDefaultRateLimitDuration + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d oauth_long_delay model=%s account=%d upstream_retry_delay=%v body=%s (model rate limit, switch account)", + p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200)) + + resetAt := time.Now().Add(rateLimitDuration) + if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) { + p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID) + } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) + + // 返回账号切换信号,让上层切换账号重试 + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + switchError: &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: modelName, + IsStickySession: p.isStickySession, + }, + } + } + + // 情况2: retryDelay < 阈值(或 MODEL_CAPACITY_EXHAUSTED),智能重试 + if shouldSmartRetry { + var lastRetryResp *http.Response + var lastRetryBody []byte + + // MODEL_CAPACITY_EXHAUSTED 使用独立的重试参数(60 次,固定 1s 间隔) + maxAttempts := antigravitySmartRetryMaxAttempts + if isModelCapacityExhausted { + maxAttempts = antigravityModelCapacityRetryMaxAttempts + waitDuration = antigravityModelCapacityRetryWait + + // 全局去重:如果其他 goroutine 已在重试同一模型且尚在 cooldown 中,直接返回 503 + if modelName != "" { + modelCapacityExhaustedMu.RLock() + cooldownUntil, exists := modelCapacityExhaustedUntil[modelName] + modelCapacityExhaustedMu.RUnlock() + if exists && time.Now().Before(cooldownUntil) { + log.Printf("%s status=%d model_capacity_exhausted_dedup model=%s account=%d cooldown_until=%v (skip retry)", + p.prefix, resp.StatusCode, modelName, p.account.ID, cooldownUntil.Format("15:04:05")) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, + } + } + } + } + + for attempt := 1; attempt <= maxAttempts; attempt++ { + log.Printf("%s status=%d oauth_smart_retry attempt=%d/%d delay=%v model=%s account=%d", + p.prefix, resp.StatusCode, attempt, maxAttempts, waitDuration, modelName, p.account.ID) + + timer := time.NewTimer(waitDuration) + select { + case <-p.ctx.Done(): + timer.Stop() + log.Printf("%s status=context_canceled_during_smart_retry", p.prefix) + return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} + case <-timer.C: + } + + // 智能重试:创建新请求 + retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=smart_retry_request_build_failed error=%v", p.prefix, err) + p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, + } + } + + retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { + log.Printf("%s status=%d smart_retry_success attempt=%d/%d", p.prefix, retryResp.StatusCode, attempt, maxAttempts) + // 重试成功,清除 MODEL_CAPACITY_EXHAUSTED cooldown + if isModelCapacityExhausted && modelName != "" { + modelCapacityExhaustedMu.Lock() + delete(modelCapacityExhaustedUntil, modelName) + modelCapacityExhaustedMu.Unlock() + } + return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} + } + + // 网络错误时,继续重试 + if retryErr != nil || retryResp == nil { + log.Printf("%s status=smart_retry_network_error attempt=%d/%d error=%v", p.prefix, attempt, maxAttempts, retryErr) + continue + } + + // 重试失败,关闭之前的响应 + if lastRetryResp != nil { + _ = lastRetryResp.Body.Close() + } + lastRetryResp = retryResp + if retryResp != nil { + lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) + _ = retryResp.Body.Close() + } + + // 解析新的重试信息,用于下次重试的等待时间(MODEL_CAPACITY_EXHAUSTED 使用固定循环,跳过) + if !isModelCapacityExhausted && attempt < maxAttempts && lastRetryBody != nil { + newShouldRetry, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) + if newShouldRetry && newWaitDuration > 0 { + waitDuration = newWaitDuration + } + } + } + + // 所有重试都失败 + rateLimitDuration := waitDuration + if rateLimitDuration <= 0 { + rateLimitDuration = antigravityDefaultRateLimitDuration + } + retryBody := lastRetryBody + if retryBody == nil { + retryBody = respBody + } + + // MODEL_CAPACITY_EXHAUSTED:模型容量不足,切换账号无意义 + // 直接返回上游错误响应,不设置模型限流,不切换账号 + if isModelCapacityExhausted { + // 设置 cooldown,让后续请求快速失败,避免重复重试 + if modelName != "" { + modelCapacityExhaustedMu.Lock() + modelCapacityExhaustedUntil[modelName] = time.Now().Add(antigravityModelCapacityCooldown) + modelCapacityExhaustedMu.Unlock() + } + log.Printf("%s status=%d smart_retry_exhausted_model_capacity attempts=%d model=%s account=%d body=%s (model capacity exhausted, not switching account)", + p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + }, + } + } + + // 单账号 503 退避重试模式:智能重试耗尽后不设限流、不切换账号, + // 直接返回 503 让 Handler 层的单账号退避循环做最终处理。 + if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d smart_retry_exhausted_single_account attempts=%d model=%s account=%d body=%s (return 503 directly)", + p.prefix, resp.StatusCode, antigravitySmartRetryMaxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + }, + } + } + + log.Printf("%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)", + p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) + + resetAt := time.Now().Add(rateLimitDuration) + s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true) + + // 清除粘性会话绑定,避免下次请求仍命中限流账号 + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) + + // 返回账号切换信号,让上层切换账号重试 + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + switchError: &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: modelName, + IsStickySession: p.isStickySession, + }, + } + } + + // 未触发智能重试,继续默认重试逻辑 + return &smartRetryResult{action: smartRetryActionContinue} +} + +// handleSingleAccountRetryInPlace 单账号 503 退避重试的原地重试逻辑。 +// +// 在多账号场景下,收到 503 + 长 retryDelay(≥ 7s)时会设置模型限流 + 切换账号; +// 但在单账号场景下,设限流毫无意义(因为切换回来的还是同一个账号,还要等限流过期)。 +// 此方法改为在 Service 层原地等待 + 重试,避免双重等待问题: +// +// 旧流程:Service 设限流 → Handler 退避等待 → Service 等限流过期 → 再请求(总耗时 = 退避 + 限流) +// 新流程:Service 直接等 retryDelay → 重试 → 成功/再等 → 重试...(总耗时 ≈ 实际 retryDelay × 重试次数) +// +// 约束: +// - 单次等待不超过 antigravitySingleAccountSmartRetryMaxWait +// - 总累计等待不超过 antigravitySingleAccountSmartRetryTotalMaxWait +// - 最多重试 antigravitySingleAccountSmartRetryMaxAttempts 次 +func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace( + p antigravityRetryLoopParams, + resp *http.Response, + respBody []byte, + baseURL string, + waitDuration time.Duration, + modelName string, +) *smartRetryResult { + // 限制单次等待时间 + if waitDuration > antigravitySingleAccountSmartRetryMaxWait { + waitDuration = antigravitySingleAccountSmartRetryMaxWait + } + if waitDuration < antigravitySmartRetryMinWait { + waitDuration = antigravitySmartRetryMinWait + } + + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_in_place model=%s account=%d upstream_retry_delay=%v (retrying in-place instead of rate-limiting)", + p.prefix, resp.StatusCode, modelName, p.account.ID, waitDuration) + + var lastRetryResp *http.Response + var lastRetryBody []byte + totalWaited := time.Duration(0) + + for attempt := 1; attempt <= antigravitySingleAccountSmartRetryMaxAttempts; attempt++ { + // 检查累计等待是否超限 + if totalWaited+waitDuration > antigravitySingleAccountSmartRetryTotalMaxWait { + remaining := antigravitySingleAccountSmartRetryTotalMaxWait - totalWaited + if remaining <= 0 { + logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: total_wait_exceeded total=%v max=%v, giving up", + p.prefix, totalWaited, antigravitySingleAccountSmartRetryTotalMaxWait) + break + } + waitDuration = remaining + } + + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry attempt=%d/%d delay=%v total_waited=%v model=%s account=%d", + p.prefix, resp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, waitDuration, totalWaited, modelName, p.account.ID) + + timer := time.NewTimer(waitDuration) + select { + case <-p.ctx.Done(): + timer.Stop() + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_single_account_retry", p.prefix) + return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} + case <-timer.C: + } + totalWaited += waitDuration + + // 创建新请求 + retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: request_build_failed error=%v", p.prefix, err) + break + } + + retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_success attempt=%d/%d total_waited=%v", + p.prefix, retryResp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited) + // 关闭之前的响应 + if lastRetryResp != nil { + _ = lastRetryResp.Body.Close() + } + return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} + } + + // 网络错误时继续重试 + if retryErr != nil || retryResp == nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: network_error attempt=%d/%d error=%v", + p.prefix, attempt, antigravitySingleAccountSmartRetryMaxAttempts, retryErr) + continue + } + + // 关闭之前的响应 + if lastRetryResp != nil { + _ = lastRetryResp.Body.Close() + } + lastRetryResp = retryResp + lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) + _ = retryResp.Body.Close() + + // 解析新的重试信息,更新下次等待时间 + if attempt < antigravitySingleAccountSmartRetryMaxAttempts && lastRetryBody != nil { + _, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) + if newWaitDuration > 0 { + waitDuration = newWaitDuration + if waitDuration > antigravitySingleAccountSmartRetryMaxWait { + waitDuration = antigravitySingleAccountSmartRetryMaxWait + } + if waitDuration < antigravitySmartRetryMinWait { + waitDuration = antigravitySmartRetryMinWait + } + } + } + } + + // 所有重试都失败,不设限流,直接返回 503 + // Handler 层的单账号退避循环会做最终处理 + retryBody := lastRetryBody + if retryBody == nil { + retryBody = respBody + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_exhausted attempts=%d total_waited=%v model=%s account=%d body=%s (return 503 directly)", + p.prefix, resp.StatusCode, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited, modelName, p.account.ID, truncateForLog(retryBody, 200)) + + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + }, + } +} + +// antigravityRetryLoop 执行带 URL fallback 的重试循环 +func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { + // 预检查:模型限流 + overages 启用 + 积分未耗尽 → 直接注入 AI Credits + overagesInjected := false + if p.requestedModel != "" && p.account.Platform == PlatformAntigravity && + p.account.IsOveragesEnabled() && !p.account.isCreditsExhausted() && + p.account.isModelRateLimitedWithContext(p.ctx, p.requestedModel) { + if creditsBody := injectEnabledCreditTypes(p.body); creditsBody != nil { + p.body = creditsBody + overagesInjected = true + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: model_rate_limited_credits_inject model=%s account=%d (injecting enabledCreditTypes)", + p.prefix, p.requestedModel, p.account.ID) + } + } + + // 预检查:如果账号已限流,直接返回切换信号 + if p.requestedModel != "" { + if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 { + // 已注入积分的请求不再受普通模型限流预检查阻断。 + if overagesInjected { + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + } else if isSingleAccountRetry(p.ctx) { + // 单账号 503 退避重试模式:跳过限流预检查,直接发请求。 + // 首次请求设的限流是为了多账号调度器跳过该账号,在单账号模式下无意义。 + // 如果上游确实还不可用,handleSmartRetry → handleSingleAccountRetryInPlace + // 会在 Service 层原地等待+重试,不需要在预检查这里等。 + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: rate_limit_switch remaining=%v model=%s account=%d", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + return nil, &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: p.requestedModel, + IsStickySession: p.isStickySession, + } + } + } + } + + baseURL := resolveAntigravityForwardBaseURL() + if baseURL == "" { + return nil, errors.New("no antigravity forward base url configured") + } + availableURLs := []string{baseURL} + + var resp *http.Response + var usedBaseURL string + logBody := p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBody + maxBytes := 2048 + if p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > 0 { + maxBytes = p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + } + getUpstreamDetail := func(body []byte) string { + if !logBody { + return "" + } + return truncateString(string(body), maxBytes) + } + +urlFallbackLoop: + for urlIdx, baseURL := range availableURLs { + usedBaseURL = baseURL + allAttemptsInternal500 := true // 追踪本轮所有 attempt 是否全部命中 INTERNAL 500 + for attempt := 1; attempt <= antigravityMaxRetries; attempt++ { + select { + case <-p.ctx.Done(): + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled error=%v", p.prefix, p.ctx.Err()) + return nil, p.ctx.Err() + default: + } + + upstreamReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) + if err != nil { + return nil, err + } + + resp, err = p.httpUpstream.Do(upstreamReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if err == nil && resp == nil { + err = errors.New("upstream returned nil response") + } + if err != nil { + safeErr := sanitizeUpstreamErrorMessage(err.Error()) + appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ + Platform: p.account.Platform, + AccountID: p.account.ID, + AccountName: p.account.Name, + UpstreamStatusCode: 0, + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "request_error", + Message: safeErr, + }) + if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { + logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (connection error): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) + continue urlFallbackLoop + } + if attempt < antigravityMaxRetries { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retry=%d/%d error=%v", p.prefix, attempt, antigravityMaxRetries, err) + if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) + return nil, p.ctx.Err() + } + continue + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retries_exhausted error=%v", p.prefix, err) + setOpsUpstreamError(p.c, 0, safeErr, "") + return nil, fmt.Errorf("upstream request failed after retries: %w", err) + } + + // 统一处理错误响应 + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + + if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { + modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel) + s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, nil) + } + + // ★ 统一入口:自定义错误码 + 临时不可调度 + if handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled { + if policyErr != nil { + return nil, policyErr + } + resp = &http.Response{ + StatusCode: outStatus, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + } + break urlFallbackLoop + } + + // 429/503 限流处理:区分 URL 级别限流、智能重试和账户配额限流 + if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable { + // 尝试智能重试处理(OAuth 账号专用) + smartResult := s.handleSmartRetry(p, resp, respBody, baseURL, urlIdx, availableURLs) + switch smartResult.action { + case smartRetryActionContinueURL: + continue urlFallbackLoop + case smartRetryActionBreakWithResp: + if smartResult.err != nil { + return nil, smartResult.err + } + // 模型限流时返回切换账号信号 + if smartResult.switchError != nil { + return nil, smartResult.switchError + } + resp = smartResult.resp + break urlFallbackLoop + } + // smartRetryActionContinue: 继续默认重试逻辑 + + // 账户/模型配额限流,重试 3 次(指数退避)- 默认逻辑(非 OAuth 账号或解析失败) + if attempt < antigravityMaxRetries { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ + Platform: p.account.Platform, + AccountID: p.account.ID, + AccountName: p.account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "retry", + Message: upstreamMsg, + Detail: getUpstreamDetail(respBody), + }) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 200)) + if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) + return nil, p.ctx.Err() + } + continue + } + + // 重试用尽,标记账户限流 + p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited base_url=%s body=%s", p.prefix, resp.StatusCode, baseURL, truncateForLog(respBody, 200)) + resp = &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + } + break urlFallbackLoop + } + + // 其他可重试错误(500/502/504/529,不包括 429 和 503) + if shouldRetryAntigravityError(resp.StatusCode) { + if attempt < antigravityMaxRetries { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ + Platform: p.account.Platform, + AccountID: p.account.ID, + AccountName: p.account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "retry", + Message: upstreamMsg, + Detail: getUpstreamDetail(respBody), + }) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 500)) + if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) + return nil, p.ctx.Err() + } + // 追踪 INTERNAL 500:非匹配的 attempt 清除标记 + if !isAntigravityInternalServerError(resp.StatusCode, respBody) { + allAttemptsInternal500 = false + } + continue + } + } + + // INTERNAL 500 渐进惩罚:3 次重试全部命中特定 500 时递增计数器并惩罚 + if allAttemptsInternal500 && isAntigravityInternalServerError(resp.StatusCode, respBody) { + s.handleInternal500RetryExhausted(p.ctx, p.prefix, p.account) + } + + // 其他 4xx 错误或重试用尽,直接返回 + resp = &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + } + break urlFallbackLoop + } + + // 成功响应(< 400) + break urlFallbackLoop + } + } + + if resp != nil && resp.StatusCode < 400 && usedBaseURL != "" { + antigravity.DefaultURLAvailability.MarkSuccess(usedBaseURL) + } + + // 成功响应时清零 INTERNAL 500 连续失败计数器(覆盖所有成功路径,含 smart retry) + if resp != nil && resp.StatusCode < 400 { + s.resetInternal500Counter(p.ctx, p.prefix, p.account.ID) + } + + return &antigravityRetryLoopResult{resp: resp}, nil +} + +// shouldRetryAntigravityError 判断是否应该重试 +func shouldRetryAntigravityError(statusCode int) bool { + switch statusCode { + case 429, 500, 502, 503, 504, 529: + return true + default: + return false + } +} + +// isURLLevelRateLimit 判断是否为 URL 级别的限流(应切换 URL 重试) +// "Resource has been exhausted" 是 URL/节点级别限流,切换 URL 可能成功 +// "exhausted your capacity on this model" 是账户/模型配额限流,切换 URL 无效 +func isURLLevelRateLimit(body []byte) bool { + // 快速检查:包含 "Resource has been exhausted" 且不包含 "capacity on this model" + bodyStr := string(body) + return strings.Contains(bodyStr, "Resource has been exhausted") && + !strings.Contains(bodyStr, "capacity on this model") +} + +// isAntigravityConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝) +func isAntigravityConnectionError(err error) bool { + if err == nil { + return false + } + + // 检查超时错误 + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return true + } + + // 检查连接错误(DNS 失败、连接拒绝) + var opErr *net.OpError + return errors.As(err, &opErr) +} + +// shouldAntigravityFallbackToNextURL 判断是否应切换到下一个 URL +// 仅连接错误和 HTTP 429 触发 URL 降级 +func shouldAntigravityFallbackToNextURL(err error, statusCode int) bool { + if isAntigravityConnectionError(err) { + return true + } + return statusCode == http.StatusTooManyRequests +} + +// getSessionID 从 gin.Context 获取 session_id(用于日志追踪) +func getSessionID(c *gin.Context) string { + if c == nil { + return "" + } + return c.GetHeader("session_id") +} + +// logPrefix 生成统一的日志前缀 +func logPrefix(sessionID, accountName string) string { + if sessionID != "" { + return fmt.Sprintf("[antigravity-Forward] session=%s account=%s", sessionID, accountName) + } + return fmt.Sprintf("[antigravity-Forward] account=%s", accountName) +} + +func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int) bool { + switch statusCode { + case 401, 403, 429, 529: + return true + default: + return statusCode >= 500 + } +} + +// isGoogleProjectConfigError 判断(已提取的小写)错误消息是否属于 Google 服务端配置类问题。 +// 只精确匹配已知的服务端侧错误,避免对客户端请求错误做无意义重试。 +// 适用于所有走 Google 后端的平台(Antigravity、Gemini)。 +func isGoogleProjectConfigError(lowerMsg string) bool { + // Google 间歇性 Bug:Project ID 有效但被临时识别失败 + return strings.Contains(lowerMsg, "invalid project resource name") +} + +// googleConfigErrorCooldown 服务端配置类 400 错误的临时封禁时长 +const googleConfigErrorCooldown = 1 * time.Minute + +// tempUnscheduleGoogleConfigError 对服务端配置类 400 错误触发临时封禁, +// 避免短时间内反复调度到同一个有问题的账号。 +func tempUnscheduleGoogleConfigError(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { + until := time.Now().Add(googleConfigErrorCooldown) + reason := "400: invalid project resource name (auto temp-unschedule 1m)" + if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { + log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) + } else { + log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) + } +} + +// emptyResponseCooldown 空流式响应的临时封禁时长 +const emptyResponseCooldown = 1 * time.Minute + +// tempUnscheduleEmptyResponse 对空流式响应触发临时封禁, +// 避免短时间内反复调度到同一个返回空响应的账号。 +func tempUnscheduleEmptyResponse(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { + until := time.Now().Add(emptyResponseCooldown) + reason := "empty stream response (auto temp-unschedule 1m)" + if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { + log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) + } else { + log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) + } +} + +// sleepAntigravityBackoffWithContext 带 context 取消检查的退避等待 +// 返回 true 表示正常完成等待,false 表示 context 已取消 +func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool { + delay := antigravityRetryBaseDelay * time.Duration(1< antigravityRetryMaxDelay { + delay = antigravityRetryMaxDelay + } + + // +/- 20% jitter + r := mathrand.New(mathrand.NewSource(time.Now().UnixNano())) + jitter := time.Duration(float64(delay) * 0.2 * (r.Float64()*2 - 1)) + sleepFor := delay + jitter + if sleepFor < 0 { + sleepFor = 0 + } + + timer := time.NewTimer(sleepFor) + select { + case <-ctx.Done(): + timer.Stop() + return false + case <-timer.C: + return true + } +} + +// isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记 +func isSingleAccountRetry(ctx context.Context) bool { + v, _ := SingleAccountRetryFromContext(ctx) + return v +} + +// setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流 +// 直接使用上游返回的模型 ID(如 claude-sonnet-4-5)作为限流 key +// 返回是否已成功设置(若模型名为空或 repo 为 nil 将返回 false) +func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, accountID int64, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if repo == nil || modelName == "" { + return false + } + // 直接使用官方模型 ID 作为 key,不再转换为 scope + if err := repo.SetModelRateLimit(ctx, accountID, modelName, resetAt); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", prefix, statusCode, modelName, err) + return false + } + if afterSmartRetry { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) + } + return true +} + +func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if account == nil || repo == nil { + return false + } + keys := antigravityModelRateLimitKeys(modelName) + if len(keys) == 0 { + return false + } + + success := false + for _, key := range keys { + if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) { + s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt) + success = true + } + } + return success +} + +func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { + if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { + return + } + if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err) + } +} + +func antigravityFallbackCooldownSeconds() (time.Duration, bool) { + raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv)) + if raw == "" { + return 0, false + } + seconds, err := strconv.Atoi(raw) + if err != nil || seconds <= 0 { + return 0, false + } + return time.Duration(seconds) * time.Second, true +} + +// antigravitySmartRetryInfo 智能重试所需的信息 +type antigravitySmartRetryInfo struct { + RetryDelay time.Duration // 重试延迟时间 + ModelName string // 限流的模型名称(如 "claude-sonnet-4-5") + IsModelCapacityExhausted bool // 是否为模型容量不足(MODEL_CAPACITY_EXHAUSTED) +} + +// parseAntigravitySmartRetryInfo 解析 Google RPC RetryInfo 和 ErrorInfo 信息 +// 返回解析结果,如果解析失败或不满足条件返回 nil +// +// 支持两种情况: +// 1. 429 RESOURCE_EXHAUSTED + RATE_LIMIT_EXCEEDED: +// - error.status == "RESOURCE_EXHAUSTED" +// - error.details[].reason == "RATE_LIMIT_EXCEEDED" +// +// 2. 503 UNAVAILABLE + MODEL_CAPACITY_EXHAUSTED: +// - error.status == "UNAVAILABLE" +// - error.details[].reason == "MODEL_CAPACITY_EXHAUSTED" +// +// 必须满足以下条件才会返回有效值: +// - error.details[] 中存在 @type == "type.googleapis.com/google.rpc.RetryInfo" 的元素 +// - 该元素包含 retryDelay 字段,格式为 "数字s"(如 "0.201506475s") +func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { + var parsed map[string]any + if err := json.Unmarshal(body, &parsed); err != nil { + return nil + } + + errObj, ok := parsed["error"].(map[string]any) + if !ok { + return nil + } + + // 检查 status 是否符合条件 + // 情况1: 429 RESOURCE_EXHAUSTED (需要进一步检查 reason == RATE_LIMIT_EXCEEDED) + // 情况2: 503 UNAVAILABLE (需要进一步检查 reason == MODEL_CAPACITY_EXHAUSTED) + status, _ := errObj["status"].(string) + isResourceExhausted := status == googleRPCStatusResourceExhausted + isUnavailable := status == googleRPCStatusUnavailable + + if !isResourceExhausted && !isUnavailable { + return nil + } + + details, ok := errObj["details"].([]any) + if !ok { + return nil + } + + var retryDelay time.Duration + var modelName string + var hasRateLimitExceeded bool // 429 需要此 reason + var hasModelCapacityExhausted bool // 503 需要此 reason + + for _, d := range details { + dm, ok := d.(map[string]any) + if !ok { + continue + } + + atType, _ := dm["@type"].(string) + + // 从 ErrorInfo 提取模型名称和 reason + if atType == googleRPCTypeErrorInfo { + if meta, ok := dm["metadata"].(map[string]any); ok { + if model, ok := meta["model"].(string); ok { + modelName = normalizeAntigravityModelName(model) + } + } + // 检查 reason + if reason, ok := dm["reason"].(string); ok { + if reason == googleRPCReasonModelCapacityExhausted { + hasModelCapacityExhausted = true + } + if reason == googleRPCReasonRateLimitExceeded { + hasRateLimitExceeded = true + } + } + continue + } + + // 从 RetryInfo 提取重试延迟 + if atType == googleRPCTypeRetryInfo { + delay, ok := dm["retryDelay"].(string) + if !ok || delay == "" { + continue + } + // 使用 time.ParseDuration 解析,支持所有 Go duration 格式 + // 例如: "0.5s", "10s", "4m50s", "1h30m", "200ms" 等 + dur, err := time.ParseDuration(delay) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] failed to parse retryDelay: %s error=%v", delay, err) + continue + } + retryDelay = dur + } + } + + // 验证条件 + // 情况1: RESOURCE_EXHAUSTED 需要有 RATE_LIMIT_EXCEEDED reason + // 情况2: UNAVAILABLE 需要有 MODEL_CAPACITY_EXHAUSTED reason + if isResourceExhausted && !hasRateLimitExceeded { + return nil + } + if isUnavailable && !hasModelCapacityExhausted { + return nil + } + + // 必须有模型名才返回有效结果 + if modelName == "" { + return nil + } + + // 如果上游未提供 retryDelay,使用默认限流时间 + if retryDelay <= 0 { + retryDelay = antigravityDefaultRateLimitDuration + } + + return &antigravitySmartRetryInfo{ + RetryDelay: retryDelay, + ModelName: modelName, + IsModelCapacityExhausted: hasModelCapacityExhausted, + } +} + +// shouldTriggerAntigravitySmartRetry 判断是否应该触发智能重试 +// 返回: +// - shouldRetry: 是否应该智能重试(retryDelay < antigravityRateLimitThreshold,或 MODEL_CAPACITY_EXHAUSTED) +// - shouldRateLimitModel: 是否应该限流模型并切换账号(仅 RATE_LIMIT_EXCEEDED 且 retryDelay >= 阈值) +// - waitDuration: 等待时间 +// - modelName: 限流的模型名称 +// - isModelCapacityExhausted: 是否为模型容量不足(MODEL_CAPACITY_EXHAUSTED) +func shouldTriggerAntigravitySmartRetry(account *Account, respBody []byte) (shouldRetry bool, shouldRateLimitModel bool, waitDuration time.Duration, modelName string, isModelCapacityExhausted bool) { + if account.Platform != PlatformAntigravity { + return false, false, 0, "", false + } + + info := parseAntigravitySmartRetryInfo(respBody) + if info == nil { + return false, false, 0, "", false + } + + // MODEL_CAPACITY_EXHAUSTED(模型容量不足):所有账号共享同一模型容量池 + // 切换账号无意义,使用固定 1s 间隔重试 + if info.IsModelCapacityExhausted { + return true, false, antigravityModelCapacityRetryWait, info.ModelName, true + } + + // RATE_LIMIT_EXCEEDED(账号级限流): + // retryDelay >= 阈值:直接限流模型,不重试 + // 注意:如果上游未提供 retryDelay,parseAntigravitySmartRetryInfo 已设置为默认 30s + if info.RetryDelay >= antigravityRateLimitThreshold { + return false, true, info.RetryDelay, info.ModelName, false + } + + // retryDelay < 阈值:智能重试 + waitDuration = info.RetryDelay + if waitDuration < antigravitySmartRetryMinWait { + waitDuration = antigravitySmartRetryMinWait + } + + return true, false, waitDuration, info.ModelName, false +} + +// handleModelRateLimitParams 模型级限流处理参数 +type handleModelRateLimitParams struct { + ctx context.Context + prefix string + account *Account + statusCode int + body []byte + cache GatewayCache + groupID int64 + sessionHash string + isStickySession bool +} + +// handleModelRateLimitResult 模型级限流处理结果 +type handleModelRateLimitResult struct { + Handled bool // 是否已处理 + ShouldRetry bool // 是否等待后重试 + WaitDuration time.Duration // 等待时间 + SwitchError *AntigravityAccountSwitchError // 账号切换错误 +} + +// handleModelRateLimit 处理模型级限流(在原有逻辑之前调用) +// 仅处理 429/503,解析模型名和 retryDelay +// - MODEL_CAPACITY_EXHAUSTED: 返回 Handled=true(实际重试由 handleSmartRetry 处理) +// - RATE_LIMIT_EXCEEDED + retryDelay < 阈值: 返回 ShouldRetry=true,由调用方等待后重试 +// - RATE_LIMIT_EXCEEDED + retryDelay >= 阈值: 设置模型限流 + 清除粘性会话 + 返回 SwitchError +func (s *AntigravityGatewayService) handleModelRateLimit(p *handleModelRateLimitParams) *handleModelRateLimitResult { + if p.statusCode != 429 && p.statusCode != 503 { + return &handleModelRateLimitResult{Handled: false} + } + + info := parseAntigravitySmartRetryInfo(p.body) + if info == nil || info.ModelName == "" { + return &handleModelRateLimitResult{Handled: false} + } + + // MODEL_CAPACITY_EXHAUSTED:模型容量不足,所有账号共享同一容量池 + // 切换账号无意义,不设置模型限流(实际重试由 handleSmartRetry 处理) + if info.IsModelCapacityExhausted { + log.Printf("%s status=%d model_capacity_exhausted model=%s (not switching account, retry handled by smart retry)", + p.prefix, p.statusCode, info.ModelName) + return &handleModelRateLimitResult{ + Handled: true, + } + } + + // RATE_LIMIT_EXCEEDED: < antigravityRateLimitThreshold: 等待后重试 + if info.RetryDelay < antigravityRateLimitThreshold { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_wait model=%s wait=%v", + p.prefix, p.statusCode, info.ModelName, info.RetryDelay) + return &handleModelRateLimitResult{ + Handled: true, + ShouldRetry: true, + WaitDuration: info.RetryDelay, + } + } + + // RATE_LIMIT_EXCEEDED: >= antigravityRateLimitThreshold: 设置限流 + 清除粘性会话 + 切换账号 + s.setModelRateLimitAndClearSession(p, info) + + return &handleModelRateLimitResult{ + Handled: true, + SwitchError: &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: info.ModelName, + IsStickySession: p.isStickySession, + }, + } +} + +// setModelRateLimitAndClearSession 设置模型限流并清除粘性会话 +func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleModelRateLimitParams, info *antigravitySmartRetryInfo) { + resetAt := time.Now().Add(info.RetryDelay) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", + p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) + + s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false) + + // 清除粘性会话绑定 + if p.cache != nil && p.sessionHash != "" { + _ = p.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash) + } +} + +// updateAccountModelRateLimitInCache 立即更新 Redis 中账号的模型限流状态 +func (s *AntigravityGatewayService) updateAccountModelRateLimitInCache(ctx context.Context, account *Account, modelKey string, resetAt time.Time) { + if s.schedulerSnapshot == nil || account == nil || modelKey == "" { + return + } + + // 更新账号对象的 Extra 字段 + if account.Extra == nil { + account.Extra = make(map[string]any) + } + + limits, _ := account.Extra["model_rate_limits"].(map[string]any) + if limits == nil { + limits = make(map[string]any) + account.Extra["model_rate_limits"] = limits + } + + limits[modelKey] = map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339), + } + + // 更新 Redis 快照 + if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] cache_update_failed account=%d model=%s err=%v", account.ID, modelKey, err) + } +} + +func (s *AntigravityGatewayService) handleUpstreamError( + ctx context.Context, prefix string, account *Account, + statusCode int, headers http.Header, body []byte, + requestedModel string, + groupID int64, sessionHash string, isStickySession bool, +) *handleModelRateLimitResult { + // 遵守自定义错误码策略:未命中则跳过所有限流处理 + if !account.ShouldHandleErrorCode(statusCode) { + return nil + } + // 模型级限流处理(优先) + result := s.handleModelRateLimit(&handleModelRateLimitParams{ + ctx: ctx, + prefix: prefix, + account: account, + statusCode: statusCode, + body: body, + cache: s.cache, + groupID: groupID, + sessionHash: sessionHash, + isStickySession: isStickySession, + }) + if result.Handled { + return result + } + + // 503 仅处理模型限流(MODEL_CAPACITY_EXHAUSTED),非模型限流不做额外处理 + // 避免将普通的 503 错误误判为账号问题 + if statusCode == 503 { + return nil + } + + // 429:尝试解析模型级限流,解析失败时兜底为账号级限流 + if statusCode == 429 { + if logBody, maxBytes := s.getLogConfig(); logBody { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity-Debug] 429 response body: %s", truncateString(string(body), maxBytes)) + } + + resetAt := ParseGeminiRateLimitResetTime(body) + defaultDur := s.getDefaultRateLimitDuration() + + // 尝试解析模型 key 并设置模型级限流 + // + // 注意:requestedModel 可能是"映射前"的请求模型名(例如 claude-opus-4-6), + // 调度与限流判定使用的是 Antigravity 最终模型名(包含映射与 thinking 后缀)。 + // 因此这里必须写入最终模型 key,确保后续调度能正确避开已限流模型。 + modelKey := resolveFinalAntigravityModelKey(ctx, account, requestedModel) + if strings.TrimSpace(modelKey) == "" { + // 极少数情况下无法映射(理论上不应发生:能转发成功说明映射已通过), + // 保持旧行为作为兜底,避免完全丢失模型级限流记录。 + modelKey = resolveAntigravityModelKey(requestedModel) + } + if modelKey != "" { + ra := s.resolveResetTime(resetAt, defaultDur) + if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v", + prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) + } + return nil + } + + // 无法解析模型 key,兜底为账号级限流 + ra := s.resolveResetTime(resetAt, defaultDur) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited account=%d reset_at=%v reset_in=%v (fallback)", + prefix, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) + if err := s.accountRepo.SetRateLimited(ctx, account.ID, ra); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limit_set_failed account=%d error=%v", prefix, account.ID, err) + } + return nil + } + // 其他错误码继续使用 rateLimitService + if s.rateLimitService == nil { + return nil + } + shouldDisable := s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body) + if shouldDisable { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d marked_error", prefix, statusCode) + } + return nil +} + +// getDefaultRateLimitDuration 获取默认限流时间 +func (s *AntigravityGatewayService) getDefaultRateLimitDuration() time.Duration { + defaultDur := antigravityDefaultRateLimitDuration + if s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes > 0 { + defaultDur = time.Duration(s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes) * time.Minute + } + if override, ok := antigravityFallbackCooldownSeconds(); ok { + defaultDur = override + } + return defaultDur +} + +// resolveResetTime 根据解析的重置时间或默认时长计算重置时间点 +func (s *AntigravityGatewayService) resolveResetTime(resetAt *int64, defaultDur time.Duration) time.Time { + if resetAt != nil { + return time.Unix(*resetAt, 0) + } + return time.Now().Add(defaultDur) +} diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index aa4cab22d7..9e0cc804d5 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -1,28 +1,20 @@ package service import ( - "bufio" "bytes" "context" "encoding/json" "errors" "fmt" "io" - "log" "log/slog" - mathrand "math/rand" - "net" "net/http" - "os" - "strconv" "strings" "sync" - "sync/atomic" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/tidwall/gjson" ) @@ -127,735 +119,6 @@ func (e *PromptTooLongError) Error() string { return fmt.Sprintf("prompt too long: status=%d", e.StatusCode) } -// antigravityRetryLoopParams 重试循环的参数 -type antigravityRetryLoopParams struct { - ctx context.Context - prefix string - account *Account - proxyURL string - accessToken string - action string - body []byte - c *gin.Context - httpUpstream HTTPUpstream - settingService *SettingService - accountRepo AccountRepository // 用于智能重试的模型级别限流 - handleError func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult - requestedModel string // 用于限流检查的原始请求模型 - isStickySession bool // 是否为粘性会话(用于账号切换时的缓存计费判断) - groupID int64 // 用于模型级限流时清除粘性会话 - sessionHash string // 用于模型级限流时清除粘性会话 -} - -// antigravityRetryLoopResult 重试循环的结果 -type antigravityRetryLoopResult struct { - resp *http.Response -} - -// resolveAntigravityForwardBaseURL 解析转发用 base URL。 -// 默认使用 daily(ForwardBaseURLs 的首个地址);当环境变量为 prod 时使用第二个地址。 -func resolveAntigravityForwardBaseURL() string { - baseURLs := antigravity.ForwardBaseURLs() - if len(baseURLs) == 0 { - return "" - } - mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv))) - if mode == "prod" && len(baseURLs) > 1 { - return baseURLs[1] - } - return baseURLs[0] -} - -// smartRetryAction 智能重试的处理结果 -type smartRetryAction int - -const ( - smartRetryActionContinue smartRetryAction = iota // 继续默认重试逻辑 - smartRetryActionBreakWithResp // 结束循环并返回 resp - smartRetryActionContinueURL // 继续 URL fallback 循环 -) - -// smartRetryResult 智能重试的结果 -type smartRetryResult struct { - action smartRetryAction - resp *http.Response - err error - switchError *AntigravityAccountSwitchError // 模型限流时返回账号切换信号 -} - -// handleSmartRetry 处理 OAuth 账号的智能重试逻辑 -// 将 429/503 限流处理逻辑抽取为独立函数,减少 antigravityRetryLoop 的复杂度 -func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParams, resp *http.Response, respBody []byte, baseURL string, urlIdx int, availableURLs []string) *smartRetryResult { - // "Resource has been exhausted" 是 URL 级别限流,切换 URL(仅 429) - if resp.StatusCode == http.StatusTooManyRequests && isURLLevelRateLimit(respBody) && urlIdx < len(availableURLs)-1 { - logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (429): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) - return &smartRetryResult{action: smartRetryActionContinueURL} - } - - category := antigravity429Unknown - if resp.StatusCode == http.StatusTooManyRequests { - category = classifyAntigravity429(respBody) - } - - // 判断是否触发智能重试 - shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) - - // AI Credits 超量请求: - // 仅在上游明确返回免费配额耗尽时才允许切换到 credits。 - if resp.StatusCode == http.StatusTooManyRequests && - category == antigravity429QuotaExhausted && - p.account.IsOveragesEnabled() && - !p.account.isCreditsExhausted() { - result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody) - if result.handled && result.resp != nil { - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: result.resp, - } - } - } - - // 情况1: retryDelay >= 阈值,限流模型并切换账号 - if shouldRateLimitModel { - // 单账号 503 退避重试模式:不设限流、不切换账号,改为原地等待+重试 - // 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。 - // 多账号场景下切换账号是最优选择,但单账号场景下设限流毫无意义(只会导致双重等待)。 - if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { - return s.handleSingleAccountRetryInPlace(p, resp, respBody, baseURL, waitDuration, modelName) - } - - rateLimitDuration := waitDuration - if rateLimitDuration <= 0 { - rateLimitDuration = antigravityDefaultRateLimitDuration - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d oauth_long_delay model=%s account=%d upstream_retry_delay=%v body=%s (model rate limit, switch account)", - p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200)) - - resetAt := time.Now().Add(rateLimitDuration) - if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) { - p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID) - } - s.clearStickySession(p.ctx, p.groupID, p.sessionHash) - - // 返回账号切换信号,让上层切换账号重试 - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - switchError: &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: modelName, - IsStickySession: p.isStickySession, - }, - } - } - - // 情况2: retryDelay < 阈值(或 MODEL_CAPACITY_EXHAUSTED),智能重试 - if shouldSmartRetry { - var lastRetryResp *http.Response - var lastRetryBody []byte - - // MODEL_CAPACITY_EXHAUSTED 使用独立的重试参数(60 次,固定 1s 间隔) - maxAttempts := antigravitySmartRetryMaxAttempts - if isModelCapacityExhausted { - maxAttempts = antigravityModelCapacityRetryMaxAttempts - waitDuration = antigravityModelCapacityRetryWait - - // 全局去重:如果其他 goroutine 已在重试同一模型且尚在 cooldown 中,直接返回 503 - if modelName != "" { - modelCapacityExhaustedMu.RLock() - cooldownUntil, exists := modelCapacityExhaustedUntil[modelName] - modelCapacityExhaustedMu.RUnlock() - if exists && time.Now().Before(cooldownUntil) { - log.Printf("%s status=%d model_capacity_exhausted_dedup model=%s account=%d cooldown_until=%v (skip retry)", - p.prefix, resp.StatusCode, modelName, p.account.ID, cooldownUntil.Format("15:04:05")) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - }, - } - } - } - } - - for attempt := 1; attempt <= maxAttempts; attempt++ { - log.Printf("%s status=%d oauth_smart_retry attempt=%d/%d delay=%v model=%s account=%d", - p.prefix, resp.StatusCode, attempt, maxAttempts, waitDuration, modelName, p.account.ID) - - timer := time.NewTimer(waitDuration) - select { - case <-p.ctx.Done(): - timer.Stop() - log.Printf("%s status=context_canceled_during_smart_retry", p.prefix) - return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} - case <-timer.C: - } - - // 智能重试:创建新请求 - retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=smart_retry_request_build_failed error=%v", p.prefix, err) - p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - }, - } - } - - retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) - if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { - log.Printf("%s status=%d smart_retry_success attempt=%d/%d", p.prefix, retryResp.StatusCode, attempt, maxAttempts) - // 重试成功,清除 MODEL_CAPACITY_EXHAUSTED cooldown - if isModelCapacityExhausted && modelName != "" { - modelCapacityExhaustedMu.Lock() - delete(modelCapacityExhaustedUntil, modelName) - modelCapacityExhaustedMu.Unlock() - } - return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} - } - - // 网络错误时,继续重试 - if retryErr != nil || retryResp == nil { - log.Printf("%s status=smart_retry_network_error attempt=%d/%d error=%v", p.prefix, attempt, maxAttempts, retryErr) - continue - } - - // 重试失败,关闭之前的响应 - if lastRetryResp != nil { - _ = lastRetryResp.Body.Close() - } - lastRetryResp = retryResp - if retryResp != nil { - lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) - _ = retryResp.Body.Close() - } - - // 解析新的重试信息,用于下次重试的等待时间(MODEL_CAPACITY_EXHAUSTED 使用固定循环,跳过) - if !isModelCapacityExhausted && attempt < maxAttempts && lastRetryBody != nil { - newShouldRetry, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) - if newShouldRetry && newWaitDuration > 0 { - waitDuration = newWaitDuration - } - } - } - - // 所有重试都失败 - rateLimitDuration := waitDuration - if rateLimitDuration <= 0 { - rateLimitDuration = antigravityDefaultRateLimitDuration - } - retryBody := lastRetryBody - if retryBody == nil { - retryBody = respBody - } - - // MODEL_CAPACITY_EXHAUSTED:模型容量不足,切换账号无意义 - // 直接返回上游错误响应,不设置模型限流,不切换账号 - if isModelCapacityExhausted { - // 设置 cooldown,让后续请求快速失败,避免重复重试 - if modelName != "" { - modelCapacityExhaustedMu.Lock() - modelCapacityExhaustedUntil[modelName] = time.Now().Add(antigravityModelCapacityCooldown) - modelCapacityExhaustedMu.Unlock() - } - log.Printf("%s status=%d smart_retry_exhausted_model_capacity attempts=%d model=%s account=%d body=%s (model capacity exhausted, not switching account)", - p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - }, - } - } - - // 单账号 503 退避重试模式:智能重试耗尽后不设限流、不切换账号, - // 直接返回 503 让 Handler 层的单账号退避循环做最终处理。 - if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d smart_retry_exhausted_single_account attempts=%d model=%s account=%d body=%s (return 503 directly)", - p.prefix, resp.StatusCode, antigravitySmartRetryMaxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - }, - } - } - - log.Printf("%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)", - p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) - - resetAt := time.Now().Add(rateLimitDuration) - s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true) - - // 清除粘性会话绑定,避免下次请求仍命中限流账号 - s.clearStickySession(p.ctx, p.groupID, p.sessionHash) - - // 返回账号切换信号,让上层切换账号重试 - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - switchError: &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: modelName, - IsStickySession: p.isStickySession, - }, - } - } - - // 未触发智能重试,继续默认重试逻辑 - return &smartRetryResult{action: smartRetryActionContinue} -} - -// handleSingleAccountRetryInPlace 单账号 503 退避重试的原地重试逻辑。 -// -// 在多账号场景下,收到 503 + 长 retryDelay(≥ 7s)时会设置模型限流 + 切换账号; -// 但在单账号场景下,设限流毫无意义(因为切换回来的还是同一个账号,还要等限流过期)。 -// 此方法改为在 Service 层原地等待 + 重试,避免双重等待问题: -// -// 旧流程:Service 设限流 → Handler 退避等待 → Service 等限流过期 → 再请求(总耗时 = 退避 + 限流) -// 新流程:Service 直接等 retryDelay → 重试 → 成功/再等 → 重试...(总耗时 ≈ 实际 retryDelay × 重试次数) -// -// 约束: -// - 单次等待不超过 antigravitySingleAccountSmartRetryMaxWait -// - 总累计等待不超过 antigravitySingleAccountSmartRetryTotalMaxWait -// - 最多重试 antigravitySingleAccountSmartRetryMaxAttempts 次 -func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace( - p antigravityRetryLoopParams, - resp *http.Response, - respBody []byte, - baseURL string, - waitDuration time.Duration, - modelName string, -) *smartRetryResult { - // 限制单次等待时间 - if waitDuration > antigravitySingleAccountSmartRetryMaxWait { - waitDuration = antigravitySingleAccountSmartRetryMaxWait - } - if waitDuration < antigravitySmartRetryMinWait { - waitDuration = antigravitySmartRetryMinWait - } - - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_in_place model=%s account=%d upstream_retry_delay=%v (retrying in-place instead of rate-limiting)", - p.prefix, resp.StatusCode, modelName, p.account.ID, waitDuration) - - var lastRetryResp *http.Response - var lastRetryBody []byte - totalWaited := time.Duration(0) - - for attempt := 1; attempt <= antigravitySingleAccountSmartRetryMaxAttempts; attempt++ { - // 检查累计等待是否超限 - if totalWaited+waitDuration > antigravitySingleAccountSmartRetryTotalMaxWait { - remaining := antigravitySingleAccountSmartRetryTotalMaxWait - totalWaited - if remaining <= 0 { - logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: total_wait_exceeded total=%v max=%v, giving up", - p.prefix, totalWaited, antigravitySingleAccountSmartRetryTotalMaxWait) - break - } - waitDuration = remaining - } - - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry attempt=%d/%d delay=%v total_waited=%v model=%s account=%d", - p.prefix, resp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, waitDuration, totalWaited, modelName, p.account.ID) - - timer := time.NewTimer(waitDuration) - select { - case <-p.ctx.Done(): - timer.Stop() - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_single_account_retry", p.prefix) - return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} - case <-timer.C: - } - totalWaited += waitDuration - - // 创建新请求 - retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: request_build_failed error=%v", p.prefix, err) - break - } - - retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) - if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_success attempt=%d/%d total_waited=%v", - p.prefix, retryResp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited) - // 关闭之前的响应 - if lastRetryResp != nil { - _ = lastRetryResp.Body.Close() - } - return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} - } - - // 网络错误时继续重试 - if retryErr != nil || retryResp == nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: network_error attempt=%d/%d error=%v", - p.prefix, attempt, antigravitySingleAccountSmartRetryMaxAttempts, retryErr) - continue - } - - // 关闭之前的响应 - if lastRetryResp != nil { - _ = lastRetryResp.Body.Close() - } - lastRetryResp = retryResp - lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) - _ = retryResp.Body.Close() - - // 解析新的重试信息,更新下次等待时间 - if attempt < antigravitySingleAccountSmartRetryMaxAttempts && lastRetryBody != nil { - _, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) - if newWaitDuration > 0 { - waitDuration = newWaitDuration - if waitDuration > antigravitySingleAccountSmartRetryMaxWait { - waitDuration = antigravitySingleAccountSmartRetryMaxWait - } - if waitDuration < antigravitySmartRetryMinWait { - waitDuration = antigravitySmartRetryMinWait - } - } - } - } - - // 所有重试都失败,不设限流,直接返回 503 - // Handler 层的单账号退避循环会做最终处理 - retryBody := lastRetryBody - if retryBody == nil { - retryBody = respBody - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_exhausted attempts=%d total_waited=%v model=%s account=%d body=%s (return 503 directly)", - p.prefix, resp.StatusCode, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited, modelName, p.account.ID, truncateForLog(retryBody, 200)) - - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - }, - } -} - -// antigravityRetryLoop 执行带 URL fallback 的重试循环 -func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { - // 预检查:模型限流 + overages 启用 + 积分未耗尽 → 直接注入 AI Credits - overagesInjected := false - if p.requestedModel != "" && p.account.Platform == PlatformAntigravity && - p.account.IsOveragesEnabled() && !p.account.isCreditsExhausted() && - p.account.isModelRateLimitedWithContext(p.ctx, p.requestedModel) { - if creditsBody := injectEnabledCreditTypes(p.body); creditsBody != nil { - p.body = creditsBody - overagesInjected = true - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: model_rate_limited_credits_inject model=%s account=%d (injecting enabledCreditTypes)", - p.prefix, p.requestedModel, p.account.ID) - } - } - - // 预检查:如果账号已限流,直接返回切换信号 - if p.requestedModel != "" { - if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 { - // 已注入积分的请求不再受普通模型限流预检查阻断。 - if overagesInjected { - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - } else if isSingleAccountRetry(p.ctx) { - // 单账号 503 退避重试模式:跳过限流预检查,直接发请求。 - // 首次请求设的限流是为了多账号调度器跳过该账号,在单账号模式下无意义。 - // 如果上游确实还不可用,handleSmartRetry → handleSingleAccountRetryInPlace - // 会在 Service 层原地等待+重试,不需要在预检查这里等。 - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: rate_limit_switch remaining=%v model=%s account=%d", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - return nil, &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: p.requestedModel, - IsStickySession: p.isStickySession, - } - } - } - } - - baseURL := resolveAntigravityForwardBaseURL() - if baseURL == "" { - return nil, errors.New("no antigravity forward base url configured") - } - availableURLs := []string{baseURL} - - var resp *http.Response - var usedBaseURL string - logBody := p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBody - maxBytes := 2048 - if p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > 0 { - maxBytes = p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - } - getUpstreamDetail := func(body []byte) string { - if !logBody { - return "" - } - return truncateString(string(body), maxBytes) - } - -urlFallbackLoop: - for urlIdx, baseURL := range availableURLs { - usedBaseURL = baseURL - allAttemptsInternal500 := true // 追踪本轮所有 attempt 是否全部命中 INTERNAL 500 - for attempt := 1; attempt <= antigravityMaxRetries; attempt++ { - select { - case <-p.ctx.Done(): - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled error=%v", p.prefix, p.ctx.Err()) - return nil, p.ctx.Err() - default: - } - - upstreamReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) - if err != nil { - return nil, err - } - - resp, err = p.httpUpstream.Do(upstreamReq, p.proxyURL, p.account.ID, p.account.Concurrency) - if err == nil && resp == nil { - err = errors.New("upstream returned nil response") - } - if err != nil { - safeErr := sanitizeUpstreamErrorMessage(err.Error()) - appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ - Platform: p.account.Platform, - AccountID: p.account.ID, - AccountName: p.account.Name, - UpstreamStatusCode: 0, - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "request_error", - Message: safeErr, - }) - if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { - logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (connection error): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) - continue urlFallbackLoop - } - if attempt < antigravityMaxRetries { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retry=%d/%d error=%v", p.prefix, attempt, antigravityMaxRetries, err) - if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) - return nil, p.ctx.Err() - } - continue - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retries_exhausted error=%v", p.prefix, err) - setOpsUpstreamError(p.c, 0, safeErr, "") - return nil, fmt.Errorf("upstream request failed after retries: %w", err) - } - - // 统一处理错误响应 - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - _ = resp.Body.Close() - - if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { - modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel) - s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - }, nil) - } - - // ★ 统一入口:自定义错误码 + 临时不可调度 - if handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled { - if policyErr != nil { - return nil, policyErr - } - resp = &http.Response{ - StatusCode: outStatus, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - } - break urlFallbackLoop - } - - // 429/503 限流处理:区分 URL 级别限流、智能重试和账户配额限流 - if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable { - // 尝试智能重试处理(OAuth 账号专用) - smartResult := s.handleSmartRetry(p, resp, respBody, baseURL, urlIdx, availableURLs) - switch smartResult.action { - case smartRetryActionContinueURL: - continue urlFallbackLoop - case smartRetryActionBreakWithResp: - if smartResult.err != nil { - return nil, smartResult.err - } - // 模型限流时返回切换账号信号 - if smartResult.switchError != nil { - return nil, smartResult.switchError - } - resp = smartResult.resp - break urlFallbackLoop - } - // smartRetryActionContinue: 继续默认重试逻辑 - - // 账户/模型配额限流,重试 3 次(指数退避)- 默认逻辑(非 OAuth 账号或解析失败) - if attempt < antigravityMaxRetries { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ - Platform: p.account.Platform, - AccountID: p.account.ID, - AccountName: p.account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "retry", - Message: upstreamMsg, - Detail: getUpstreamDetail(respBody), - }) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 200)) - if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) - return nil, p.ctx.Err() - } - continue - } - - // 重试用尽,标记账户限流 - p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited base_url=%s body=%s", p.prefix, resp.StatusCode, baseURL, truncateForLog(respBody, 200)) - resp = &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - } - break urlFallbackLoop - } - - // 其他可重试错误(500/502/504/529,不包括 429 和 503) - if shouldRetryAntigravityError(resp.StatusCode) { - if attempt < antigravityMaxRetries { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ - Platform: p.account.Platform, - AccountID: p.account.ID, - AccountName: p.account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "retry", - Message: upstreamMsg, - Detail: getUpstreamDetail(respBody), - }) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 500)) - if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) - return nil, p.ctx.Err() - } - // 追踪 INTERNAL 500:非匹配的 attempt 清除标记 - if !isAntigravityInternalServerError(resp.StatusCode, respBody) { - allAttemptsInternal500 = false - } - continue - } - } - - // INTERNAL 500 渐进惩罚:3 次重试全部命中特定 500 时递增计数器并惩罚 - if allAttemptsInternal500 && isAntigravityInternalServerError(resp.StatusCode, respBody) { - s.handleInternal500RetryExhausted(p.ctx, p.prefix, p.account) - } - - // 其他 4xx 错误或重试用尽,直接返回 - resp = &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - } - break urlFallbackLoop - } - - // 成功响应(< 400) - break urlFallbackLoop - } - } - - if resp != nil && resp.StatusCode < 400 && usedBaseURL != "" { - antigravity.DefaultURLAvailability.MarkSuccess(usedBaseURL) - } - - // 成功响应时清零 INTERNAL 500 连续失败计数器(覆盖所有成功路径,含 smart retry) - if resp != nil && resp.StatusCode < 400 { - s.resetInternal500Counter(p.ctx, p.prefix, p.account.ID) - } - - return &antigravityRetryLoopResult{resp: resp}, nil -} - -// shouldRetryAntigravityError 判断是否应该重试 -func shouldRetryAntigravityError(statusCode int) bool { - switch statusCode { - case 429, 500, 502, 503, 504, 529: - return true - default: - return false - } -} - -// isURLLevelRateLimit 判断是否为 URL 级别的限流(应切换 URL 重试) -// "Resource has been exhausted" 是 URL/节点级别限流,切换 URL 可能成功 -// "exhausted your capacity on this model" 是账户/模型配额限流,切换 URL 无效 -func isURLLevelRateLimit(body []byte) bool { - // 快速检查:包含 "Resource has been exhausted" 且不包含 "capacity on this model" - bodyStr := string(body) - return strings.Contains(bodyStr, "Resource has been exhausted") && - !strings.Contains(bodyStr, "capacity on this model") -} - -// isAntigravityConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝) -func isAntigravityConnectionError(err error) bool { - if err == nil { - return false - } - - // 检查超时错误 - var netErr net.Error - if errors.As(err, &netErr) && netErr.Timeout() { - return true - } - - // 检查连接错误(DNS 失败、连接拒绝) - var opErr *net.OpError - return errors.As(err, &opErr) -} - -// shouldAntigravityFallbackToNextURL 判断是否应切换到下一个 URL -// 仅连接错误和 HTTP 429 触发 URL 降级 -func shouldAntigravityFallbackToNextURL(err error, statusCode int) bool { - if isAntigravityConnectionError(err) { - return true - } - return statusCode == http.StatusTooManyRequests -} - -// getSessionID 从 gin.Context 获取 session_id(用于日志追踪) -func getSessionID(c *gin.Context) string { - if c == nil { - return "" - } - return c.GetHeader("session_id") -} - -// logPrefix 生成统一的日志前缀 -func logPrefix(sessionID, accountName string) string { - if sessionID != "" { - return fmt.Sprintf("[antigravity-Forward] session=%s account=%s", sessionID, accountName) - } - return fmt.Sprintf("[antigravity-Forward] account=%s", accountName) -} - // AntigravityGatewayService 处理 Antigravity 平台的 API 转发 type AntigravityGatewayService struct { accountRepo AccountRepository @@ -1374,3281 +637,3 @@ func (s *AntigravityGatewayService) unwrapV1InternalResponse(body []byte) ([]byt } return body, nil } - -// Forward 转发 Claude 协议请求(Claude → Gemini 转换) -// -// 限流处理流程: -// -// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 -// ├─ 成功 → 正常返回 -// └─ 429/503 → handleSmartRetry -// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 -// └─ retryDelay < 7s → 等待后重试 1 次 -// ├─ 成功 → 正常返回 -// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 -func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte, isStickySession bool) (*ForwardResult, error) { - // 上游透传账号直接转发,不走 OAuth token 刷新 - if account.Type == AccountTypeUpstream { - return s.ForwardUpstream(ctx, c, account, body) - } - - startTime := time.Now() - - sessionID := getSessionID(c) - prefix := logPrefix(sessionID, account.Name) - - // 解析 Claude 请求 - var claudeReq antigravity.ClaudeRequest - if err := json.Unmarshal(body, &claudeReq); err != nil { - return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") - } - if strings.TrimSpace(claudeReq.Model) == "" { - return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Missing model") - } - - originalModel := claudeReq.Model - mappedModel := s.getMappedModel(account, claudeReq.Model) - if mappedModel == "" { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - return nil, s.writeClaudeError(c, http.StatusForbidden, "permission_error", fmt.Sprintf("model %s not in whitelist", claudeReq.Model)) - } - // 应用 thinking 模式自动后缀:如果 thinking 开启且目标是 claude-sonnet-4-5,自动改为 thinking 版本 - thinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive") - mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled) - billingModel := mappedModel - - // 获取 access_token - if s.tokenProvider == nil { - return nil, s.writeClaudeError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`), - } - } - - projectID, err := resolveAntigravityProjectID(account) - if err != nil { - _ = s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) - return nil, err - } - - // 代理 URL - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - // 获取转换选项 - // Antigravity 上游要求必须包含身份提示词,否则会返回 429 - transformOpts := s.getClaudeTransformOptions(ctx) - transformOpts.EnableIdentityPatch = true // 强制启用,Antigravity 上游必需 - - // 转换 Claude 请求为 Gemini 格式 - geminiBody, err := antigravity.TransformClaudeToGeminiWithOptions(&claudeReq, projectID, mappedModel, transformOpts) - if err != nil { - return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request") - } - - // Antigravity 上游只支持流式请求,统一使用 streamGenerateContent - // 如果客户端请求非流式,在响应处理阶段会收集完整流式响应后转换返回 - action := "streamGenerateContent" - - // 执行带重试的请求 - result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: action, - body: geminiBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, // Forward 由上层判断粘性会话 - groupID: 0, // Forward 方法没有 groupID,由上层处理粘性会话清除 - sessionHash: "", // Forward 方法没有 sessionHash,由上层处理粘性会话清除 - }) - if err != nil { - // 检查是否是账号切换信号,转换为 UpstreamFailoverError 让 Handler 切换账号 - if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusServiceUnavailable, - ForceCacheBilling: switchErr.IsStickySession, - } - } - // 区分客户端取消和真正的上游失败,返回更准确的错误消息 - if c.Request.Context().Err() != nil { - return nil, s.writeClaudeError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response") - } - return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries") - } - resp := result.resp - defer func() { _ = resp.Body.Close() }() - - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - - // 优先检测 thinking block 的 signature 相关错误(400)并重试一次: - // Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验, - // 当历史消息携带的 signature 不合法时会直接 400;去除 thinking 后可继续完成请求。 - if resp.StatusCode == http.StatusBadRequest && isSignatureRelatedError(respBody) && s.settingService.IsSignatureRectifierEnabled(ctx) { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - logBody, maxBytes := s.getLogConfig() - upstreamDetail := s.getUpstreamErrorDetail(respBody) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "signature_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - // Conservative two-stage fallback: - // 1) Disable top-level thinking + thinking->text - // 2) Only if still signature-related 400: also downgrade tool_use/tool_result to text. - - retryStages := []struct { - name string - strip func(*antigravity.ClaudeRequest) (bool, error) - }{ - {name: "thinking-only", strip: stripThinkingFromClaudeRequest}, - {name: "thinking+tools", strip: stripSignatureSensitiveBlocksFromClaudeRequest}, - } - - for _, stage := range retryStages { - retryClaudeReq := claudeReq - retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) - - stripped, stripErr := stage.strip(&retryClaudeReq) - if stripErr != nil || !stripped { - continue - } - - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected signature-related 400, retrying once (%s)", account.ID, stage.name) - - retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, s.getClaudeTransformOptions(ctx)) - if txErr != nil { - continue - } - retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: action, - body: retryGeminiBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, - groupID: 0, // Forward 方法没有 groupID,由上层处理粘性会话清除 - sessionHash: "", // Forward 方法没有 sessionHash,由上层处理粘性会话清除 - }) - if retryErr != nil { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - Kind: "signature_retry_request_error", - Message: sanitizeUpstreamErrorMessage(retryErr.Error()), - }) - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: signature retry request failed (%s): %v", account.ID, stage.name, retryErr) - continue - } - - retryResp := retryResult.resp - if retryResp.StatusCode < 400 { - _ = resp.Body.Close() - resp = retryResp - respBody = nil - break - } - - retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) - _ = retryResp.Body.Close() - if retryResp.StatusCode == http.StatusTooManyRequests { - retryBaseURL := "" - if retryResp.Request != nil && retryResp.Request.URL != nil { - retryBaseURL = retryResp.Request.URL.Scheme + "://" + retryResp.Request.URL.Host - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited base_url=%s retry_stage=%s body=%s", prefix, retryBaseURL, stage.name, truncateForLog(retryBody, 200)) - } - kind := "signature_retry" - if strings.TrimSpace(stage.name) != "" { - kind = "signature_retry_" + strings.ReplaceAll(stage.name, "+", "_") - } - retryUpstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(retryBody)) - retryUpstreamMsg = sanitizeUpstreamErrorMessage(retryUpstreamMsg) - retryUpstreamDetail := "" - if logBody { - retryUpstreamDetail = truncateString(string(retryBody), maxBytes) - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: retryResp.StatusCode, - UpstreamRequestID: retryResp.Header.Get("x-request-id"), - Kind: kind, - Message: retryUpstreamMsg, - Detail: retryUpstreamDetail, - }) - - // If this stage fixed the signature issue, we stop; otherwise we may try the next stage. - if retryResp.StatusCode != http.StatusBadRequest || !isSignatureRelatedError(retryBody) { - respBody = retryBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - } - break - } - - // Still signature-related; capture context and allow next stage. - respBody = retryBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - } - } - } - - // Budget 整流:检测 budget_tokens 约束错误并自动修正重试 - if resp.StatusCode == http.StatusBadRequest && respBody != nil && !isSignatureRelatedError(respBody) { - errMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - if isThinkingBudgetConstraintError(errMsg) && s.settingService.IsBudgetRectifierEnabled(ctx) { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "budget_constraint_error", - Message: errMsg, - Detail: s.getUpstreamErrorDetail(respBody), - }) - - // 修正 claudeReq 的 thinking 参数(adaptive 模式不修正) - if claudeReq.Thinking == nil || claudeReq.Thinking.Type != "adaptive" { - retryClaudeReq := claudeReq - retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) - // 创建新的 ThinkingConfig 避免修改原始 claudeReq.Thinking 指针 - retryClaudeReq.Thinking = &antigravity.ThinkingConfig{ - Type: "enabled", - BudgetTokens: BudgetRectifyBudgetTokens, - } - if retryClaudeReq.MaxTokens < BudgetRectifyMinMaxTokens { - retryClaudeReq.MaxTokens = BudgetRectifyMaxTokens - } - - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) - - retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, transformOpts) - if txErr == nil { - retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: action, - body: retryGeminiBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, - groupID: 0, - sessionHash: "", - }) - if retryErr == nil { - retryResp := retryResult.resp - if retryResp.StatusCode < 400 { - _ = resp.Body.Close() - resp = retryResp - respBody = nil - } else { - retryBody := s.readUpstreamErrorBody(retryResp) - _ = retryResp.Body.Close() - respBody = retryBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - } - } - } else { - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: budget rectifier retry failed: %v", account.ID, retryErr) - } - } - } - } - } - - // 处理错误响应(重试后仍失败或不触发重试) - if resp.StatusCode >= 400 { - // 检测 prompt too long 错误,返回特殊错误类型供上层 fallback - if resp.StatusCode == http.StatusBadRequest && isPromptTooLongError(respBody) { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.getUpstreamErrorDetail(respBody) - logBody, maxBytes := s.getLogConfig() - if logBody { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=400 prompt_too_long=true upstream_message=%q request_id=%s body=%s", prefix, upstreamMsg, resp.Header.Get("x-request-id"), truncateForLog(respBody, maxBytes)) - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "prompt_too_long", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &PromptTooLongError{ - StatusCode: resp.StatusCode, - RequestID: resp.Header.Get("x-request-id"), - Body: respBody, - } - } - - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession) - - // 精确匹配服务端配置类 400 错误,触发同账号重试 + failover - if resp.StatusCode == http.StatusBadRequest { - msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - if isGoogleProjectConfigError(msg) { - upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - upstreamDetail := s.getUpstreamErrorDetail(respBody) - log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: true} - } - } - - if s.shouldFailoverUpstreamError(resp.StatusCode) { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.getUpstreamErrorDetail(respBody) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} - } - - return nil, s.writeMappedClaudeError(c, account, resp.StatusCode, resp.Header.Get("x-request-id"), respBody) - } - } - - requestID := resp.Header.Get("x-request-id") - if requestID != "" { - c.Header("x-request-id", requestID) - } - - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - if claudeReq.Stream { - // 客户端要求流式,直接透传转换 - streamRes, err := s.handleClaudeStreamingResponse(c, resp, startTime, originalModel) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - clientDisconnect = streamRes.clientDisconnect - } else { - // 客户端要求非流式,收集流式响应后转换返回 - streamRes, err := s.handleClaudeStreamToNonStreaming(c, resp, startTime, originalModel) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - } - - return &ForwardResult{ - RequestID: requestID, - Usage: *usage, - Model: originalModel, - UpstreamModel: billingModel, - Stream: claudeReq.Stream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - }, nil -} - -func isSignatureRelatedError(respBody []byte) bool { - msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - if msg == "" { - // Fallback: best-effort scan of the raw payload. - msg = strings.ToLower(string(respBody)) - } - - // Keep this intentionally broad: different upstreams may use "signature" or "thought_signature". - if strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature") { - return true - } - - // Also detect thinking block structural errors: - // "Expected `thinking` or `redacted_thinking`, but found `text`" - if strings.Contains(msg, "expected") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { - return true - } - - return false -} - -// isPromptTooLongError 检测是否为 prompt too long 错误 -func isPromptTooLongError(respBody []byte) bool { - msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - if msg == "" { - msg = strings.ToLower(string(respBody)) - } - return strings.Contains(msg, "prompt is too long") || - strings.Contains(msg, "request is too long") || - strings.Contains(msg, "context length exceeded") || - strings.Contains(msg, "max_tokens") -} - -// isPassthroughErrorMessage 检查错误消息是否在透传白名单中 -func isPassthroughErrorMessage(msg string) bool { - lower := strings.ToLower(msg) - for _, pattern := range antigravityPassthroughErrorMessages { - if strings.Contains(lower, pattern) { - return true - } - } - return false -} - -// getPassthroughOrDefault 若消息在白名单内则返回原始消息,否则返回默认消息 -func getPassthroughOrDefault(upstreamMsg, defaultMsg string) string { - if isPassthroughErrorMessage(upstreamMsg) { - return upstreamMsg - } - return defaultMsg -} - -func extractAntigravityErrorMessage(body []byte) string { - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return "" - } - - // Google-style: {"error": {"message": "..."}} - if errObj, ok := payload["error"].(map[string]any); ok { - if msg, ok := errObj["message"].(string); ok && strings.TrimSpace(msg) != "" { - return msg - } - } - - // Fallback: top-level message - if msg, ok := payload["message"].(string); ok && strings.TrimSpace(msg) != "" { - return msg - } - - return "" -} - -// stripThinkingFromClaudeRequest converts thinking blocks to text blocks in a Claude Messages request. -// This preserves the thinking content while avoiding signature validation errors. -// Note: redacted_thinking blocks are removed because they cannot be converted to text. -// It also disables top-level `thinking` to avoid upstream structural constraints for thinking mode. -func stripThinkingFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { - if req == nil { - return false, nil - } - - changed := false - if req.Thinking != nil { - req.Thinking = nil - changed = true - } - - for i := range req.Messages { - raw := req.Messages[i].Content - if len(raw) == 0 { - continue - } - - // If content is a string, nothing to strip. - var str string - if json.Unmarshal(raw, &str) == nil { - continue - } - - // Otherwise treat as an array of blocks and convert thinking blocks to text. - var blocks []map[string]any - if err := json.Unmarshal(raw, &blocks); err != nil { - continue - } - - filtered := make([]map[string]any, 0, len(blocks)) - modifiedAny := false - for _, block := range blocks { - t, _ := block["type"].(string) - switch t { - case "thinking": - thinkingText, _ := block["thinking"].(string) - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - case "redacted_thinking": - modifiedAny = true - case "": - if thinkingText, hasThinking := block["thinking"].(string); hasThinking { - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - } else { - filtered = append(filtered, block) - } - default: - filtered = append(filtered, block) - } - } - - if !modifiedAny { - continue - } - - if len(filtered) == 0 { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": "(content removed)", - }) - } - - newRaw, err := json.Marshal(filtered) - if err != nil { - return changed, err - } - req.Messages[i].Content = newRaw - changed = true - } - - return changed, nil -} - -// stripSignatureSensitiveBlocksFromClaudeRequest is a stronger retry degradation that additionally converts -// tool blocks to plain text. Use this only after a thinking-only retry still fails with signature errors. -func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { - if req == nil { - return false, nil - } - - changed := false - if req.Thinking != nil { - req.Thinking = nil - changed = true - } - - for i := range req.Messages { - raw := req.Messages[i].Content - if len(raw) == 0 { - continue - } - - // If content is a string, nothing to strip. - var str string - if json.Unmarshal(raw, &str) == nil { - continue - } - - // Otherwise treat as an array of blocks and convert signature-sensitive blocks to text. - var blocks []map[string]any - if err := json.Unmarshal(raw, &blocks); err != nil { - continue - } - - filtered := make([]map[string]any, 0, len(blocks)) - modifiedAny := false - for _, block := range blocks { - t, _ := block["type"].(string) - switch t { - case "thinking": - // Convert thinking to text, skip if empty - thinkingText, _ := block["thinking"].(string) - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - case "redacted_thinking": - // Remove redacted_thinking (cannot convert encrypted content) - modifiedAny = true - case "tool_use": - // Convert tool_use to text to avoid upstream signature/thought_signature validation errors. - // This is a retry-only degradation path, so we prioritise request validity over tool semantics. - name, _ := block["name"].(string) - id, _ := block["id"].(string) - input := block["input"] - inputJSON, _ := json.Marshal(input) - text := "(tool_use)" - if name != "" { - text += " name=" + name - } - if id != "" { - text += " id=" + id - } - if len(inputJSON) > 0 && string(inputJSON) != "null" { - text += " input=" + string(inputJSON) - } - filtered = append(filtered, map[string]any{ - "type": "text", - "text": text, - }) - modifiedAny = true - case "tool_result": - // Convert tool_result to text so it stays consistent when tool_use is downgraded. - toolUseID, _ := block["tool_use_id"].(string) - isError, _ := block["is_error"].(bool) - content := block["content"] - contentJSON, _ := json.Marshal(content) - text := "(tool_result)" - if toolUseID != "" { - text += " tool_use_id=" + toolUseID - } - if isError { - text += " is_error=true" - } - if len(contentJSON) > 0 && string(contentJSON) != "null" { - text += "\n" + string(contentJSON) - } - filtered = append(filtered, map[string]any{ - "type": "text", - "text": text, - }) - modifiedAny = true - case "": - // Handle untyped block with "thinking" field - if thinkingText, hasThinking := block["thinking"].(string); hasThinking { - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - } else { - filtered = append(filtered, block) - } - default: - filtered = append(filtered, block) - } - } - - if !modifiedAny { - continue - } - - if len(filtered) == 0 { - // Keep request valid: upstream rejects empty content arrays. - filtered = append(filtered, map[string]any{ - "type": "text", - "text": "(content removed)", - }) - } - - newRaw, err := json.Marshal(filtered) - if err != nil { - return changed, err - } - req.Messages[i].Content = newRaw - changed = true - } - - return changed, nil -} - -// ForwardGemini 转发 Gemini 协议请求 -// -// 限流处理流程: -// -// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 -// ├─ 成功 → 正常返回 -// └─ 429/503 → handleSmartRetry -// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 -// └─ retryDelay < 7s → 等待后重试 1 次 -// ├─ 成功 → 正常返回 -// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 -type ForwardGeminiOption func(*forwardGeminiOptions) - -type forwardGeminiOptions struct { - groupID int64 - sessionHash string -} - -func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption { - return func(opts *forwardGeminiOptions) { - opts.groupID = groupID - opts.sessionHash = sessionHash - } -} - -func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) { - startTime := time.Now() - forwardOpts := forwardGeminiOptions{} - for _, apply := range options { - if apply != nil { - apply(&forwardOpts) - } - } - - sessionID := getSessionID(c) - prefix := logPrefix(sessionID, account.Name) - - if strings.TrimSpace(originalModel) == "" { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing model in URL") - } - if strings.TrimSpace(action) == "" { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing action in URL") - } - if len(body) == 0 { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty") - } - - // 解析请求以获取 image_size(用于图片计费) - imageInputSize := s.extractImageInputSize(body) - imageSize := normalizeOpenAIImageSizeTier(imageInputSize) - - switch action { - case "generateContent", "streamGenerateContent": - // ok - case "countTokens": - // 直接返回空值,不透传上游 - c.JSON(http.StatusOK, map[string]any{"totalTokens": 0}) - return &ForwardResult{ - RequestID: "", - Usage: ClaudeUsage{}, - Model: originalModel, - Stream: false, - Duration: time.Since(startTime), - FirstTokenMs: nil, - }, nil - default: - return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action) - } - - mappedModel := s.getMappedModel(account, originalModel) - if mappedModel == "" { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - return nil, s.writeGoogleError(c, http.StatusForbidden, fmt.Sprintf("model %s not in whitelist", originalModel)) - } - billingModel := mappedModel - - // 获取 access_token - if s.tokenProvider == nil { - return nil, s.writeGoogleError(c, http.StatusBadGateway, "Antigravity token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":{"message":"Failed to get upstream access token","status":"UNAVAILABLE"}}`), - } - } - - projectID, err := resolveAntigravityProjectID(account) - if err != nil { - _ = s.writeGoogleError(c, http.StatusBadRequest, err.Error()) - return nil, err - } - - // 代理 URL - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - // Antigravity 上游要求必须包含身份提示词,注入到请求中 - injectedBody, err := injectIdentityPatchToGeminiRequest(body) - if err != nil { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Invalid request body") - } - - // 清理 Schema - if cleanedBody, err := cleanGeminiRequest(injectedBody); err == nil { - injectedBody = cleanedBody - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Cleaned request schema in forwarded request for account %s", account.Name) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Failed to clean schema: %v", err) - } - - // 包装请求 - wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody) - if err != nil { - return nil, s.writeGoogleError(c, http.StatusInternalServerError, "Failed to build upstream request") - } - - // Antigravity 上游只支持流式请求,统一使用 streamGenerateContent - // 如果客户端请求非流式,在响应处理阶段会收集完整流式响应后返回 - upstreamAction := "streamGenerateContent" - - // 执行带重试的请求 - result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: upstreamAction, - body: wrappedBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话 - groupID: forwardOpts.groupID, - sessionHash: forwardOpts.sessionHash, - }) - if err != nil { - // 检查是否是账号切换信号,转换为 UpstreamFailoverError 让 Handler 切换账号 - if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusServiceUnavailable, - ForceCacheBilling: switchErr.IsStickySession, - } - } - // 区分客户端取消和真正的上游失败,返回更准确的错误消息 - if c.Request.Context().Err() != nil { - return nil, s.writeGoogleError(c, http.StatusBadGateway, "Client disconnected before upstream response") - } - return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries") - } - resp := result.resp - defer func() { - if resp != nil && resp.Body != nil { - _ = resp.Body.Close() - } - }() - - // 处理错误响应 - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - contentType := resp.Header.Get("Content-Type") - // 尽早关闭原始响应体,释放连接;后续逻辑仍可能需要读取 body,因此用内存副本重新包装。 - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - - // 模型兜底:模型不存在且开启 fallback 时,自动用 fallback 模型重试一次 - if s.settingService != nil && s.settingService.IsModelFallbackEnabled(ctx) && - isModelNotFoundError(resp.StatusCode, respBody) { - fallbackModel := s.settingService.GetFallbackModel(ctx, PlatformAntigravity) - if fallbackModel != "" && fallbackModel != mappedModel { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Model not found (%s), retrying with fallback model %s (account: %s)", mappedModel, fallbackModel, account.Name) - - fallbackWrapped, err := s.wrapV1InternalRequest(projectID, fallbackModel, injectedBody) - if err == nil { - fallbackReq, err := antigravity.NewAPIRequest(ctx, upstreamAction, accessToken, fallbackWrapped) - if err == nil { - fallbackResp, err := s.httpUpstream.Do(fallbackReq, proxyURL, account.ID, account.Concurrency) - if err == nil && fallbackResp.StatusCode < 400 { - _ = resp.Body.Close() - resp = fallbackResp - } else if fallbackResp != nil { - _ = fallbackResp.Body.Close() - } - } - } - } - } - - // Gemini 原生请求中的 thoughtSignature 可能来自旧上下文/旧账号,触发上游严格校验后返回 - // "Corrupted thought signature."。检测到此类 400 时,将 thoughtSignature 清理为 dummy 值后重试一次。 - signatureCheckBody := respBody - if unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody); unwrapErr == nil && len(unwrapped) > 0 { - signatureCheckBody = unwrapped - } - if resp.StatusCode == http.StatusBadRequest && - s.settingService != nil && - s.settingService.IsSignatureRectifierEnabled(ctx) && - isSignatureRelatedError(signatureCheckBody) && - bytes.Contains(injectedBody, []byte(`"thoughtSignature"`)) { - upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(signatureCheckBody))) - upstreamDetail := s.getUpstreamErrorDetail(signatureCheckBody) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "signature_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: detected signature-related 400, retrying with cleaned thought signatures", account.ID) - - cleanedInjectedBody := CleanGeminiNativeThoughtSignatures(injectedBody) - retryWrappedBody, wrapErr := s.wrapV1InternalRequest(projectID, mappedModel, cleanedInjectedBody) - if wrapErr == nil { - retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: upstreamAction, - body: retryWrappedBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, - groupID: forwardOpts.groupID, - sessionHash: forwardOpts.sessionHash, - }) - if retryErr == nil { - retryResp := retryResult.resp - if retryResp.StatusCode < 400 { - resp = retryResp - } else { - retryRespBody := s.readUpstreamErrorBody(retryResp) - _ = retryResp.Body.Close() - retryOpsBody := retryRespBody - if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 { - retryOpsBody = retryUnwrapped - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: retryResp.StatusCode, - UpstreamRequestID: retryResp.Header.Get("x-request-id"), - Kind: "signature_retry", - Message: sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(retryOpsBody))), - Detail: s.getUpstreamErrorDetail(retryOpsBody), - }) - respBody = retryRespBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryRespBody)), - } - contentType = resp.Header.Get("Content-Type") - } - } else { - if switchErr, ok := IsAntigravityAccountSwitchError(retryErr); ok { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: http.StatusServiceUnavailable, - Kind: "failover", - Message: sanitizeUpstreamErrorMessage(retryErr.Error()), - }) - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusServiceUnavailable, - ForceCacheBilling: switchErr.IsStickySession, - } - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - Kind: "signature_retry_request_error", - Message: sanitizeUpstreamErrorMessage(retryErr.Error()), - }) - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry request failed: %v", account.ID, retryErr) - } - } else { - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry wrap failed: %v", account.ID, wrapErr) - } - } - - // fallback 成功:继续按正常响应处理 - if resp.StatusCode < 400 { - goto handleSuccess - } - - requestID := resp.Header.Get("x-request-id") - if requestID != "" { - c.Header("x-request-id", requestID) - } - - unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody) - unwrappedForOps := unwrapped - if unwrapErr != nil || len(unwrappedForOps) == 0 { - unwrappedForOps = respBody - } - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession) - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps) - - // Always record upstream context for Ops error logs, even when we will failover. - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - - // 精确匹配服务端配置类 400 错误,触发同账号重试 + failover - if resp.StatusCode == http.StatusBadRequest && isGoogleProjectConfigError(strings.ToLower(upstreamMsg)) { - log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps, RetryableOnSameAccount: true} - } - - if s.shouldFailoverUpstreamError(resp.StatusCode) { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps} - } - if contentType == "" { - contentType = "application/json" - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream error status=%d body=%s", resp.StatusCode, truncateForLog(unwrappedForOps, 500)) - MarkResponseCommitted(c) - c.Data(resp.StatusCode, contentType, unwrappedForOps) - return nil, fmt.Errorf("antigravity upstream error: %d", resp.StatusCode) - } - -handleSuccess: - requestID := resp.Header.Get("x-request-id") - if requestID != "" { - c.Header("x-request-id", requestID) - } - - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - - if stream { - // 客户端要求流式,直接透传 - streamRes, err := s.handleGeminiStreamingResponse(c, resp, startTime) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - clientDisconnect = streamRes.clientDisconnect - } else { - // 客户端要求非流式,收集流式响应后返回 - streamRes, err := s.handleGeminiStreamToNonStreaming(c, resp, startTime) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - } - - if usage == nil { - usage = &ClaudeUsage{} - } - - // 判断是否为图片生成模型 - imageCount := 0 - if isImageGenerationModel(mappedModel) { - // Gemini 图片生成 API 每次请求只生成一张图片(API 限制) - imageCount = 1 - } - - return &ForwardResult{ - RequestID: requestID, - Usage: *usage, - Model: originalModel, - UpstreamModel: billingModel, - Stream: stream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - ImageCount: imageCount, - ImageSize: imageSize, - ImageInputSize: imageInputSize, - }, nil -} - -func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int) bool { - switch statusCode { - case 401, 403, 429, 529: - return true - default: - return statusCode >= 500 - } -} - -// isGoogleProjectConfigError 判断(已提取的小写)错误消息是否属于 Google 服务端配置类问题。 -// 只精确匹配已知的服务端侧错误,避免对客户端请求错误做无意义重试。 -// 适用于所有走 Google 后端的平台(Antigravity、Gemini)。 -func isGoogleProjectConfigError(lowerMsg string) bool { - // Google 间歇性 Bug:Project ID 有效但被临时识别失败 - return strings.Contains(lowerMsg, "invalid project resource name") -} - -// googleConfigErrorCooldown 服务端配置类 400 错误的临时封禁时长 -const googleConfigErrorCooldown = 1 * time.Minute - -// tempUnscheduleGoogleConfigError 对服务端配置类 400 错误触发临时封禁, -// 避免短时间内反复调度到同一个有问题的账号。 -func tempUnscheduleGoogleConfigError(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { - until := time.Now().Add(googleConfigErrorCooldown) - reason := "400: invalid project resource name (auto temp-unschedule 1m)" - if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { - log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) - } else { - log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) - } -} - -// emptyResponseCooldown 空流式响应的临时封禁时长 -const emptyResponseCooldown = 1 * time.Minute - -// tempUnscheduleEmptyResponse 对空流式响应触发临时封禁, -// 避免短时间内反复调度到同一个返回空响应的账号。 -func tempUnscheduleEmptyResponse(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { - until := time.Now().Add(emptyResponseCooldown) - reason := "empty stream response (auto temp-unschedule 1m)" - if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { - log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) - } else { - log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) - } -} - -// sleepAntigravityBackoffWithContext 带 context 取消检查的退避等待 -// 返回 true 表示正常完成等待,false 表示 context 已取消 -func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool { - delay := antigravityRetryBaseDelay * time.Duration(1< antigravityRetryMaxDelay { - delay = antigravityRetryMaxDelay - } - - // +/- 20% jitter - r := mathrand.New(mathrand.NewSource(time.Now().UnixNano())) - jitter := time.Duration(float64(delay) * 0.2 * (r.Float64()*2 - 1)) - sleepFor := delay + jitter - if sleepFor < 0 { - sleepFor = 0 - } - - timer := time.NewTimer(sleepFor) - select { - case <-ctx.Done(): - timer.Stop() - return false - case <-timer.C: - return true - } -} - -// isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记 -func isSingleAccountRetry(ctx context.Context) bool { - v, _ := SingleAccountRetryFromContext(ctx) - return v -} - -// setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流 -// 直接使用上游返回的模型 ID(如 claude-sonnet-4-5)作为限流 key -// 返回是否已成功设置(若模型名为空或 repo 为 nil 将返回 false) -func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, accountID int64, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { - if repo == nil || modelName == "" { - return false - } - // 直接使用官方模型 ID 作为 key,不再转换为 scope - if err := repo.SetModelRateLimit(ctx, accountID, modelName, resetAt); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", prefix, statusCode, modelName, err) - return false - } - if afterSmartRetry { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) - } - return true -} - -func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { - if account == nil || repo == nil { - return false - } - keys := antigravityModelRateLimitKeys(modelName) - if len(keys) == 0 { - return false - } - - success := false - for _, key := range keys { - if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) { - s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt) - success = true - } - } - return success -} - -func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { - if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { - return - } - if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err) - } -} - -func antigravityFallbackCooldownSeconds() (time.Duration, bool) { - raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv)) - if raw == "" { - return 0, false - } - seconds, err := strconv.Atoi(raw) - if err != nil || seconds <= 0 { - return 0, false - } - return time.Duration(seconds) * time.Second, true -} - -// antigravitySmartRetryInfo 智能重试所需的信息 -type antigravitySmartRetryInfo struct { - RetryDelay time.Duration // 重试延迟时间 - ModelName string // 限流的模型名称(如 "claude-sonnet-4-5") - IsModelCapacityExhausted bool // 是否为模型容量不足(MODEL_CAPACITY_EXHAUSTED) -} - -// parseAntigravitySmartRetryInfo 解析 Google RPC RetryInfo 和 ErrorInfo 信息 -// 返回解析结果,如果解析失败或不满足条件返回 nil -// -// 支持两种情况: -// 1. 429 RESOURCE_EXHAUSTED + RATE_LIMIT_EXCEEDED: -// - error.status == "RESOURCE_EXHAUSTED" -// - error.details[].reason == "RATE_LIMIT_EXCEEDED" -// -// 2. 503 UNAVAILABLE + MODEL_CAPACITY_EXHAUSTED: -// - error.status == "UNAVAILABLE" -// - error.details[].reason == "MODEL_CAPACITY_EXHAUSTED" -// -// 必须满足以下条件才会返回有效值: -// - error.details[] 中存在 @type == "type.googleapis.com/google.rpc.RetryInfo" 的元素 -// - 该元素包含 retryDelay 字段,格式为 "数字s"(如 "0.201506475s") -func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { - var parsed map[string]any - if err := json.Unmarshal(body, &parsed); err != nil { - return nil - } - - errObj, ok := parsed["error"].(map[string]any) - if !ok { - return nil - } - - // 检查 status 是否符合条件 - // 情况1: 429 RESOURCE_EXHAUSTED (需要进一步检查 reason == RATE_LIMIT_EXCEEDED) - // 情况2: 503 UNAVAILABLE (需要进一步检查 reason == MODEL_CAPACITY_EXHAUSTED) - status, _ := errObj["status"].(string) - isResourceExhausted := status == googleRPCStatusResourceExhausted - isUnavailable := status == googleRPCStatusUnavailable - - if !isResourceExhausted && !isUnavailable { - return nil - } - - details, ok := errObj["details"].([]any) - if !ok { - return nil - } - - var retryDelay time.Duration - var modelName string - var hasRateLimitExceeded bool // 429 需要此 reason - var hasModelCapacityExhausted bool // 503 需要此 reason - - for _, d := range details { - dm, ok := d.(map[string]any) - if !ok { - continue - } - - atType, _ := dm["@type"].(string) - - // 从 ErrorInfo 提取模型名称和 reason - if atType == googleRPCTypeErrorInfo { - if meta, ok := dm["metadata"].(map[string]any); ok { - if model, ok := meta["model"].(string); ok { - modelName = normalizeAntigravityModelName(model) - } - } - // 检查 reason - if reason, ok := dm["reason"].(string); ok { - if reason == googleRPCReasonModelCapacityExhausted { - hasModelCapacityExhausted = true - } - if reason == googleRPCReasonRateLimitExceeded { - hasRateLimitExceeded = true - } - } - continue - } - - // 从 RetryInfo 提取重试延迟 - if atType == googleRPCTypeRetryInfo { - delay, ok := dm["retryDelay"].(string) - if !ok || delay == "" { - continue - } - // 使用 time.ParseDuration 解析,支持所有 Go duration 格式 - // 例如: "0.5s", "10s", "4m50s", "1h30m", "200ms" 等 - dur, err := time.ParseDuration(delay) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] failed to parse retryDelay: %s error=%v", delay, err) - continue - } - retryDelay = dur - } - } - - // 验证条件 - // 情况1: RESOURCE_EXHAUSTED 需要有 RATE_LIMIT_EXCEEDED reason - // 情况2: UNAVAILABLE 需要有 MODEL_CAPACITY_EXHAUSTED reason - if isResourceExhausted && !hasRateLimitExceeded { - return nil - } - if isUnavailable && !hasModelCapacityExhausted { - return nil - } - - // 必须有模型名才返回有效结果 - if modelName == "" { - return nil - } - - // 如果上游未提供 retryDelay,使用默认限流时间 - if retryDelay <= 0 { - retryDelay = antigravityDefaultRateLimitDuration - } - - return &antigravitySmartRetryInfo{ - RetryDelay: retryDelay, - ModelName: modelName, - IsModelCapacityExhausted: hasModelCapacityExhausted, - } -} - -// shouldTriggerAntigravitySmartRetry 判断是否应该触发智能重试 -// 返回: -// - shouldRetry: 是否应该智能重试(retryDelay < antigravityRateLimitThreshold,或 MODEL_CAPACITY_EXHAUSTED) -// - shouldRateLimitModel: 是否应该限流模型并切换账号(仅 RATE_LIMIT_EXCEEDED 且 retryDelay >= 阈值) -// - waitDuration: 等待时间 -// - modelName: 限流的模型名称 -// - isModelCapacityExhausted: 是否为模型容量不足(MODEL_CAPACITY_EXHAUSTED) -func shouldTriggerAntigravitySmartRetry(account *Account, respBody []byte) (shouldRetry bool, shouldRateLimitModel bool, waitDuration time.Duration, modelName string, isModelCapacityExhausted bool) { - if account.Platform != PlatformAntigravity { - return false, false, 0, "", false - } - - info := parseAntigravitySmartRetryInfo(respBody) - if info == nil { - return false, false, 0, "", false - } - - // MODEL_CAPACITY_EXHAUSTED(模型容量不足):所有账号共享同一模型容量池 - // 切换账号无意义,使用固定 1s 间隔重试 - if info.IsModelCapacityExhausted { - return true, false, antigravityModelCapacityRetryWait, info.ModelName, true - } - - // RATE_LIMIT_EXCEEDED(账号级限流): - // retryDelay >= 阈值:直接限流模型,不重试 - // 注意:如果上游未提供 retryDelay,parseAntigravitySmartRetryInfo 已设置为默认 30s - if info.RetryDelay >= antigravityRateLimitThreshold { - return false, true, info.RetryDelay, info.ModelName, false - } - - // retryDelay < 阈值:智能重试 - waitDuration = info.RetryDelay - if waitDuration < antigravitySmartRetryMinWait { - waitDuration = antigravitySmartRetryMinWait - } - - return true, false, waitDuration, info.ModelName, false -} - -// handleModelRateLimitParams 模型级限流处理参数 -type handleModelRateLimitParams struct { - ctx context.Context - prefix string - account *Account - statusCode int - body []byte - cache GatewayCache - groupID int64 - sessionHash string - isStickySession bool -} - -// handleModelRateLimitResult 模型级限流处理结果 -type handleModelRateLimitResult struct { - Handled bool // 是否已处理 - ShouldRetry bool // 是否等待后重试 - WaitDuration time.Duration // 等待时间 - SwitchError *AntigravityAccountSwitchError // 账号切换错误 -} - -// handleModelRateLimit 处理模型级限流(在原有逻辑之前调用) -// 仅处理 429/503,解析模型名和 retryDelay -// - MODEL_CAPACITY_EXHAUSTED: 返回 Handled=true(实际重试由 handleSmartRetry 处理) -// - RATE_LIMIT_EXCEEDED + retryDelay < 阈值: 返回 ShouldRetry=true,由调用方等待后重试 -// - RATE_LIMIT_EXCEEDED + retryDelay >= 阈值: 设置模型限流 + 清除粘性会话 + 返回 SwitchError -func (s *AntigravityGatewayService) handleModelRateLimit(p *handleModelRateLimitParams) *handleModelRateLimitResult { - if p.statusCode != 429 && p.statusCode != 503 { - return &handleModelRateLimitResult{Handled: false} - } - - info := parseAntigravitySmartRetryInfo(p.body) - if info == nil || info.ModelName == "" { - return &handleModelRateLimitResult{Handled: false} - } - - // MODEL_CAPACITY_EXHAUSTED:模型容量不足,所有账号共享同一容量池 - // 切换账号无意义,不设置模型限流(实际重试由 handleSmartRetry 处理) - if info.IsModelCapacityExhausted { - log.Printf("%s status=%d model_capacity_exhausted model=%s (not switching account, retry handled by smart retry)", - p.prefix, p.statusCode, info.ModelName) - return &handleModelRateLimitResult{ - Handled: true, - } - } - - // RATE_LIMIT_EXCEEDED: < antigravityRateLimitThreshold: 等待后重试 - if info.RetryDelay < antigravityRateLimitThreshold { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_wait model=%s wait=%v", - p.prefix, p.statusCode, info.ModelName, info.RetryDelay) - return &handleModelRateLimitResult{ - Handled: true, - ShouldRetry: true, - WaitDuration: info.RetryDelay, - } - } - - // RATE_LIMIT_EXCEEDED: >= antigravityRateLimitThreshold: 设置限流 + 清除粘性会话 + 切换账号 - s.setModelRateLimitAndClearSession(p, info) - - return &handleModelRateLimitResult{ - Handled: true, - SwitchError: &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: info.ModelName, - IsStickySession: p.isStickySession, - }, - } -} - -// setModelRateLimitAndClearSession 设置模型限流并清除粘性会话 -func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleModelRateLimitParams, info *antigravitySmartRetryInfo) { - resetAt := time.Now().Add(info.RetryDelay) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", - p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) - - s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false) - - // 清除粘性会话绑定 - if p.cache != nil && p.sessionHash != "" { - _ = p.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash) - } -} - -// updateAccountModelRateLimitInCache 立即更新 Redis 中账号的模型限流状态 -func (s *AntigravityGatewayService) updateAccountModelRateLimitInCache(ctx context.Context, account *Account, modelKey string, resetAt time.Time) { - if s.schedulerSnapshot == nil || account == nil || modelKey == "" { - return - } - - // 更新账号对象的 Extra 字段 - if account.Extra == nil { - account.Extra = make(map[string]any) - } - - limits, _ := account.Extra["model_rate_limits"].(map[string]any) - if limits == nil { - limits = make(map[string]any) - account.Extra["model_rate_limits"] = limits - } - - limits[modelKey] = map[string]any{ - "rate_limited_at": time.Now().UTC().Format(time.RFC3339), - "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339), - } - - // 更新 Redis 快照 - if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] cache_update_failed account=%d model=%s err=%v", account.ID, modelKey, err) - } -} - -func (s *AntigravityGatewayService) handleUpstreamError( - ctx context.Context, prefix string, account *Account, - statusCode int, headers http.Header, body []byte, - requestedModel string, - groupID int64, sessionHash string, isStickySession bool, -) *handleModelRateLimitResult { - // 遵守自定义错误码策略:未命中则跳过所有限流处理 - if !account.ShouldHandleErrorCode(statusCode) { - return nil - } - // 模型级限流处理(优先) - result := s.handleModelRateLimit(&handleModelRateLimitParams{ - ctx: ctx, - prefix: prefix, - account: account, - statusCode: statusCode, - body: body, - cache: s.cache, - groupID: groupID, - sessionHash: sessionHash, - isStickySession: isStickySession, - }) - if result.Handled { - return result - } - - // 503 仅处理模型限流(MODEL_CAPACITY_EXHAUSTED),非模型限流不做额外处理 - // 避免将普通的 503 错误误判为账号问题 - if statusCode == 503 { - return nil - } - - // 429:尝试解析模型级限流,解析失败时兜底为账号级限流 - if statusCode == 429 { - if logBody, maxBytes := s.getLogConfig(); logBody { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity-Debug] 429 response body: %s", truncateString(string(body), maxBytes)) - } - - resetAt := ParseGeminiRateLimitResetTime(body) - defaultDur := s.getDefaultRateLimitDuration() - - // 尝试解析模型 key 并设置模型级限流 - // - // 注意:requestedModel 可能是"映射前"的请求模型名(例如 claude-opus-4-6), - // 调度与限流判定使用的是 Antigravity 最终模型名(包含映射与 thinking 后缀)。 - // 因此这里必须写入最终模型 key,确保后续调度能正确避开已限流模型。 - modelKey := resolveFinalAntigravityModelKey(ctx, account, requestedModel) - if strings.TrimSpace(modelKey) == "" { - // 极少数情况下无法映射(理论上不应发生:能转发成功说明映射已通过), - // 保持旧行为作为兜底,避免完全丢失模型级限流记录。 - modelKey = resolveAntigravityModelKey(requestedModel) - } - if modelKey != "" { - ra := s.resolveResetTime(resetAt, defaultDur) - if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v", - prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) - } - return nil - } - - // 无法解析模型 key,兜底为账号级限流 - ra := s.resolveResetTime(resetAt, defaultDur) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited account=%d reset_at=%v reset_in=%v (fallback)", - prefix, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) - if err := s.accountRepo.SetRateLimited(ctx, account.ID, ra); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limit_set_failed account=%d error=%v", prefix, account.ID, err) - } - return nil - } - // 其他错误码继续使用 rateLimitService - if s.rateLimitService == nil { - return nil - } - shouldDisable := s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body) - if shouldDisable { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d marked_error", prefix, statusCode) - } - return nil -} - -// getDefaultRateLimitDuration 获取默认限流时间 -func (s *AntigravityGatewayService) getDefaultRateLimitDuration() time.Duration { - defaultDur := antigravityDefaultRateLimitDuration - if s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes > 0 { - defaultDur = time.Duration(s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes) * time.Minute - } - if override, ok := antigravityFallbackCooldownSeconds(); ok { - defaultDur = override - } - return defaultDur -} - -// resolveResetTime 根据解析的重置时间或默认时长计算重置时间点 -func (s *AntigravityGatewayService) resolveResetTime(resetAt *int64, defaultDur time.Duration) time.Time { - if resetAt != nil { - return time.Unix(*resetAt, 0) - } - return time.Now().Add(defaultDur) -} - -type antigravityStreamResult struct { - usage *ClaudeUsage - firstTokenMs *int - clientDisconnect bool // 客户端是否在流式传输过程中断开 -} - -// antigravityClientWriter 封装流式响应的客户端写入,自动检测断开并标记。 -// 断开后所有写入操作变为 no-op,调用方通过 Disconnected() 判断是否继续 drain 上游。 -type antigravityClientWriter struct { - w gin.ResponseWriter - flusher http.Flusher - disconnected bool - prefix string // 日志前缀,标识来源方法 -} - -func newAntigravityClientWriter(w gin.ResponseWriter, flusher http.Flusher, prefix string) *antigravityClientWriter { - return &antigravityClientWriter{w: w, flusher: flusher, prefix: prefix} -} - -// Write 写入数据到客户端,写入失败时标记断开并返回 false -func (cw *antigravityClientWriter) Write(p []byte) bool { - if cw.disconnected { - return false - } - if _, err := cw.w.Write(p); err != nil { - cw.markDisconnected() - return false - } - cw.flusher.Flush() - return true -} - -// Fprintf 格式化写入数据到客户端,写入失败时标记断开并返回 false -func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool { - if cw.disconnected { - return false - } - if _, err := fmt.Fprintf(cw.w, format, args...); err != nil { - cw.markDisconnected() - return false - } - cw.flusher.Flush() - return true -} - -func (cw *antigravityClientWriter) Disconnected() bool { return cw.disconnected } - -func (cw *antigravityClientWriter) markDisconnected() { - cw.disconnected = true - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during streaming (%s), continuing to drain upstream for billing", cw.prefix) -} - -// handleStreamReadError 处理上游读取错误的通用逻辑。 -// 返回 (clientDisconnect, handled):handled=true 表示错误已处理,调用方应返回已收集的 usage。 -func handleStreamReadError(err error, clientDisconnected bool, prefix string) (disconnect bool, handled bool) { - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - logger.LegacyPrintf("service.antigravity_gateway", "Context canceled during streaming (%s), returning collected usage", prefix) - return true, true - } - if clientDisconnected { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream read error after client disconnect (%s): %v, returning collected usage", prefix, err) - return true, true - } - return false, false -} - -func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { - c.Status(resp.StatusCode) - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - - contentType := resp.Header.Get("Content-Type") - if contentType == "" { - contentType = "text/event-stream; charset=utf-8" - } - c.Header("Content-Type", contentType) - - flusher, ok := c.Writer.(http.Flusher) - if !ok { - return nil, errors.New("streaming not supported") - } - - // 使用 Scanner 并限制单行大小,避免 ReadString 无上限导致 OOM - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - usage := &ClaudeUsage{} - var firstTokenMs *int - - type scanEvent struct { - line string - err error - } - // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - // 上游数据间隔超时保护(防止上游挂起长期占用连接) - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开 - keepaliveInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTicker *time.Ticker - if keepaliveInterval > 0 { - keepaliveTicker = time.NewTicker(keepaliveInterval) - defer keepaliveTicker.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTicker != nil { - keepaliveCh = keepaliveTicker.C - } - lastDataAt := time.Now() - - cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity gemini") - - // 仅发送一次错误事件,避免多次写入导致协议混乱 - errorEventSent := false - sendErrorEvent := func(reason string) { - if errorEventSent || cw.Disconnected() { - return - } - errorEventSent = true - _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) - flusher.Flush() - } - - for { - select { - case ev, ok := <-events: - if !ok { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil - } - if ev.err != nil { - if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity gemini"); handled { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil - } - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) - sendErrorEvent("response_too_large") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err - } - sendErrorEvent("stream_read_error") - return nil, ev.err - } - - lastDataAt = time.Now() - - line := ev.line - trimmed := strings.TrimRight(line, "\r\n") - if strings.HasPrefix(trimmed, "data:") { - payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) - if payload == "" || payload == "[DONE]" { - cw.Fprintf("%s\n", line) - continue - } - - // 解包 v1internal 响应 - inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) - if parseErr == nil && inner != nil { - payload = string(inner) - } - - // 解析 usage - if u := extractGeminiUsage(inner); u != nil { - usage = u - } - var parsed map[string]any - if json.Unmarshal(inner, &parsed) == nil { - // Check for MALFORMED_FUNCTION_CALL - if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { - if cand, ok := candidates[0].(map[string]any); ok { - if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward stream") - if content, ok := cand["content"]; ok { - if b, err := json.Marshal(content); err == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) - } - } - } - } - } - } - - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - cw.Fprintf("data: %s\n\n", payload) - continue - } - - cw.Fprintf("%s\n", line) - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if cw.Disconnected() { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity gemini), returning collected usage") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") - sendErrorEvent("stream_timeout") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") - - case <-keepaliveCh: - if cw.Disconnected() { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - continue - } - // SSE ping/keepalive:保持连接活跃防止 Cloudflare Tunnel 等代理断开 - if !cw.Fprintf(":\n\n") { - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity gemini), continuing to drain upstream for billing") - continue - } - } - } -} - -// handleGeminiStreamToNonStreaming 读取上游流式响应,合并为非流式响应返回给客户端 -// Gemini 流式响应是增量的,需要累积所有 chunk 的内容 -func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - usage := &ClaudeUsage{} - var firstTokenMs *int - var last map[string]any - var lastWithParts map[string]any - var collectedImageParts []map[string]any // 收集所有包含图片的 parts - var collectedTextParts []string // 收集所有文本片段 - - type scanEvent struct { - line string - err error - } - - // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - // 上游数据间隔超时保护(防止上游挂起长期占用连接) - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 流结束,返回收集的响应 - goto returnResponse - } - if ev.err != nil { - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity non-stream): max_size=%d error=%v", maxLineSize, ev.err) - } - return nil, ev.err - } - - line := ev.line - trimmed := strings.TrimRight(line, "\r\n") - - if !strings.HasPrefix(trimmed, "data:") { - continue - } - - payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) - if payload == "" || payload == "[DONE]" { - continue - } - - // 解包 v1internal 响应 - inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) - if parseErr != nil { - continue - } - - var parsed map[string]any - if err := json.Unmarshal(inner, &parsed); err != nil { - continue - } - - // 记录首 token 时间 - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - last = parsed - - // 提取 usage - if u := extractGeminiUsage(inner); u != nil { - usage = u - } - - // Check for MALFORMED_FUNCTION_CALL - if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { - if cand, ok := candidates[0].(map[string]any); ok { - if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward non-stream collect") - if content, ok := cand["content"]; ok { - if b, err := json.Marshal(content); err == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) - } - } - } - } - } - - // 保留最后一个有 parts 的响应 - if parts := extractGeminiParts(parsed); len(parts) > 0 { - lastWithParts = parsed - // 收集包含图片和文本的 parts - for _, part := range parts { - if inlineData, ok := part["inlineData"].(map[string]any); ok { - collectedImageParts = append(collectedImageParts, part) - _ = inlineData // 避免 unused 警告 - } - if text, ok := part["text"].(string); ok && text != "" { - collectedTextParts = append(collectedTextParts, text) - } - } - } - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity non-stream)") - return nil, fmt.Errorf("stream data interval timeout") - } - } - -returnResponse: - // 选择最后一个有效响应 - finalResponse := pickGeminiCollectResult(last, lastWithParts) - - // 处理空响应情况 — 触发同账号重试 + failover 切换账号 - if last == nil && lastWithParts == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (gemini non-stream), triggering failover") - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), - RetryableOnSameAccount: true, - } - } - - // 如果收集到了图片 parts,需要合并到最终响应中 - if len(collectedImageParts) > 0 { - finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts) - } - - // 如果收集到了文本,需要合并到最终响应中 - if len(collectedTextParts) > 0 { - finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts) - } - - respBody, err := json.Marshal(finalResponse) - if err != nil { - return nil, fmt.Errorf("failed to marshal response: %w", err) - } - c.Data(http.StatusOK, "application/json", respBody) - - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil -} - -// getOrCreateGeminiParts 获取 Gemini 响应的 parts 结构,返回深拷贝和更新回调 -func getOrCreateGeminiParts(response map[string]any) (result map[string]any, existingParts []any, setParts func([]any)) { - // 深拷贝 response - result = make(map[string]any) - for k, v := range response { - result[k] = v - } - - // 获取或创建 candidates - candidates, ok := result["candidates"].([]any) - if !ok || len(candidates) == 0 { - candidates = []any{map[string]any{}} - } - - // 获取第一个 candidate - candidate, ok := candidates[0].(map[string]any) - if !ok { - candidate = make(map[string]any) - candidates[0] = candidate - } - - // 获取或创建 content - content, ok := candidate["content"].(map[string]any) - if !ok { - content = map[string]any{"role": "model"} - candidate["content"] = content - } - - // 获取现有 parts - existingParts, ok = content["parts"].([]any) - if !ok { - existingParts = []any{} - } - - // 返回更新回调 - setParts = func(newParts []any) { - content["parts"] = newParts - result["candidates"] = candidates - } - - return result, existingParts, setParts -} - -// mergeCollectedPartsToResponse 将收集的所有 parts 合并到 Gemini 响应中 -// 这个函数会合并所有类型的 parts:text、thinking、functionCall、inlineData 等 -// 保持原始顺序,只合并连续的普通 text parts -func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map[string]any) map[string]any { - if len(collectedParts) == 0 { - return response - } - - result, _, setParts := getOrCreateGeminiParts(response) - - // 合并策略: - // 1. 保持原始顺序 - // 2. 连续的普通 text parts 合并为一个 - // 3. thinking、functionCall、inlineData 等保持原样 - var mergedParts []any - var textBuffer strings.Builder - - flushTextBuffer := func() { - if textBuffer.Len() > 0 { - mergedParts = append(mergedParts, map[string]any{ - "text": textBuffer.String(), - }) - textBuffer.Reset() - } - } - - for _, part := range collectedParts { - // 检查是否是普通 text part - if text, ok := part["text"].(string); ok { - // 检查是否有 thought 标记 - if thought, _ := part["thought"].(bool); thought { - // thinking part,先刷新 text buffer,然后保留原样 - flushTextBuffer() - mergedParts = append(mergedParts, part) - } else { - // 普通 text,累积到 buffer - _, _ = textBuffer.WriteString(text) - } - } else { - // 非 text part(functionCall、inlineData 等),先刷新 text buffer,然后保留原样 - flushTextBuffer() - mergedParts = append(mergedParts, part) - } - } - - // 刷新剩余的 text - flushTextBuffer() - - setParts(mergedParts) - return result -} - -// mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中 -func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any { - if len(imageParts) == 0 { - return response - } - - result, existingParts, setParts := getOrCreateGeminiParts(response) - - // 检查现有 parts 中是否已经有图片 - for _, p := range existingParts { - if pm, ok := p.(map[string]any); ok { - if _, hasInline := pm["inlineData"]; hasInline { - return result // 已有图片,不重复添加 - } - } - } - - // 添加收集到的图片 parts - for _, imgPart := range imageParts { - existingParts = append(existingParts, imgPart) - } - setParts(existingParts) - return result -} - -// mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中 -func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any { - if len(textParts) == 0 { - return response - } - - mergedText := strings.Join(textParts, "") - result, existingParts, setParts := getOrCreateGeminiParts(response) - - // 查找并更新第一个 text part,或创建新的 - newParts := make([]any, 0, len(existingParts)+1) - textUpdated := false - - for _, p := range existingParts { - pm, ok := p.(map[string]any) - if !ok { - newParts = append(newParts, p) - continue - } - if _, hasText := pm["text"]; hasText && !textUpdated { - // 用累积的文本替换 - newPart := make(map[string]any) - for k, v := range pm { - newPart[k] = v - } - newPart["text"] = mergedText - newParts = append(newParts, newPart) - textUpdated = true - } else { - newParts = append(newParts, pm) - } - } - - if !textUpdated { - newParts = append([]any{map[string]any{"text": mergedText}}, newParts...) - } - - setParts(newParts) - return result -} - -func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error { - MarkResponseCommitted(c) - c.JSON(status, gin.H{ - "type": "error", - "error": gin.H{"type": errType, "message": message}, - }) - return fmt.Errorf("%s", message) -} - -// WriteMappedClaudeError 导出版本,供 handler 层使用(如 fallback 错误处理) -func (s *AntigravityGatewayService) WriteMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { - return s.writeMappedClaudeError(c, account, upstreamStatus, upstreamRequestID, body) -} - -func (s *AntigravityGatewayService) writeMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { - MarkResponseCommitted(c) - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - logBody, maxBytes := s.getLogConfig() - upstreamDetail := s.getUpstreamErrorDetail(body) - setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: upstreamStatus, - UpstreamRequestID: upstreamRequestID, - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - // 记录上游错误详情便于排障(可选:由配置控制;不回显到客户端) - if logBody { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream_error status=%d body=%s", upstreamStatus, truncateForLog(body, maxBytes)) - } - - // 检查错误透传规则 - if ptStatus, ptErrType, ptErrMsg, matched := applyErrorPassthroughRule( - c, account.Platform, upstreamStatus, body, - 0, "", "", - ); matched { - c.JSON(ptStatus, gin.H{ - "type": "error", - "error": gin.H{"type": ptErrType, "message": ptErrMsg}, - }) - if upstreamMsg == "" { - return fmt.Errorf("upstream error: %d", upstreamStatus) - } - return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) - } - - var statusCode int - var errType, errMsg string - - switch upstreamStatus { - case 400: - statusCode = http.StatusBadRequest - errType = "invalid_request_error" - errMsg = getPassthroughOrDefault(upstreamMsg, "Invalid request") - case 401: - statusCode = http.StatusBadGateway - errType = "authentication_error" - errMsg = "Upstream authentication failed" - case 403: - statusCode = http.StatusBadGateway - errType = "permission_error" - errMsg = "Upstream access forbidden" - case 429: - statusCode = http.StatusTooManyRequests - errType = "rate_limit_error" - errMsg = "Upstream rate limit exceeded" - case 529: - statusCode = http.StatusServiceUnavailable - errType = "overloaded_error" - errMsg = "Upstream service overloaded" - default: - statusCode = http.StatusBadGateway - errType = "upstream_error" - errMsg = "Upstream request failed" - } - - c.JSON(statusCode, gin.H{ - "type": "error", - "error": gin.H{"type": errType, "message": errMsg}, - }) - if upstreamMsg == "" { - return fmt.Errorf("upstream error: %d", upstreamStatus) - } - return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) -} - -func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int, message string) error { - MarkResponseCommitted(c) - statusStr := "UNKNOWN" - switch status { - case 400: - statusStr = "INVALID_ARGUMENT" - case 404: - statusStr = "NOT_FOUND" - case 429: - statusStr = "RESOURCE_EXHAUSTED" - case 500: - statusStr = "INTERNAL" - case 502, 503: - statusStr = "UNAVAILABLE" - } - - c.JSON(status, gin.H{ - "error": gin.H{ - "code": status, - "message": message, - "status": statusStr, - }, - }) - return fmt.Errorf("%s", message) -} - -// handleClaudeStreamToNonStreaming 收集上游流式响应,转换为 Claude 非流式格式返回 -// 用于处理客户端非流式请求但上游只支持流式的情况 -func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - var firstTokenMs *int - var last map[string]any - var lastWithParts map[string]any - var collectedParts []map[string]any // 收集所有 parts(包括 text、thinking、functionCall、inlineData 等) - - type scanEvent struct { - line string - err error - } - - // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - // 上游数据间隔超时保护(防止上游挂起长期占用连接) - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 流结束,转换并返回响应 - goto returnResponse - } - if ev.err != nil { - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err) - } - return nil, ev.err - } - - line := ev.line - trimmed := strings.TrimRight(line, "\r\n") - - if !strings.HasPrefix(trimmed, "data:") { - continue - } - - payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) - if payload == "" || payload == "[DONE]" { - continue - } - - // 解包 v1internal 响应 - inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) - if parseErr != nil { - continue - } - - var parsed map[string]any - if err := json.Unmarshal(inner, &parsed); err != nil { - continue - } - - // 记录首 token 时间 - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - last = parsed - - // 保留最后一个有 parts 的响应,并收集所有 parts - if parts := extractGeminiParts(parsed); len(parts) > 0 { - lastWithParts = parsed - - // 收集所有 parts(text、thinking、functionCall、inlineData 等) - collectedParts = append(collectedParts, parts...) - } - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity claude non-stream)") - return nil, fmt.Errorf("stream data interval timeout") - } - } - -returnResponse: - // 选择最后一个有效响应 - finalResponse := pickGeminiCollectResult(last, lastWithParts) - - // 处理空响应情况 — 触发同账号重试 + failover 切换账号 - if last == nil && lastWithParts == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (claude non-stream), triggering failover") - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), - RetryableOnSameAccount: true, - } - } - - // 将收集的所有 parts 合并到最终响应中 - if len(collectedParts) > 0 { - finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts) - } - - // 序列化为 JSON(Gemini 格式) - geminiBody, err := json.Marshal(finalResponse) - if err != nil { - return nil, fmt.Errorf("failed to marshal gemini response: %w", err) - } - - // 转换 Gemini 响应为 Claude 格式 - claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody)) - return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") - } - - c.Data(http.StatusOK, "application/json", claudeResp) - - // 转换为 service.ClaudeUsage - usage := &ClaudeUsage{ - InputTokens: agUsage.InputTokens, - OutputTokens: agUsage.OutputTokens, - CacheCreationInputTokens: agUsage.CacheCreationInputTokens, - CacheReadInputTokens: agUsage.CacheReadInputTokens, - } - - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil -} - -// handleClaudeStreamingResponse 处理 Claude 流式响应(Gemini SSE → Claude SSE 转换) -func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { - c.Header("Content-Type", "text/event-stream") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - c.Status(http.StatusOK) - - flusher, ok := c.Writer.(http.Flusher) - if !ok { - return nil, errors.New("streaming not supported") - } - - processor := antigravity.NewStreamingProcessor(originalModel) - var firstTokenMs *int - // 使用 Scanner 并限制单行大小,避免 ReadString 无上限导致 OOM - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - // 辅助函数:转换 antigravity.ClaudeUsage 到 service.ClaudeUsage - convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage { - if agUsage == nil { - return &ClaudeUsage{} - } - return &ClaudeUsage{ - InputTokens: agUsage.InputTokens, - OutputTokens: agUsage.OutputTokens, - CacheCreationInputTokens: agUsage.CacheCreationInputTokens, - CacheReadInputTokens: agUsage.CacheReadInputTokens, - } - } - - type scanEvent struct { - line string - err error - } - // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开 - keepaliveInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTicker *time.Ticker - if keepaliveInterval > 0 { - keepaliveTicker = time.NewTicker(keepaliveInterval) - defer keepaliveTicker.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTicker != nil { - keepaliveCh = keepaliveTicker.C - } - lastDataAt := time.Now() - - cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity claude") - - // 仅发送一次错误事件,避免多次写入导致协议混乱 - errorEventSent := false - sendErrorEvent := func(reason string) { - if errorEventSent || cw.Disconnected() { - return - } - errorEventSent = true - _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) - flusher.Flush() - } - - // finishUsage 是获取 processor 最终 usage 的辅助函数 - finishUsage := func() *ClaudeUsage { - _, agUsage := processor.Finish() - return convertUsage(agUsage) - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 上游完成,发送结束事件 - finalEvents, agUsage := processor.Finish() - if len(finalEvents) > 0 { - cw.Write(finalEvents) - } else if !processor.MessageStartSent() && !cw.Disconnected() { - // 整个流未收到任何可解析的上游数据(全部 SSE 行均无法被 JSON 解析), - // 触发 failover 在同账号重试,避免向客户端发出缺少 message_start 的残缺流 - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Claude-Stream] empty stream response (no valid events parsed), triggering failover") - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), - RetryableOnSameAccount: true, - } - } - return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil - } - if ev.err != nil { - if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity claude"); handled { - return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil - } - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) - sendErrorEvent("response_too_large") - return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err - } - sendErrorEvent("stream_read_error") - return nil, fmt.Errorf("stream read error: %w", ev.err) - } - - lastDataAt = time.Now() - - // 处理 SSE 行,转换为 Claude 格式 - claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n")) - if len(claudeEvents) > 0 { - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - cw.Write(claudeEvents) - } - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if cw.Disconnected() { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity claude), returning collected usage") - return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: true}, nil - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") - sendErrorEvent("stream_timeout") - return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") - - case <-keepaliveCh: - if cw.Disconnected() { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - continue - } - // SSE ping 事件:Anthropic 原生格式,客户端会正确处理, - // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 - if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity claude), continuing to drain upstream for billing") - continue - } - } - } -} - -func (s *AntigravityGatewayService) extractImageInputSize(body []byte) string { - var req antigravity.GeminiRequest - if err := json.Unmarshal(body, &req); err != nil { - return "" - } - - if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil { - return strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize) - } - - return "" -} - -// isImageGenerationModel 判断模型是否为图片生成模型 -// 支持的模型:gemini-3.1-flash-image, gemini-3-pro-image, gemini-2.5-flash-image 等 -func isImageGenerationModel(model string) bool { - modelLower := strings.ToLower(model) - // 移除 models/ 前缀 - modelLower = strings.TrimPrefix(modelLower, "models/") - - // 精确匹配或前缀匹配 - return modelLower == "gemini-3.1-flash-image" || - modelLower == "gemini-3.1-flash-image-preview" || - strings.HasPrefix(modelLower, "gemini-3.1-flash-image-") || - modelLower == "gemini-3-pro-image" || - modelLower == "gemini-3-pro-image-preview" || - strings.HasPrefix(modelLower, "gemini-3-pro-image-") || - modelLower == "gemini-2.5-flash-image" || - modelLower == "gemini-2.5-flash-image-preview" || - strings.HasPrefix(modelLower, "gemini-2.5-flash-image-") -} - -// cleanGeminiRequest 清理 Gemini 请求体中的 Schema -func cleanGeminiRequest(body []byte) ([]byte, error) { - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, err - } - - modified := false - - // 1. 清理 Tools - if tools, ok := payload["tools"].([]any); ok && len(tools) > 0 { - for _, t := range tools { - toolMap, ok := t.(map[string]any) - if !ok { - continue - } - - // function_declarations (snake_case) or functionDeclarations (camelCase) - var funcs []any - if f, ok := toolMap["functionDeclarations"].([]any); ok { - funcs = f - } else if f, ok := toolMap["function_declarations"].([]any); ok { - funcs = f - } - - if len(funcs) == 0 { - continue - } - - for _, f := range funcs { - funcMap, ok := f.(map[string]any) - if !ok { - continue - } - - if params, ok := funcMap["parameters"].(map[string]any); ok { - antigravity.DeepCleanUndefined(params) - cleaned := antigravity.CleanJSONSchema(params) - funcMap["parameters"] = cleaned - modified = true - } - } - } - } - - if !modified { - return body, nil - } - - return json.Marshal(payload) -} - -// filterEmptyPartsFromGeminiRequest 过滤掉 parts 为空的消息 -// Gemini API 不接受空 parts,需要在请求前过滤 -func filterEmptyPartsFromGeminiRequest(body []byte) ([]byte, error) { - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, err - } - - contents, ok := payload["contents"].([]any) - if !ok || len(contents) == 0 { - return body, nil - } - - filtered := make([]any, 0, len(contents)) - modified := false - - for _, c := range contents { - contentMap, ok := c.(map[string]any) - if !ok { - filtered = append(filtered, c) - continue - } - - parts, hasParts := contentMap["parts"] - if !hasParts { - filtered = append(filtered, c) - continue - } - - partsSlice, ok := parts.([]any) - if !ok { - filtered = append(filtered, c) - continue - } - - // 跳过 parts 为空数组的消息 - if len(partsSlice) == 0 { - modified = true - continue - } - - filtered = append(filtered, c) - } - - if !modified { - return body, nil - } - - payload["contents"] = filtered - return json.Marshal(payload) -} - -// ForwardUpstream 使用 base_url + /v1/messages + 双 header 认证透传上游 Claude 请求 -func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { - startTime := time.Now() - sessionID := getSessionID(c) - prefix := logPrefix(sessionID, account.Name) - - // 获取上游配置 - baseURL := strings.TrimSpace(account.GetCredential("base_url")) - apiKey := strings.TrimSpace(account.GetCredential("api_key")) - if baseURL == "" || apiKey == "" { - return nil, fmt.Errorf("upstream account missing base_url or api_key") - } - baseURL = strings.TrimSuffix(baseURL, "/") - - // 解析请求获取模型信息 - var claudeReq antigravity.ClaudeRequest - if err := json.Unmarshal(body, &claudeReq); err != nil { - return nil, fmt.Errorf("parse claude request: %w", err) - } - if strings.TrimSpace(claudeReq.Model) == "" { - return nil, fmt.Errorf("missing model") - } - originalModel := claudeReq.Model - - // 构建上游请求 URL - upstreamURL := baseURL + "/v1/messages" - - // 能力维度 sanitize:Anthropic-compatible 上游透传路径也需要保证 body↔beta header - // 对称。客户端 anthropic-beta header 不含 context-management-2025-06-27 但 body 带 - // context_management 时 strip,与 Anthropic 直连 / Bedrock / Vertex 路径保持一致。 - clientBeta := c.GetHeader("anthropic-beta") - if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { - body = sanitized - } - - // 创建请求 - req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("create upstream request: %w", err) - } - - // 设置请求头 - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+apiKey) - req.Header.Set("x-api-key", apiKey) // Claude API 兼容 - - // 透传 Claude 相关 headers - if v := c.GetHeader("anthropic-version"); v != "" { - req.Header.Set("anthropic-version", v) - } - if v := clientBeta; v != "" { - req.Header.Set("anthropic-beta", v) - } - - // 代理 URL - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - // 发送请求 - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s upstream request failed: %v", prefix, err) - return nil, fmt.Errorf("upstream request failed: %w", err) - } - defer func() { _ = resp.Body.Close() }() - - // 处理错误响应 - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - - // 429 错误时标记账号限流 - if resp.StatusCode == http.StatusTooManyRequests { - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", false) - } - - // 透传上游错误 - c.Header("Content-Type", resp.Header.Get("Content-Type")) - c.Status(resp.StatusCode) - _, _ = c.Writer.Write(respBody) - - return &ForwardResult{ - Model: originalModel, - }, nil - } - - // 处理成功响应(流式/非流式) - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - - if claudeReq.Stream { - // 流式响应:透传 - c.Header("Content-Type", "text/event-stream") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - c.Status(http.StatusOK) - - streamRes := s.streamUpstreamResponse(c, resp, startTime) - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - clientDisconnect = streamRes.clientDisconnect - } else { - // 非流式响应:直接透传 - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("read upstream response: %w", err) - } - - // 提取 usage - usage = s.extractClaudeUsage(respBody) - - c.Header("Content-Type", resp.Header.Get("Content-Type")) - c.Status(http.StatusOK) - _, _ = c.Writer.Write(respBody) - } - - // 构建计费结果 - duration := time.Since(startTime) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=success duration_ms=%d", prefix, duration.Milliseconds()) - - return &ForwardResult{ - Model: originalModel, - Stream: claudeReq.Stream, - Duration: duration, - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - Usage: ClaudeUsage{ - InputTokens: usage.InputTokens, - OutputTokens: usage.OutputTokens, - CacheReadInputTokens: usage.CacheReadInputTokens, - CacheCreationInputTokens: usage.CacheCreationInputTokens, - }, - }, nil -} - -// streamUpstreamResponse 透传上游 SSE 流并提取 Claude usage -func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp *http.Response, startTime time.Time) *antigravityStreamResult { - usage := &ClaudeUsage{} - var firstTokenMs *int - - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanner.Buffer(make([]byte, 64*1024), maxLineSize) - - type scanEvent struct { - line string - err error - } - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func() { - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }() - defer close(done) - - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开 - keepaliveInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTicker *time.Ticker - if keepaliveInterval > 0 { - keepaliveTicker = time.NewTicker(keepaliveInterval) - defer keepaliveTicker.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTicker != nil { - keepaliveCh = keepaliveTicker.C - } - lastDataAt := time.Now() - - flusher, _ := c.Writer.(http.Flusher) - cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity upstream") - - for { - select { - case ev, ok := <-events: - if !ok { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()} - } - if ev.err != nil { - if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity upstream"); handled { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect} - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream read error (antigravity upstream): %v", ev.err) - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} - } - - lastDataAt = time.Now() - - line := ev.line - - // 记录首 token 时间 - if firstTokenMs == nil && len(line) > 0 { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - // 尝试从 message_delta 或 message_stop 事件提取 usage - s.extractSSEUsage(line, usage) - - // 透传行 - cw.Fprintf("%s\n", line) - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if cw.Disconnected() { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity upstream), returning collected usage") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true} - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity upstream)") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} - - case <-keepaliveCh: - if cw.Disconnected() { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - continue - } - // SSE ping 事件:Anthropic 原生格式,客户端会正确处理, - // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 - if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity upstream), continuing to drain upstream for billing") - continue - } - } - } -} - -// extractSSEUsage 从 SSE data 行中提取 Claude usage(用于流式透传场景) -// -// Anthropic streaming 的 usage 字段分布在两类事件中: -// - message_start:嵌套在 event.message.usage(input_tokens、cache_creation_input_tokens、 -// cache_read_input_tokens 等输入侧字段) -// - message_delta:位于顶层 event.usage(流结束时的最终 output_tokens) -// -// 仅读取顶层 event.usage 会漏掉 message_start 的输入侧字段,导致流式透传请求落库的 -// usage_logs 记录 input_tokens=0。 -func (s *AntigravityGatewayService) extractSSEUsage(line string, usage *ClaudeUsage) { - if !strings.HasPrefix(line, "data: ") { - return - } - dataStr := strings.TrimPrefix(line, "data: ") - var event map[string]any - if json.Unmarshal([]byte(dataStr), &event) != nil { - return - } - var u map[string]any - if eventType, _ := event["type"].(string); eventType == "message_start" { - if msg, ok := event["message"].(map[string]any); ok { - u, _ = msg["usage"].(map[string]any) - } - } else { - u, _ = event["usage"].(map[string]any) - } - if u == nil { - return - } - if v, ok := u["input_tokens"].(float64); ok && int(v) > 0 { - usage.InputTokens = int(v) - } - if v, ok := u["output_tokens"].(float64); ok && int(v) > 0 { - usage.OutputTokens = int(v) - } - if v, ok := u["cache_read_input_tokens"].(float64); ok && int(v) > 0 { - usage.CacheReadInputTokens = int(v) - } - if v, ok := u["cache_creation_input_tokens"].(float64); ok && int(v) > 0 { - usage.CacheCreationInputTokens = int(v) - } - // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 - if cc, ok := u["cache_creation"].(map[string]any); ok { - if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { - usage.CacheCreation5mTokens = int(v) - } - if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { - usage.CacheCreation1hTokens = int(v) - } - } -} - -// extractClaudeUsage 从非流式 Claude 响应提取 usage -func (s *AntigravityGatewayService) extractClaudeUsage(body []byte) *ClaudeUsage { - usage := &ClaudeUsage{} - var resp map[string]any - if json.Unmarshal(body, &resp) != nil { - return usage - } - if u, ok := resp["usage"].(map[string]any); ok { - if v, ok := u["input_tokens"].(float64); ok { - usage.InputTokens = int(v) - } - if v, ok := u["output_tokens"].(float64); ok { - usage.OutputTokens = int(v) - } - if v, ok := u["cache_read_input_tokens"].(float64); ok { - usage.CacheReadInputTokens = int(v) - } - if v, ok := u["cache_creation_input_tokens"].(float64); ok { - usage.CacheCreationInputTokens = int(v) - } - // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 - if cc, ok := u["cache_creation"].(map[string]any); ok { - if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { - usage.CacheCreation5mTokens = int(v) - } - if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { - usage.CacheCreation1hTokens = int(v) - } - } - } - return usage -} diff --git a/backend/internal/service/antigravity_gateway_streaming.go b/backend/internal/service/antigravity_gateway_streaming.go new file mode 100644 index 0000000000..1a6c59f617 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_streaming.go @@ -0,0 +1,1150 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +type antigravityStreamResult struct { + usage *ClaudeUsage + firstTokenMs *int + clientDisconnect bool // 客户端是否在流式传输过程中断开 +} + +// antigravityClientWriter 封装流式响应的客户端写入,自动检测断开并标记。 +// 断开后所有写入操作变为 no-op,调用方通过 Disconnected() 判断是否继续 drain 上游。 +type antigravityClientWriter struct { + w gin.ResponseWriter + flusher http.Flusher + disconnected bool + prefix string // 日志前缀,标识来源方法 +} + +func newAntigravityClientWriter(w gin.ResponseWriter, flusher http.Flusher, prefix string) *antigravityClientWriter { + return &antigravityClientWriter{w: w, flusher: flusher, prefix: prefix} +} + +// Write 写入数据到客户端,写入失败时标记断开并返回 false +func (cw *antigravityClientWriter) Write(p []byte) bool { + if cw.disconnected { + return false + } + if _, err := cw.w.Write(p); err != nil { + cw.markDisconnected() + return false + } + cw.flusher.Flush() + return true +} + +// Fprintf 格式化写入数据到客户端,写入失败时标记断开并返回 false +func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool { + if cw.disconnected { + return false + } + if _, err := fmt.Fprintf(cw.w, format, args...); err != nil { + cw.markDisconnected() + return false + } + cw.flusher.Flush() + return true +} + +func (cw *antigravityClientWriter) Disconnected() bool { return cw.disconnected } + +func (cw *antigravityClientWriter) markDisconnected() { + cw.disconnected = true + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during streaming (%s), continuing to drain upstream for billing", cw.prefix) +} + +// handleStreamReadError 处理上游读取错误的通用逻辑。 +// 返回 (clientDisconnect, handled):handled=true 表示错误已处理,调用方应返回已收集的 usage。 +func handleStreamReadError(err error, clientDisconnected bool, prefix string) (disconnect bool, handled bool) { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("service.antigravity_gateway", "Context canceled during streaming (%s), returning collected usage", prefix) + return true, true + } + if clientDisconnected { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream read error after client disconnect (%s): %v, returning collected usage", prefix, err) + return true, true + } + return false, false +} + +func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { + c.Status(resp.StatusCode) + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + contentType = "text/event-stream; charset=utf-8" + } + c.Header("Content-Type", contentType) + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + return nil, errors.New("streaming not supported") + } + + // 使用 Scanner 并限制单行大小,避免 ReadString 无上限导致 OOM + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + usage := &ClaudeUsage{} + var firstTokenMs *int + + type scanEvent struct { + line string + err error + } + // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + // 上游数据间隔超时保护(防止上游挂起长期占用连接) + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开 + keepaliveInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTicker *time.Ticker + if keepaliveInterval > 0 { + keepaliveTicker = time.NewTicker(keepaliveInterval) + defer keepaliveTicker.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTicker != nil { + keepaliveCh = keepaliveTicker.C + } + lastDataAt := time.Now() + + cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity gemini") + + // 仅发送一次错误事件,避免多次写入导致协议混乱 + errorEventSent := false + sendErrorEvent := func(reason string) { + if errorEventSent || cw.Disconnected() { + return + } + errorEventSent = true + _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) + flusher.Flush() + } + + for { + select { + case ev, ok := <-events: + if !ok { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil + } + if ev.err != nil { + if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity gemini"); handled { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil + } + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) + sendErrorEvent("response_too_large") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err + } + sendErrorEvent("stream_read_error") + return nil, ev.err + } + + lastDataAt = time.Now() + + line := ev.line + trimmed := strings.TrimRight(line, "\r\n") + if strings.HasPrefix(trimmed, "data:") { + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if payload == "" || payload == "[DONE]" { + cw.Fprintf("%s\n", line) + continue + } + + // 解包 v1internal 响应 + inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) + if parseErr == nil && inner != nil { + payload = string(inner) + } + + // 解析 usage + if u := extractGeminiUsage(inner); u != nil { + usage = u + } + var parsed map[string]any + if json.Unmarshal(inner, &parsed) == nil { + // Check for MALFORMED_FUNCTION_CALL + if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { + if cand, ok := candidates[0].(map[string]any); ok { + if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward stream") + if content, ok := cand["content"]; ok { + if b, err := json.Marshal(content); err == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) + } + } + } + } + } + } + + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + cw.Fprintf("data: %s\n\n", payload) + continue + } + + cw.Fprintf("%s\n", line) + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if cw.Disconnected() { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity gemini), returning collected usage") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") + sendErrorEvent("stream_timeout") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") + + case <-keepaliveCh: + if cw.Disconnected() { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + continue + } + // SSE ping/keepalive:保持连接活跃防止 Cloudflare Tunnel 等代理断开 + if !cw.Fprintf(":\n\n") { + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity gemini), continuing to drain upstream for billing") + continue + } + } + } +} + +// handleGeminiStreamToNonStreaming 读取上游流式响应,合并为非流式响应返回给客户端 +// Gemini 流式响应是增量的,需要累积所有 chunk 的内容 +func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + usage := &ClaudeUsage{} + var firstTokenMs *int + var last map[string]any + var lastWithParts map[string]any + var collectedImageParts []map[string]any // 收集所有包含图片的 parts + var collectedTextParts []string // 收集所有文本片段 + + type scanEvent struct { + line string + err error + } + + // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + // 上游数据间隔超时保护(防止上游挂起长期占用连接) + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 流结束,返回收集的响应 + goto returnResponse + } + if ev.err != nil { + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity non-stream): max_size=%d error=%v", maxLineSize, ev.err) + } + return nil, ev.err + } + + line := ev.line + trimmed := strings.TrimRight(line, "\r\n") + + if !strings.HasPrefix(trimmed, "data:") { + continue + } + + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if payload == "" || payload == "[DONE]" { + continue + } + + // 解包 v1internal 响应 + inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) + if parseErr != nil { + continue + } + + var parsed map[string]any + if err := json.Unmarshal(inner, &parsed); err != nil { + continue + } + + // 记录首 token 时间 + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + last = parsed + + // 提取 usage + if u := extractGeminiUsage(inner); u != nil { + usage = u + } + + // Check for MALFORMED_FUNCTION_CALL + if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { + if cand, ok := candidates[0].(map[string]any); ok { + if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward non-stream collect") + if content, ok := cand["content"]; ok { + if b, err := json.Marshal(content); err == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) + } + } + } + } + } + + // 保留最后一个有 parts 的响应 + if parts := extractGeminiParts(parsed); len(parts) > 0 { + lastWithParts = parsed + // 收集包含图片和文本的 parts + for _, part := range parts { + if inlineData, ok := part["inlineData"].(map[string]any); ok { + collectedImageParts = append(collectedImageParts, part) + _ = inlineData // 避免 unused 警告 + } + if text, ok := part["text"].(string); ok && text != "" { + collectedTextParts = append(collectedTextParts, text) + } + } + } + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity non-stream)") + return nil, fmt.Errorf("stream data interval timeout") + } + } + +returnResponse: + // 选择最后一个有效响应 + finalResponse := pickGeminiCollectResult(last, lastWithParts) + + // 处理空响应情况 — 触发同账号重试 + failover 切换账号 + if last == nil && lastWithParts == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (gemini non-stream), triggering failover") + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), + RetryableOnSameAccount: true, + } + } + + // 如果收集到了图片 parts,需要合并到最终响应中 + if len(collectedImageParts) > 0 { + finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts) + } + + // 如果收集到了文本,需要合并到最终响应中 + if len(collectedTextParts) > 0 { + finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts) + } + + respBody, err := json.Marshal(finalResponse) + if err != nil { + return nil, fmt.Errorf("failed to marshal response: %w", err) + } + c.Data(http.StatusOK, "application/json", respBody) + + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil +} + +// getOrCreateGeminiParts 获取 Gemini 响应的 parts 结构,返回深拷贝和更新回调 +func getOrCreateGeminiParts(response map[string]any) (result map[string]any, existingParts []any, setParts func([]any)) { + // 深拷贝 response + result = make(map[string]any) + for k, v := range response { + result[k] = v + } + + // 获取或创建 candidates + candidates, ok := result["candidates"].([]any) + if !ok || len(candidates) == 0 { + candidates = []any{map[string]any{}} + } + + // 获取第一个 candidate + candidate, ok := candidates[0].(map[string]any) + if !ok { + candidate = make(map[string]any) + candidates[0] = candidate + } + + // 获取或创建 content + content, ok := candidate["content"].(map[string]any) + if !ok { + content = map[string]any{"role": "model"} + candidate["content"] = content + } + + // 获取现有 parts + existingParts, ok = content["parts"].([]any) + if !ok { + existingParts = []any{} + } + + // 返回更新回调 + setParts = func(newParts []any) { + content["parts"] = newParts + result["candidates"] = candidates + } + + return result, existingParts, setParts +} + +// mergeCollectedPartsToResponse 将收集的所有 parts 合并到 Gemini 响应中 +// 这个函数会合并所有类型的 parts:text、thinking、functionCall、inlineData 等 +// 保持原始顺序,只合并连续的普通 text parts +func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map[string]any) map[string]any { + if len(collectedParts) == 0 { + return response + } + + result, _, setParts := getOrCreateGeminiParts(response) + + // 合并策略: + // 1. 保持原始顺序 + // 2. 连续的普通 text parts 合并为一个 + // 3. thinking、functionCall、inlineData 等保持原样 + var mergedParts []any + var textBuffer strings.Builder + + flushTextBuffer := func() { + if textBuffer.Len() > 0 { + mergedParts = append(mergedParts, map[string]any{ + "text": textBuffer.String(), + }) + textBuffer.Reset() + } + } + + for _, part := range collectedParts { + // 检查是否是普通 text part + if text, ok := part["text"].(string); ok { + // 检查是否有 thought 标记 + if thought, _ := part["thought"].(bool); thought { + // thinking part,先刷新 text buffer,然后保留原样 + flushTextBuffer() + mergedParts = append(mergedParts, part) + } else { + // 普通 text,累积到 buffer + _, _ = textBuffer.WriteString(text) + } + } else { + // 非 text part(functionCall、inlineData 等),先刷新 text buffer,然后保留原样 + flushTextBuffer() + mergedParts = append(mergedParts, part) + } + } + + // 刷新剩余的 text + flushTextBuffer() + + setParts(mergedParts) + return result +} + +// mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中 +func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any { + if len(imageParts) == 0 { + return response + } + + result, existingParts, setParts := getOrCreateGeminiParts(response) + + // 检查现有 parts 中是否已经有图片 + for _, p := range existingParts { + if pm, ok := p.(map[string]any); ok { + if _, hasInline := pm["inlineData"]; hasInline { + return result // 已有图片,不重复添加 + } + } + } + + // 添加收集到的图片 parts + for _, imgPart := range imageParts { + existingParts = append(existingParts, imgPart) + } + setParts(existingParts) + return result +} + +// mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中 +func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any { + if len(textParts) == 0 { + return response + } + + mergedText := strings.Join(textParts, "") + result, existingParts, setParts := getOrCreateGeminiParts(response) + + // 查找并更新第一个 text part,或创建新的 + newParts := make([]any, 0, len(existingParts)+1) + textUpdated := false + + for _, p := range existingParts { + pm, ok := p.(map[string]any) + if !ok { + newParts = append(newParts, p) + continue + } + if _, hasText := pm["text"]; hasText && !textUpdated { + // 用累积的文本替换 + newPart := make(map[string]any) + for k, v := range pm { + newPart[k] = v + } + newPart["text"] = mergedText + newParts = append(newParts, newPart) + textUpdated = true + } else { + newParts = append(newParts, pm) + } + } + + if !textUpdated { + newParts = append([]any{map[string]any{"text": mergedText}}, newParts...) + } + + setParts(newParts) + return result +} + +func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error { + MarkResponseCommitted(c) + c.JSON(status, gin.H{ + "type": "error", + "error": gin.H{"type": errType, "message": message}, + }) + return fmt.Errorf("%s", message) +} + +// WriteMappedClaudeError 导出版本,供 handler 层使用(如 fallback 错误处理) +func (s *AntigravityGatewayService) WriteMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { + return s.writeMappedClaudeError(c, account, upstreamStatus, upstreamRequestID, body) +} + +func (s *AntigravityGatewayService) writeMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { + MarkResponseCommitted(c) + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + logBody, maxBytes := s.getLogConfig() + upstreamDetail := s.getUpstreamErrorDetail(body) + setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: upstreamStatus, + UpstreamRequestID: upstreamRequestID, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + // 记录上游错误详情便于排障(可选:由配置控制;不回显到客户端) + if logBody { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream_error status=%d body=%s", upstreamStatus, truncateForLog(body, maxBytes)) + } + + // 检查错误透传规则 + if ptStatus, ptErrType, ptErrMsg, matched := applyErrorPassthroughRule( + c, account.Platform, upstreamStatus, body, + 0, "", "", + ); matched { + c.JSON(ptStatus, gin.H{ + "type": "error", + "error": gin.H{"type": ptErrType, "message": ptErrMsg}, + }) + if upstreamMsg == "" { + return fmt.Errorf("upstream error: %d", upstreamStatus) + } + return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) + } + + var statusCode int + var errType, errMsg string + + switch upstreamStatus { + case 400: + statusCode = http.StatusBadRequest + errType = "invalid_request_error" + errMsg = getPassthroughOrDefault(upstreamMsg, "Invalid request") + case 401: + statusCode = http.StatusBadGateway + errType = "authentication_error" + errMsg = "Upstream authentication failed" + case 403: + statusCode = http.StatusBadGateway + errType = "permission_error" + errMsg = "Upstream access forbidden" + case 429: + statusCode = http.StatusTooManyRequests + errType = "rate_limit_error" + errMsg = "Upstream rate limit exceeded" + case 529: + statusCode = http.StatusServiceUnavailable + errType = "overloaded_error" + errMsg = "Upstream service overloaded" + default: + statusCode = http.StatusBadGateway + errType = "upstream_error" + errMsg = "Upstream request failed" + } + + c.JSON(statusCode, gin.H{ + "type": "error", + "error": gin.H{"type": errType, "message": errMsg}, + }) + if upstreamMsg == "" { + return fmt.Errorf("upstream error: %d", upstreamStatus) + } + return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) +} + +func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int, message string) error { + MarkResponseCommitted(c) + statusStr := "UNKNOWN" + switch status { + case 400: + statusStr = "INVALID_ARGUMENT" + case 404: + statusStr = "NOT_FOUND" + case 429: + statusStr = "RESOURCE_EXHAUSTED" + case 500: + statusStr = "INTERNAL" + case 502, 503: + statusStr = "UNAVAILABLE" + } + + c.JSON(status, gin.H{ + "error": gin.H{ + "code": status, + "message": message, + "status": statusStr, + }, + }) + return fmt.Errorf("%s", message) +} + +// handleClaudeStreamToNonStreaming 收集上游流式响应,转换为 Claude 非流式格式返回 +// 用于处理客户端非流式请求但上游只支持流式的情况 +func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + var firstTokenMs *int + var last map[string]any + var lastWithParts map[string]any + var collectedParts []map[string]any // 收集所有 parts(包括 text、thinking、functionCall、inlineData 等) + + type scanEvent struct { + line string + err error + } + + // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + // 上游数据间隔超时保护(防止上游挂起长期占用连接) + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 流结束,转换并返回响应 + goto returnResponse + } + if ev.err != nil { + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err) + } + return nil, ev.err + } + + line := ev.line + trimmed := strings.TrimRight(line, "\r\n") + + if !strings.HasPrefix(trimmed, "data:") { + continue + } + + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if payload == "" || payload == "[DONE]" { + continue + } + + // 解包 v1internal 响应 + inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) + if parseErr != nil { + continue + } + + var parsed map[string]any + if err := json.Unmarshal(inner, &parsed); err != nil { + continue + } + + // 记录首 token 时间 + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + last = parsed + + // 保留最后一个有 parts 的响应,并收集所有 parts + if parts := extractGeminiParts(parsed); len(parts) > 0 { + lastWithParts = parsed + + // 收集所有 parts(text、thinking、functionCall、inlineData 等) + collectedParts = append(collectedParts, parts...) + } + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity claude non-stream)") + return nil, fmt.Errorf("stream data interval timeout") + } + } + +returnResponse: + // 选择最后一个有效响应 + finalResponse := pickGeminiCollectResult(last, lastWithParts) + + // 处理空响应情况 — 触发同账号重试 + failover 切换账号 + if last == nil && lastWithParts == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (claude non-stream), triggering failover") + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), + RetryableOnSameAccount: true, + } + } + + // 将收集的所有 parts 合并到最终响应中 + if len(collectedParts) > 0 { + finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts) + } + + // 序列化为 JSON(Gemini 格式) + geminiBody, err := json.Marshal(finalResponse) + if err != nil { + return nil, fmt.Errorf("failed to marshal gemini response: %w", err) + } + + // 转换 Gemini 响应为 Claude 格式 + claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody)) + return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") + } + + c.Data(http.StatusOK, "application/json", claudeResp) + + // 转换为 service.ClaudeUsage + usage := &ClaudeUsage{ + InputTokens: agUsage.InputTokens, + OutputTokens: agUsage.OutputTokens, + CacheCreationInputTokens: agUsage.CacheCreationInputTokens, + CacheReadInputTokens: agUsage.CacheReadInputTokens, + } + + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil +} + +// handleClaudeStreamingResponse 处理 Claude 流式响应(Gemini SSE → Claude SSE 转换) +func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Status(http.StatusOK) + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + return nil, errors.New("streaming not supported") + } + + processor := antigravity.NewStreamingProcessor(originalModel) + var firstTokenMs *int + // 使用 Scanner 并限制单行大小,避免 ReadString 无上限导致 OOM + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + // 辅助函数:转换 antigravity.ClaudeUsage 到 service.ClaudeUsage + convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage { + if agUsage == nil { + return &ClaudeUsage{} + } + return &ClaudeUsage{ + InputTokens: agUsage.InputTokens, + OutputTokens: agUsage.OutputTokens, + CacheCreationInputTokens: agUsage.CacheCreationInputTokens, + CacheReadInputTokens: agUsage.CacheReadInputTokens, + } + } + + type scanEvent struct { + line string + err error + } + // 独立 goroutine 读取上游,避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开 + keepaliveInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTicker *time.Ticker + if keepaliveInterval > 0 { + keepaliveTicker = time.NewTicker(keepaliveInterval) + defer keepaliveTicker.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTicker != nil { + keepaliveCh = keepaliveTicker.C + } + lastDataAt := time.Now() + + cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity claude") + + // 仅发送一次错误事件,避免多次写入导致协议混乱 + errorEventSent := false + sendErrorEvent := func(reason string) { + if errorEventSent || cw.Disconnected() { + return + } + errorEventSent = true + _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) + flusher.Flush() + } + + // finishUsage 是获取 processor 最终 usage 的辅助函数 + finishUsage := func() *ClaudeUsage { + _, agUsage := processor.Finish() + return convertUsage(agUsage) + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 上游完成,发送结束事件 + finalEvents, agUsage := processor.Finish() + if len(finalEvents) > 0 { + cw.Write(finalEvents) + } else if !processor.MessageStartSent() && !cw.Disconnected() { + // 整个流未收到任何可解析的上游数据(全部 SSE 行均无法被 JSON 解析), + // 触发 failover 在同账号重试,避免向客户端发出缺少 message_start 的残缺流 + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Claude-Stream] empty stream response (no valid events parsed), triggering failover") + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), + RetryableOnSameAccount: true, + } + } + return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil + } + if ev.err != nil { + if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity claude"); handled { + return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil + } + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) + sendErrorEvent("response_too_large") + return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err + } + sendErrorEvent("stream_read_error") + return nil, fmt.Errorf("stream read error: %w", ev.err) + } + + lastDataAt = time.Now() + + // 处理 SSE 行,转换为 Claude 格式 + claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n")) + if len(claudeEvents) > 0 { + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + cw.Write(claudeEvents) + } + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if cw.Disconnected() { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity claude), returning collected usage") + return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: true}, nil + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") + sendErrorEvent("stream_timeout") + return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") + + case <-keepaliveCh: + if cw.Disconnected() { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + continue + } + // SSE ping 事件:Anthropic 原生格式,客户端会正确处理, + // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 + if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity claude), continuing to drain upstream for billing") + continue + } + } + } +} + +func (s *AntigravityGatewayService) extractImageInputSize(body []byte) string { + var req antigravity.GeminiRequest + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + + if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil { + return strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize) + } + + return "" +} + +// isImageGenerationModel 判断模型是否为图片生成模型 +// 支持的模型:gemini-3.1-flash-image, gemini-3-pro-image, gemini-2.5-flash-image 等 +func isImageGenerationModel(model string) bool { + modelLower := strings.ToLower(model) + // 移除 models/ 前缀 + modelLower = strings.TrimPrefix(modelLower, "models/") + + // 精确匹配或前缀匹配 + return modelLower == "gemini-3.1-flash-image" || + modelLower == "gemini-3.1-flash-image-preview" || + strings.HasPrefix(modelLower, "gemini-3.1-flash-image-") || + modelLower == "gemini-3-pro-image" || + modelLower == "gemini-3-pro-image-preview" || + strings.HasPrefix(modelLower, "gemini-3-pro-image-") || + modelLower == "gemini-2.5-flash-image" || + modelLower == "gemini-2.5-flash-image-preview" || + strings.HasPrefix(modelLower, "gemini-2.5-flash-image-") +} diff --git a/backend/internal/service/antigravity_gateway_upstream.go b/backend/internal/service/antigravity_gateway_upstream.go new file mode 100644 index 0000000000..2914863131 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_upstream.go @@ -0,0 +1,375 @@ +package service + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// ForwardUpstream 使用 base_url + /v1/messages + 双 header 认证透传上游 Claude 请求 +func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { + startTime := time.Now() + sessionID := getSessionID(c) + prefix := logPrefix(sessionID, account.Name) + + // 获取上游配置 + baseURL := strings.TrimSpace(account.GetCredential("base_url")) + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if baseURL == "" || apiKey == "" { + return nil, fmt.Errorf("upstream account missing base_url or api_key") + } + baseURL = strings.TrimSuffix(baseURL, "/") + + // 解析请求获取模型信息 + var claudeReq antigravity.ClaudeRequest + if err := json.Unmarshal(body, &claudeReq); err != nil { + return nil, fmt.Errorf("parse claude request: %w", err) + } + if strings.TrimSpace(claudeReq.Model) == "" { + return nil, fmt.Errorf("missing model") + } + originalModel := claudeReq.Model + + // 构建上游请求 URL + upstreamURL := baseURL + "/v1/messages" + + // 能力维度 sanitize:Anthropic-compatible 上游透传路径也需要保证 body↔beta header + // 对称。客户端 anthropic-beta header 不含 context-management-2025-06-27 但 body 带 + // context_management 时 strip,与 Anthropic 直连 / Bedrock / Vertex 路径保持一致。 + clientBeta := c.GetHeader("anthropic-beta") + if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { + body = sanitized + } + + // 创建请求 + req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create upstream request: %w", err) + } + + // 设置请求头 + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("x-api-key", apiKey) // Claude API 兼容 + + // 透传 Claude 相关 headers + if v := c.GetHeader("anthropic-version"); v != "" { + req.Header.Set("anthropic-version", v) + } + if v := clientBeta; v != "" { + req.Header.Set("anthropic-beta", v) + } + + // 代理 URL + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + // 发送请求 + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s upstream request failed: %v", prefix, err) + return nil, fmt.Errorf("upstream request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + // 处理错误响应 + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + + // 429 错误时标记账号限流 + if resp.StatusCode == http.StatusTooManyRequests { + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", false) + } + + // 透传上游错误 + c.Header("Content-Type", resp.Header.Get("Content-Type")) + c.Status(resp.StatusCode) + _, _ = c.Writer.Write(respBody) + + return &ForwardResult{ + Model: originalModel, + }, nil + } + + // 处理成功响应(流式/非流式) + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + + if claudeReq.Stream { + // 流式响应:透传 + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Status(http.StatusOK) + + streamRes := s.streamUpstreamResponse(c, resp, startTime) + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + } else { + // 非流式响应:直接透传 + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read upstream response: %w", err) + } + + // 提取 usage + usage = s.extractClaudeUsage(respBody) + + c.Header("Content-Type", resp.Header.Get("Content-Type")) + c.Status(http.StatusOK) + _, _ = c.Writer.Write(respBody) + } + + // 构建计费结果 + duration := time.Since(startTime) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=success duration_ms=%d", prefix, duration.Milliseconds()) + + return &ForwardResult{ + Model: originalModel, + Stream: claudeReq.Stream, + Duration: duration, + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + Usage: ClaudeUsage{ + InputTokens: usage.InputTokens, + OutputTokens: usage.OutputTokens, + CacheReadInputTokens: usage.CacheReadInputTokens, + CacheCreationInputTokens: usage.CacheCreationInputTokens, + }, + }, nil +} + +// streamUpstreamResponse 透传上游 SSE 流并提取 Claude usage +func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp *http.Response, startTime time.Time) *antigravityStreamResult { + usage := &ClaudeUsage{} + var firstTokenMs *int + + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanner.Buffer(make([]byte, 64*1024), maxLineSize) + + type scanEvent struct { + line string + err error + } + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func() { + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }() + defer close(done) + + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + // 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开 + keepaliveInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTicker *time.Ticker + if keepaliveInterval > 0 { + keepaliveTicker = time.NewTicker(keepaliveInterval) + defer keepaliveTicker.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTicker != nil { + keepaliveCh = keepaliveTicker.C + } + lastDataAt := time.Now() + + flusher, _ := c.Writer.(http.Flusher) + cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity upstream") + + for { + select { + case ev, ok := <-events: + if !ok { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()} + } + if ev.err != nil { + if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity upstream"); handled { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect} + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream read error (antigravity upstream): %v", ev.err) + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} + } + + lastDataAt = time.Now() + + line := ev.line + + // 记录首 token 时间 + if firstTokenMs == nil && len(line) > 0 { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + // 尝试从 message_delta 或 message_stop 事件提取 usage + s.extractSSEUsage(line, usage) + + // 透传行 + cw.Fprintf("%s\n", line) + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if cw.Disconnected() { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity upstream), returning collected usage") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true} + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity upstream)") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} + + case <-keepaliveCh: + if cw.Disconnected() { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + continue + } + // SSE ping 事件:Anthropic 原生格式,客户端会正确处理, + // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 + if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity upstream), continuing to drain upstream for billing") + continue + } + } + } +} + +// extractSSEUsage 从 SSE data 行中提取 Claude usage(用于流式透传场景) +// +// Anthropic streaming 的 usage 字段分布在两类事件中: +// - message_start:嵌套在 event.message.usage(input_tokens、cache_creation_input_tokens、 +// cache_read_input_tokens 等输入侧字段) +// - message_delta:位于顶层 event.usage(流结束时的最终 output_tokens) +// +// 仅读取顶层 event.usage 会漏掉 message_start 的输入侧字段,导致流式透传请求落库的 +// usage_logs 记录 input_tokens=0。 +func (s *AntigravityGatewayService) extractSSEUsage(line string, usage *ClaudeUsage) { + if !strings.HasPrefix(line, "data: ") { + return + } + dataStr := strings.TrimPrefix(line, "data: ") + var event map[string]any + if json.Unmarshal([]byte(dataStr), &event) != nil { + return + } + var u map[string]any + if eventType, _ := event["type"].(string); eventType == "message_start" { + if msg, ok := event["message"].(map[string]any); ok { + u, _ = msg["usage"].(map[string]any) + } + } else { + u, _ = event["usage"].(map[string]any) + } + if u == nil { + return + } + if v, ok := u["input_tokens"].(float64); ok && int(v) > 0 { + usage.InputTokens = int(v) + } + if v, ok := u["output_tokens"].(float64); ok && int(v) > 0 { + usage.OutputTokens = int(v) + } + if v, ok := u["cache_read_input_tokens"].(float64); ok && int(v) > 0 { + usage.CacheReadInputTokens = int(v) + } + if v, ok := u["cache_creation_input_tokens"].(float64); ok && int(v) > 0 { + usage.CacheCreationInputTokens = int(v) + } + // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 + if cc, ok := u["cache_creation"].(map[string]any); ok { + if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { + usage.CacheCreation5mTokens = int(v) + } + if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { + usage.CacheCreation1hTokens = int(v) + } + } +} + +// extractClaudeUsage 从非流式 Claude 响应提取 usage +func (s *AntigravityGatewayService) extractClaudeUsage(body []byte) *ClaudeUsage { + usage := &ClaudeUsage{} + var resp map[string]any + if json.Unmarshal(body, &resp) != nil { + return usage + } + if u, ok := resp["usage"].(map[string]any); ok { + if v, ok := u["input_tokens"].(float64); ok { + usage.InputTokens = int(v) + } + if v, ok := u["output_tokens"].(float64); ok { + usage.OutputTokens = int(v) + } + if v, ok := u["cache_read_input_tokens"].(float64); ok { + usage.CacheReadInputTokens = int(v) + } + if v, ok := u["cache_creation_input_tokens"].(float64); ok { + usage.CacheCreationInputTokens = int(v) + } + // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 + if cc, ok := u["cache_creation"].(map[string]any); ok { + if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { + usage.CacheCreation5mTokens = int(v) + } + if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { + usage.CacheCreation1hTokens = int(v) + } + } + } + return usage +} diff --git a/backend/internal/service/api_key.go b/backend/internal/service/api_key.go index dfc3ec1c5a..b92a848184 100644 --- a/backend/internal/service/api_key.go +++ b/backend/internal/service/api_key.go @@ -40,6 +40,7 @@ type APIKey struct { CompiledIPWhitelist *ip.CompiledIPRules `json:"-"` CompiledIPBlacklist *ip.CompiledIPRules `json:"-"` LastUsedAt *time.Time + LastUsedIP *string CreatedAt time.Time UpdatedAt time.Time User *User diff --git a/backend/internal/service/api_key_auth_cache.go b/backend/internal/service/api_key_auth_cache.go index 32c3910c9d..11b5246a1d 100644 --- a/backend/internal/service/api_key_auth_cache.go +++ b/backend/internal/service/api_key_auth_cache.go @@ -67,11 +67,17 @@ type APIKeyAuthGroupSnapshot struct { WeeklyLimitUSD *float64 `json:"weekly_limit_usd,omitempty"` MonthlyLimitUSD *float64 `json:"monthly_limit_usd,omitempty"` AllowImageGeneration bool `json:"allow_image_generation"` + AllowBatchImageGeneration bool `json:"allow_batch_image_generation"` ImageRateIndependent bool `json:"image_rate_independent"` ImageRateMultiplier float64 `json:"image_rate_multiplier"` ImagePrice1K *float64 `json:"image_price_1k,omitempty"` ImagePrice2K *float64 `json:"image_price_2k,omitempty"` ImagePrice4K *float64 `json:"image_price_4k,omitempty"` + VideoRateIndependent bool `json:"video_rate_independent"` + VideoRateMultiplier float64 `json:"video_rate_multiplier"` + VideoPrice480P *float64 `json:"video_price_480p,omitempty"` + VideoPrice720P *float64 `json:"video_price_720p,omitempty"` + VideoPrice1080P *float64 `json:"video_price_1080p,omitempty"` ClaudeCodeOnly bool `json:"claude_code_only"` FallbackGroupID *int64 `json:"fallback_group_id,omitempty"` FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request,omitempty"` diff --git a/backend/internal/service/api_key_auth_cache_impl.go b/backend/internal/service/api_key_auth_cache_impl.go index b5aedf271e..539c7375d9 100644 --- a/backend/internal/service/api_key_auth_cache_impl.go +++ b/backend/internal/service/api_key_auth_cache_impl.go @@ -14,7 +14,7 @@ import ( "github.com/dgraph-io/ristretto" ) -const apiKeyAuthSnapshotVersion = 13 // v13: include group peak rate fields +const apiKeyAuthSnapshotVersion = 14 // v14: include group video pricing fields type apiKeyAuthCacheConfig struct { l1Size int @@ -259,11 +259,17 @@ func (s *APIKeyService) snapshotFromAPIKey(ctx context.Context, apiKey *APIKey) WeeklyLimitUSD: apiKey.Group.WeeklyLimitUSD, MonthlyLimitUSD: apiKey.Group.MonthlyLimitUSD, AllowImageGeneration: apiKey.Group.AllowImageGeneration, + AllowBatchImageGeneration: apiKey.Group.AllowBatchImageGeneration, ImageRateIndependent: apiKey.Group.ImageRateIndependent, ImageRateMultiplier: apiKey.Group.ImageRateMultiplier, ImagePrice1K: apiKey.Group.ImagePrice1K, ImagePrice2K: apiKey.Group.ImagePrice2K, ImagePrice4K: apiKey.Group.ImagePrice4K, + VideoRateIndependent: apiKey.Group.VideoRateIndependent, + VideoRateMultiplier: apiKey.Group.VideoRateMultiplier, + VideoPrice480P: apiKey.Group.VideoPrice480P, + VideoPrice720P: apiKey.Group.VideoPrice720P, + VideoPrice1080P: apiKey.Group.VideoPrice1080P, ClaudeCodeOnly: apiKey.Group.ClaudeCodeOnly, FallbackGroupID: apiKey.Group.FallbackGroupID, FallbackGroupIDOnInvalidRequest: apiKey.Group.FallbackGroupIDOnInvalidRequest, @@ -336,11 +342,17 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho WeeklyLimitUSD: snapshot.Group.WeeklyLimitUSD, MonthlyLimitUSD: snapshot.Group.MonthlyLimitUSD, AllowImageGeneration: snapshot.Group.AllowImageGeneration, + AllowBatchImageGeneration: snapshot.Group.AllowBatchImageGeneration, ImageRateIndependent: snapshot.Group.ImageRateIndependent, ImageRateMultiplier: snapshot.Group.ImageRateMultiplier, ImagePrice1K: snapshot.Group.ImagePrice1K, ImagePrice2K: snapshot.Group.ImagePrice2K, ImagePrice4K: snapshot.Group.ImagePrice4K, + VideoRateIndependent: snapshot.Group.VideoRateIndependent, + VideoRateMultiplier: snapshot.Group.VideoRateMultiplier, + VideoPrice480P: snapshot.Group.VideoPrice480P, + VideoPrice720P: snapshot.Group.VideoPrice720P, + VideoPrice1080P: snapshot.Group.VideoPrice1080P, ClaudeCodeOnly: snapshot.Group.ClaudeCodeOnly, FallbackGroupID: snapshot.Group.FallbackGroupID, FallbackGroupIDOnInvalidRequest: snapshot.Group.FallbackGroupIDOnInvalidRequest, diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index 8903be65ee..b13ca6b3d0 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -6,6 +6,7 @@ import ( "encoding/hex" "fmt" "html" + "sort" "strconv" "strings" "sync" @@ -40,8 +41,9 @@ var ( ) const ( - apiKeyMaxErrorsPerHour = 20 - apiKeyLastUsedMinTouch = 30 * time.Second + apiKeyMaxErrorsPerHour = 20 + apiKeyLastUsedMinTouch = 30 * time.Second + apiKeySortCurrentConcurrency = "current_concurrency" // DB 写失败后的短退避,避免请求路径持续同步重试造成写风暴与高延迟。 apiKeyLastUsedFailBackoff = 5 * time.Second ) @@ -82,6 +84,10 @@ type APIKeyRepository interface { GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error) } +type apiKeyAllByUserIDLister interface { + ListAllByUserID(ctx context.Context, userID int64, filters APIKeyListFilters) ([]APIKey, error) +} + // APIKeyRateLimitData holds rate limit usage and window state for an API key. type APIKeyRateLimitData struct { Usage5h float64 @@ -437,6 +443,10 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK // List 获取用户的API Key列表 func (s *APIKeyService) List(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) { + if normalizedAPIKeySortBy(params.SortBy) == apiKeySortCurrentConcurrency { + return s.listByCurrentConcurrency(ctx, userID, params, filters) + } + keys, pagination, err := s.apiKeyRepo.ListByUserID(ctx, userID, params, filters) if err != nil { return nil, nil, fmt.Errorf("list api keys: %w", err) @@ -445,6 +455,75 @@ func (s *APIKeyService) List(ctx context.Context, userID int64, params paginatio return keys, pagination, nil } +func (s *APIKeyService) listByCurrentConcurrency(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) { + repo, ok := s.apiKeyRepo.(apiKeyAllByUserIDLister) + if !ok { + return nil, nil, fmt.Errorf("list api keys by current concurrency: repository does not support unpaginated API key listing") + } + + keys, err := repo.ListAllByUserID(ctx, userID, filters) + if err != nil { + return nil, nil, fmt.Errorf("list api keys: %w", err) + } + s.fillCurrentConcurrency(ctx, keys) + sortAPIKeysByCurrentConcurrency(keys, params.NormalizedSortOrder(pagination.SortOrderDesc)) + return paginateAPIKeys(keys, params), apiKeyPaginationResult(int64(len(keys)), params), nil +} + +func normalizedAPIKeySortBy(sortBy string) string { + return strings.ToLower(strings.TrimSpace(sortBy)) +} + +func sortAPIKeysByCurrentConcurrency(keys []APIKey, sortOrder string) { + desc := sortOrder != pagination.SortOrderAsc + sort.SliceStable(keys, func(i, j int) bool { + if keys[i].CurrentConcurrency == keys[j].CurrentConcurrency { + if desc { + return keys[i].ID > keys[j].ID + } + return keys[i].ID < keys[j].ID + } + if desc { + return keys[i].CurrentConcurrency > keys[j].CurrentConcurrency + } + return keys[i].CurrentConcurrency < keys[j].CurrentConcurrency + }) +} + +func paginateAPIKeys(keys []APIKey, params pagination.PaginationParams) []APIKey { + if len(keys) == 0 { + return []APIKey{} + } + limit := params.Limit() + page := params.Page + if page < 1 { + page = 1 + } + offset := (page - 1) * limit + if offset >= len(keys) { + return []APIKey{} + } + end := offset + limit + if end > len(keys) { + end = len(keys) + } + return keys[offset:end] +} + +func apiKeyPaginationResult(total int64, params pagination.PaginationParams) *pagination.PaginationResult { + limit := params.Limit() + pages := int(total) / limit + if int(total)%limit > 0 { + pages++ + } + return &pagination.PaginationResult{ + Total: total, + Page: params.Page, + PageSize: limit, + Pages: pages, + } +} + func (s *APIKeyService) fillCurrentConcurrency(ctx context.Context, keys []APIKey) { if s == nil || s.concurrencyService == nil || len(keys) == 0 { return diff --git a/backend/internal/service/api_key_service_delete_test.go b/backend/internal/service/api_key_service_delete_test.go index 25ad1edb15..a6ac75d60e 100644 --- a/backend/internal/service/api_key_service_delete_test.go +++ b/backend/internal/service/api_key_service_delete_test.go @@ -9,6 +9,7 @@ package service import ( "context" "errors" + "strings" "testing" "time" @@ -24,20 +25,26 @@ import ( // - deleteErr: 模拟 Delete 返回的错误 // - deletedIDs: 记录被调用删除的 API Key ID,用于断言验证 type apiKeyRepoStub struct { - apiKey *APIKey // GetKeyAndOwnerID 的返回值 - getByIDErr error // GetKeyAndOwnerID 的错误返回值 - deleteErr error // Delete 的错误返回值 - updateErr error // Update 的错误返回值 - deletedIDs []int64 // 记录已删除的 API Key ID 列表 - updatedKeys []APIKey - allowListByUserID bool - listByUserIDKeys []APIKey - listByUserIDErr error - listByUserIDCalls []int64 - listByUserIDParams []pagination.PaginationParams - updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error - touchedIDs []int64 - touchedUsedAts []time.Time + apiKey *APIKey // GetKeyAndOwnerID 的返回值 + getByIDErr error // GetKeyAndOwnerID 的错误返回值 + deleteErr error // Delete 的错误返回值 + updateErr error // Update 的错误返回值 + deletedIDs []int64 // 记录已删除的 API Key ID 列表 + updatedKeys []APIKey + allowListByUserID bool + listByUserIDKeys []APIKey + listByUserIDErr error + listByUserIDCalls []int64 + listByUserIDParams []pagination.PaginationParams + listByUserIDFilters []APIKeyListFilters + allowListAllByUserID bool + listAllByUserIDKeys []APIKey + listAllByUserIDErr error + listAllByUserIDCalls []int64 + listAllByUserIDFilters []APIKeyListFilters + updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error + touchedIDs []int64 + touchedUsedAts []time.Time } // 以下方法在本测试中不应被调用,使用 panic 确保测试失败时能快速定位问题 @@ -103,6 +110,7 @@ func (s *apiKeyRepoStub) ListByUserID(ctx context.Context, userID int64, params } s.listByUserIDCalls = append(s.listByUserIDCalls, userID) s.listByUserIDParams = append(s.listByUserIDParams, params) + s.listByUserIDFilters = append(s.listByUserIDFilters, filters) if s.listByUserIDErr != nil { return nil, nil, s.listByUserIDErr } @@ -115,6 +123,51 @@ func (s *apiKeyRepoStub) ListByUserID(ctx context.Context, userID int64, params }, nil } +func (s *apiKeyRepoStub) ListAllByUserID(ctx context.Context, userID int64, filters APIKeyListFilters) ([]APIKey, error) { + if !s.allowListAllByUserID { + panic("unexpected ListAllByUserID call") + } + s.listAllByUserIDCalls = append(s.listAllByUserIDCalls, userID) + s.listAllByUserIDFilters = append(s.listAllByUserIDFilters, filters) + if s.listAllByUserIDErr != nil { + return nil, s.listAllByUserIDErr + } + source := s.listByUserIDKeys + if s.listAllByUserIDKeys != nil { + source = s.listAllByUserIDKeys + } + return filterAPIKeyStubKeys(userID, source, filters), nil +} + +func filterAPIKeyStubKeys(userID int64, keys []APIKey, filters APIKeyListFilters) []APIKey { + result := make([]APIKey, 0, len(keys)) + search := strings.ToLower(filters.Search) + for _, key := range keys { + if key.UserID != userID { + continue + } + if search != "" && + !strings.Contains(strings.ToLower(key.Name), search) && + !strings.Contains(strings.ToLower(key.Key), search) { + continue + } + if filters.Status != "" && key.Status != filters.Status { + continue + } + if filters.GroupID != nil { + if *filters.GroupID == 0 { + if key.GroupID != nil { + continue + } + } else if key.GroupID == nil || *key.GroupID != *filters.GroupID { + continue + } + } + result = append(result, key) + } + return result +} + func (s *apiKeyRepoStub) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) { panic("unexpected VerifyOwnership call") } @@ -320,6 +373,96 @@ func TestAPIKeyService_List_FillsCurrentConcurrency(t *testing.T) { require.Equal(t, 0, keys[1].CurrentConcurrency) } +func TestAPIKeyService_List_SortByCurrentConcurrency(t *testing.T) { + groupID := int64(42) + keys := []APIKey{ + {ID: 1, UserID: 7, Key: "sk-target-1", Name: "target-one", GroupID: &groupID, Status: StatusActive}, + {ID: 2, UserID: 7, Key: "sk-target-2", Name: "target-two", GroupID: &groupID, Status: StatusActive}, + {ID: 3, UserID: 7, Key: "sk-target-3", Name: "target-three", GroupID: &groupID, Status: StatusActive}, + {ID: 4, UserID: 7, Key: "sk-target-4", Name: "target-four", GroupID: &groupID, Status: StatusActive}, + {ID: 9, UserID: 7, Key: "sk-target-9", Name: "target-inactive", GroupID: &groupID, Status: StatusDisabled}, + {ID: 10, UserID: 7, Key: "sk-other-10", Name: "other", GroupID: &groupID, Status: StatusActive}, + {ID: 11, UserID: 7, Key: "sk-target-11", Name: "target-no-group", Status: StatusActive}, + {ID: 12, UserID: 8, Key: "sk-target-12", Name: "target-other-user", GroupID: &groupID, Status: StatusActive}, + } + filters := APIKeyListFilters{ + Search: "target", + Status: StatusActive, + GroupID: &groupID, + } + repo := &apiKeyRepoStub{ + allowListAllByUserID: true, + listAllByUserIDKeys: keys, + } + concurrency := NewConcurrencyService(&stubConcurrencyCacheForTest{ + apiKeyConcurrency: map[int64]int{ + 1: 5, + 2: 5, + 3: 2, + 4: 8, + 9: 99, + 10: 99, + 11: 99, + 12: 99, + }, + }) + svc := &APIKeyService{apiKeyRepo: repo, concurrencyService: concurrency} + + got, page, err := svc.List(context.Background(), 7, pagination.PaginationParams{ + Page: 2, + PageSize: 2, + SortBy: "current_concurrency", + SortOrder: "desc", + }, filters) + require.NoError(t, err) + require.Equal(t, []int64{1, 3}, apiKeyTestIDs(got)) + require.Equal(t, int64(4), page.Total) + require.Equal(t, 2, page.Page) + require.Equal(t, 2, page.PageSize) + require.Equal(t, 2, page.Pages) + require.Empty(t, repo.listByUserIDCalls) + require.Equal(t, []int64{7}, repo.listAllByUserIDCalls) + require.Len(t, repo.listAllByUserIDFilters, 1) + require.Equal(t, filters.Search, repo.listAllByUserIDFilters[0].Search) + require.Equal(t, filters.Status, repo.listAllByUserIDFilters[0].Status) + require.NotNil(t, repo.listAllByUserIDFilters[0].GroupID) + require.Equal(t, groupID, *repo.listAllByUserIDFilters[0].GroupID) +} + +func TestAPIKeyService_List_SortByCurrentConcurrencyAscTiesByID(t *testing.T) { + repo := &apiKeyRepoStub{ + allowListAllByUserID: true, + listAllByUserIDKeys: []APIKey{ + {ID: 1, UserID: 7, Key: "sk-1", Name: "one", Status: StatusActive}, + {ID: 2, UserID: 7, Key: "sk-2", Name: "two", Status: StatusActive}, + {ID: 3, UserID: 7, Key: "sk-3", Name: "three", Status: StatusActive}, + {ID: 4, UserID: 7, Key: "sk-4", Name: "four", Status: StatusActive}, + }, + } + concurrency := NewConcurrencyService(&stubConcurrencyCacheForTest{ + apiKeyConcurrency: map[int64]int{1: 5, 2: 5, 3: 2, 4: 8}, + }) + svc := &APIKeyService{apiKeyRepo: repo, concurrencyService: concurrency} + + got, page, err := svc.List(context.Background(), 7, pagination.PaginationParams{ + Page: 1, + PageSize: 4, + SortBy: "current_concurrency", + SortOrder: "asc", + }, APIKeyListFilters{}) + require.NoError(t, err) + require.Equal(t, []int64{3, 1, 2, 4}, apiKeyTestIDs(got)) + require.Equal(t, 4, page.PageSize) +} + +func apiKeyTestIDs(keys []APIKey) []int64 { + ids := make([]int64, 0, len(keys)) + for _, key := range keys { + ids = append(ids, key.ID) + } + return ids +} + func TestAPIKeyService_GetByID_FillsCurrentConcurrency(t *testing.T) { repo := &apiKeyRepoStub{ apiKey: &APIKey{ID: 10, UserID: 7, Key: "sk-10", Name: "key-10"}, diff --git a/backend/internal/service/batch_image.go b/backend/internal/service/batch_image.go new file mode 100644 index 0000000000..af0206b8b0 --- /dev/null +++ b/backend/internal/service/batch_image.go @@ -0,0 +1,414 @@ +package service + +import ( + "context" + "crypto/rand" + "encoding/hex" + "net/http" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +const ( + BatchImageProviderGeminiAPI = "gemini_api" + BatchImageProviderVertex = "vertex" +) + +const ( + BatchImageJobStatusCreated = "created" + BatchImageJobStatusUploading = "uploading" + BatchImageJobStatusSubmitted = "submitted" + BatchImageJobStatusRunning = "running" + BatchImageJobStatusIndexing = "indexing" + BatchImageJobStatusSettling = "settling" + BatchImageJobStatusCompleted = "completed" + BatchImageJobStatusFailed = "failed" + BatchImageJobStatusCancelled = "cancelled" + BatchImageJobStatusOutputDeleted = "output_deleted" +) + +const ( + BatchImageItemStatusPending = "pending" + BatchImageItemStatusSuccess = "success" + BatchImageItemStatusFailed = "failed" + BatchImageItemStatusCancelled = "cancelled" +) + +var ( + ErrBatchImageJobNotFound = infraerrors.New(http.StatusNotFound, "BATCH_IMAGE_JOB_NOT_FOUND", "batch image job not found") + ErrBatchImageJobExists = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_JOB_EXISTS", "batch image job already exists") + ErrBatchImageItemExists = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_ITEM_EXISTS", "batch image item already exists") + + ErrBatchImageInvalidTransition = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_INVALID_TRANSITION", "invalid batch image job status transition") + ErrBatchImageInvalidProvider = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_INVALID_PROVIDER", "invalid batch image provider") + + ErrBatchImageMissingProviderJobName = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_MISSING_PROVIDER_JOB_NAME", "batch image provider job name is missing") + ErrBatchImageMissingAccountID = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_MISSING_ACCOUNT_ID", "batch image account id is missing") + ErrBatchImageUnsupportedProvider = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_UNSUPPORTED_PROVIDER", "unsupported batch image provider") + ErrBatchImageIndexOutputMissing = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_INDEX_OUTPUT_MISSING", "batch image provider output is missing") + ErrBatchImageIndexParseFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_INDEX_PARSE_FAILED", "batch image provider output parse failed") + ErrBatchImageIndexNoResultLines = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_INDEX_NO_RESULT_LINES", "batch image provider output has no result lines") + ErrBatchImageDuplicateCustomID = infraerrors.New(http.StatusBadGateway, "DUPLICATE_CUSTOM_ID_IN_OUTPUT", "batch image provider output contains duplicate custom id") + ErrBatchImageIndexStateConflict = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_INDEX_STATE_CONFLICT", "batch image job is no longer in indexing state") + + ErrBatchImageSettlementInvalidStatus = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_SETTLEMENT_INVALID_STATUS", "batch image job is not ready for settlement") + ErrBatchImageSettlementManifestConflict = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_SETTLEMENT_MANIFEST_CONFLICT", "batch image settlement manifest hash conflict") + ErrBatchImageSettlementPricingMissing = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_SETTLEMENT_PRICING_MISSING", "batch image settlement pricing is missing") + ErrBatchImageSettlementBillingFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_SETTLEMENT_BILLING_FAILED", "batch image settlement billing failed") + ErrBatchImageAlreadySettled = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_ALREADY_SETTLED", "batch image job is already settled") + ErrBatchImageSettlementMissingAPIKeyID = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_SETTLEMENT_MISSING_API_KEY_ID", "batch image settlement api key id is missing") + ErrBatchImageSettlementMissingAccountID = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_SETTLEMENT_MISSING_ACCOUNT_ID", "batch image settlement account id is missing") + ErrBatchImageSettlementInvalidCounts = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_SETTLEMENT_INVALID_COUNTS", "batch image settlement counts are invalid") + ErrBatchImageSettlementCostExceedsHold = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_SETTLEMENT_COST_EXCEEDS_HOLD", "batch image settlement cost exceeds held balance") + ErrBatchImageBillingHoldFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_BILLING_HOLD_FAILED", "batch image balance hold failed") + ErrBatchImageInsufficientBalance = infraerrors.New(http.StatusPaymentRequired, "BATCH_IMAGE_INSUFFICIENT_BALANCE", "insufficient balance for batch image hold") + + ErrBatchImageDisabled = infraerrors.New(http.StatusNotFound, "BATCH_IMAGE_DISABLED", "batch image API is disabled") + ErrBatchImageGroupDisabled = infraerrors.New(http.StatusForbidden, "BATCH_IMAGE_GROUP_DISABLED", "batch image API is disabled for this group") + ErrBatchImageInvalidModel = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_INVALID_MODEL", "batch image model is required") + ErrBatchImageNoAccountAvailable = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_NO_ACCOUNT_AVAILABLE", "no compatible batch image account is available") + ErrBatchImageInvalidItems = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_INVALID_ITEMS", "batch image items are invalid") + ErrBatchImageDuplicateCustomIDInRequest = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_DUPLICATE_CUSTOM_ID", "batch image custom ids must be unique") + ErrBatchImagePromptTooLong = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROMPT_TOO_LONG", "batch image prompt is too long") + ErrBatchImageInvalidReferenceImage = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_INVALID_REFERENCE_IMAGE", "batch image reference image is invalid") + ErrBatchImageTooManyReferenceImages = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_TOO_MANY_REFERENCE_IMAGES", "too many batch image reference images for this model") + ErrBatchImageReferenceImagesTooLarge = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_REFERENCE_IMAGES_TOO_LARGE", "batch image reference images are too large") + ErrBatchImageTooManyOutputImages = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_TOO_MANY_OUTPUT_IMAGES", "too many batch image output images") + ErrBatchImageProviderSubmitFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_PROVIDER_SUBMIT_FAILED", "batch image provider submit failed") + ErrBatchImageQueueFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_QUEUE_FAILED", "batch image queue failed") + ErrBatchImageIdempotencyConflict = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_IDEMPOTENCY_CONFLICT", "idempotency key reused with different batch image request") + ErrBatchImageCancelFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_CANCEL_FAILED", "batch image cancel failed") + ErrBatchImageVertexGCSBucketMissing = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_VERTEX_GCS_BUCKET_MISSING", "Vertex managed GCS bucket is not configured") + + ErrBatchImageNotReady = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_NOT_READY", "batch image job is not completed") + ErrBatchImageOutputDeleted = infraerrors.New(http.StatusGone, "BATCH_IMAGE_OUTPUT_DELETED", "batch image output has been deleted") + ErrBatchImageItemNotFound = infraerrors.New(http.StatusNotFound, "BATCH_IMAGE_ITEM_NOT_FOUND", "batch image item not found") + ErrBatchImageItemFailed = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_ITEM_FAILED", "batch image item did not succeed") + ErrBatchImageResultMissing = infraerrors.New(http.StatusInternalServerError, "BATCH_IMAGE_RESULT_MISSING", "batch image result is missing") + ErrBatchImageDownloadLimited = infraerrors.New(http.StatusTooManyRequests, "BATCH_IMAGE_DOWNLOAD_LIMITED", "too many batch image downloads") + ErrBatchImageDownloadFailed = infraerrors.New(http.StatusInternalServerError, "BATCH_IMAGE_DOWNLOAD_FAILED", "batch image download failed") + ErrBatchImageDownloadTooLarge = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_DOWNLOAD_TOO_LARGE", "batch image download is too large") + ErrBatchImageItemImageIndexOutOfRange = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_ITEM_IMAGE_INDEX_OUT_OF_RANGE", "batch image item image index is out of range") + ErrBatchImageZipTooManyItems = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_ZIP_TOO_MANY_ITEMS", "batch image ZIP contains too many items; use single item downloads") + ErrBatchImageOutputDeleteNotReady = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_OUTPUT_DELETE_NOT_READY", "batch image output can only be deleted after completion") + ErrBatchImageRecordDeleteNotReady = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_RECORD_DELETE_NOT_READY", "batch image record can only be deleted after the job finishes") + ErrBatchImageCleanupFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_CLEANUP_FAILED", "batch image cleanup failed") + ErrBatchImageCleanupUnsafePath = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_CLEANUP_UNSAFE_PATH", "batch image cleanup path is unsafe") + ErrBatchImageProviderCleanupFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_PROVIDER_CLEANUP_FAILED", "batch image provider cleanup failed") +) + +type BatchImageJob struct { + ID int64 + BatchID string + UserID int64 + APIKeyID *int64 + AccountID *int64 + Provider string + Model string + TaskName string + ParentBatchID *string + Status string + ProviderJobName *string + ProviderInputRef *string + ProviderOutputRef *string + GCSInputURI *string + GCSOutputURI *string + + ItemCount int + SuccessCount int + FailCount int + CancelledCount int + + EstimatedCost float64 + HoldAmount *float64 + ActualCost *float64 + BaseUnitPrice float64 + GroupRateMultiplier float64 + AccountRateMultiplier float64 + BatchDiscountMultiplier float64 + HoldMultiplier float64 + BillableUnitPrice float64 + HoldUnitPrice float64 + PricingSnapshotVersion int + Currency string + HoldID *string + + IdempotencyKey *string + RequestHash *string + ManifestHash *string + + RetryCount int + Version int + + OutputExpiresAt *time.Time + InputDeletedAt *time.Time + OutputDeletedAt *time.Time + DownloadedAt *time.Time + UserDeletedAt *time.Time + + LastErrorCode *string + LastErrorMessage *string + + CreatedAt time.Time + UpdatedAt time.Time + SubmittedAt *time.Time + StartedAt *time.Time + FinishedAt *time.Time + SettledAt *time.Time +} + +type CreateBatchImageJobParams struct { + BatchID string + UserID int64 + APIKeyID *int64 + AccountID *int64 + Provider string + Model string + TaskName string + ParentBatchID *string + Status string + ProviderJobName *string + ProviderInputRef *string + ProviderOutputRef *string + GCSInputURI *string + GCSOutputURI *string + + ItemCount int + SuccessCount int + FailCount int + CancelledCount int + + EstimatedCost float64 + HoldAmount *float64 + ActualCost *float64 + BaseUnitPrice float64 + GroupRateMultiplier float64 + AccountRateMultiplier float64 + BatchDiscountMultiplier float64 + HoldMultiplier float64 + BillableUnitPrice float64 + HoldUnitPrice float64 + PricingSnapshotVersion int + Currency string + HoldID *string + + IdempotencyKey *string + RequestHash *string + ManifestHash *string + + RetryCount int + + OutputExpiresAt *time.Time +} + +type BatchImageItem struct { + ID int64 + JobID string + CustomID string + Status string + RequestHash *string + PromptPreview *string + ProviderSourceObject *string + SourceLineNumber *int + SourceByteOffset *int64 + SourceByteLength *int64 + MimeType *string + FileExtension *string + ImageCount int + ErrorCode *string + ErrorMessage *string + BilledAmount *float64 + CreatedAt time.Time + IndexedAt *time.Time +} + +type CreateBatchImageItemParams struct { + JobID string + CustomID string + Status string + RequestHash *string + PromptPreview *string + ProviderSourceObject *string + SourceLineNumber *int + SourceByteOffset *int64 + SourceByteLength *int64 + MimeType *string + FileExtension *string + ImageCount int + ErrorCode *string + ErrorMessage *string + BilledAmount *float64 + IndexedAt *time.Time +} + +type BatchImageItemFilter struct { + Status string + Limit int + Offset int +} + +type BatchImageJobFilter struct { + Status string + TaskNameLike string + Downloaded *bool + CreatedAfter *time.Time + CreatedBefore *time.Time + ExcludeDeleted bool + Limit int + Offset int +} + +type BatchImageCounts struct { + SuccessCount int + FailCount int +} + +type UpdateBatchImageJobProviderSubmitParams struct { + BatchID string + ProviderJobName string + ProviderInputRef string + ProviderOutputRef string + GCSInputURI string + GCSOutputURI string + EventPayload any +} + +type BatchImageTransitionOptions struct { + EventType string + EventPayload any + ErrorCode *string + ErrorMessage *string + Now *time.Time +} + +type MarkBatchImageJobSettledParams struct { + BatchID string + ActualCost float64 + ManifestHash string + EventPayload any + Now *time.Time + OutputExpiresAt *time.Time +} + +type BatchImageEvent struct { + ID int64 + JobID string + EventType string + Payload []byte + EventHash *string + CreatedAt time.Time +} + +type BatchImageRepository interface { + CreateBatchImageJob(ctx context.Context, params CreateBatchImageJobParams) (*BatchImageJob, error) + GetBatchImageJobByBatchID(ctx context.Context, batchID string) (*BatchImageJob, error) + GetBatchImageJobByIdempotencyKey(ctx context.Context, userID, apiKeyID int64, key string) (*BatchImageJob, error) + GetBatchImageJobByBatchIDForOwner(ctx context.Context, userID, apiKeyID int64, batchID string) (*BatchImageJob, error) + GetBatchImageJobByID(ctx context.Context, id int64) (*BatchImageJob, error) + ListBatchImageJobsForOwner(ctx context.Context, userID, apiKeyID int64, filter BatchImageJobFilter) ([]*BatchImageJob, error) + TransitionBatchImageJobStatus(ctx context.Context, batchID, toStatus string, opts BatchImageTransitionOptions) error + // TouchBatchImageJobSubmitting 刷新未提交(created/uploading)job 的 updated_at, + // 作为慢提交期间的心跳,防止被 stale 恢复扫描误杀。 + TouchBatchImageJobSubmitting(ctx context.Context, batchID string) error + // FailStaleUnsubmittedBatchImageJob 原子地将仍处于 created/uploading 且 + // provider_job_name 为空、updated_at 早于 cutoff 的 job 转为 failed。 + // 返回 false 表示 job 已被并发推进(如已提交成功),调用方不得释放冻结。 + FailStaleUnsubmittedBatchImageJob(ctx context.Context, batchID string, cutoff time.Time, code, message string) (bool, error) + UpdateBatchImageJobProviderOutputRef(ctx context.Context, batchID, providerOutputRef string) error + UpdateBatchImageJobProviderSubmit(ctx context.Context, params UpdateBatchImageJobProviderSubmitParams) error + RecordBatchImageJobSubmitFailure(ctx context.Context, batchID, code, message string, markFailed bool) error + MarkBatchImageJobSettled(ctx context.Context, params MarkBatchImageJobSettledParams) error + SetBatchImageJobSettlementFailed(ctx context.Context, batchID, code, message string) (int, error) + CreateBatchImageItem(ctx context.Context, params CreateBatchImageItemParams) (*BatchImageItem, error) + BulkCreateBatchImageItems(ctx context.Context, params []CreateBatchImageItemParams) error + ReplaceBatchImageItemsForJob(ctx context.Context, batchID string, items []CreateBatchImageItemParams, counts BatchImageCounts) error + ListBatchImageItems(ctx context.Context, batchID string, filter BatchImageItemFilter) ([]*BatchImageItem, error) + ListBatchImageItemsForOwner(ctx context.Context, userID, apiKeyID int64, batchID string, filter BatchImageItemFilter) ([]*BatchImageItem, error) + GetBatchImageJobForDownload(ctx context.Context, userID, apiKeyID int64, batchID string) (*BatchImageJob, error) + GetBatchImageItemForDownload(ctx context.Context, batchID, customID string) (*BatchImageItem, error) + ListBatchImageItemsForDownload(ctx context.Context, batchID string, status string, limit int) ([]*BatchImageItem, error) + ListBatchImageJobsDueForInputCleanup(ctx context.Context, cutoff time.Time, limit int) ([]*BatchImageJob, error) + ListBatchImageJobsDueForOutputCleanup(ctx context.Context, now time.Time, limit int) ([]*BatchImageJob, error) + ListStaleUnsubmittedBatchImageJobs(ctx context.Context, cutoff time.Time, limit int) ([]*BatchImageJob, error) + MarkBatchImageInputDeleted(ctx context.Context, batchID string, deletedAt time.Time) error + MarkBatchImageOutputDeleted(ctx context.Context, batchID string, deletedAt time.Time) error + MarkBatchImageDownloaded(ctx context.Context, batchID string, downloadedAt time.Time) error + MarkBatchImageJobUserDeleted(ctx context.Context, userID, apiKeyID int64, batchID string, deletedAt time.Time) error + SetBatchImageOutputExpiresAt(ctx context.Context, batchID string, expiresAt time.Time) error + RecordBatchImageCleanupFailure(ctx context.Context, batchID, code, message string) error + AppendBatchImageEvent(ctx context.Context, batchID, eventType string, payload any) error +} + +func NewBatchImageID() (string, error) { + var b [16]byte + if _, err := rand.Read(b[:]); err != nil { + return "", err + } + return "imgbatch_" + hex.EncodeToString(b[:]), nil +} + +func IsSupportedBatchImageProvider(provider string) bool { + switch provider { + case BatchImageProviderGeminiAPI, BatchImageProviderVertex: + return true + default: + return false + } +} + +func IsTerminalBatchImageJobStatus(status string) bool { + switch status { + case BatchImageJobStatusCompleted, BatchImageJobStatusFailed, BatchImageJobStatusCancelled, BatchImageJobStatusOutputDeleted: + return true + default: + return false + } +} + +func CanTransitionBatchImageJob(from, to string) bool { + if from == "" || to == "" { + return false + } + if IsTerminalBatchImageJobStatus(from) { + return to == BatchImageJobStatusOutputDeleted && + from != BatchImageJobStatusOutputDeleted && + (from == BatchImageJobStatusCompleted || from == BatchImageJobStatusFailed || from == BatchImageJobStatusCancelled) + } + if to == BatchImageJobStatusFailed { + return true + } + + allowed := map[string]map[string]struct{}{ + BatchImageJobStatusCreated: { + BatchImageJobStatusUploading: {}, + BatchImageJobStatusSubmitted: {}, + BatchImageJobStatusCancelled: {}, + }, + BatchImageJobStatusUploading: { + BatchImageJobStatusSubmitted: {}, + BatchImageJobStatusCancelled: {}, + }, + BatchImageJobStatusSubmitted: { + BatchImageJobStatusRunning: {}, + BatchImageJobStatusIndexing: {}, + BatchImageJobStatusFailed: {}, + BatchImageJobStatusCancelled: {}, + }, + BatchImageJobStatusRunning: { + BatchImageJobStatusRunning: {}, + BatchImageJobStatusIndexing: {}, + BatchImageJobStatusFailed: {}, + BatchImageJobStatusCancelled: {}, + }, + BatchImageJobStatusIndexing: { + BatchImageJobStatusSettling: {}, + BatchImageJobStatusFailed: {}, + }, + BatchImageJobStatusSettling: { + BatchImageJobStatusCompleted: {}, + }, + } + _, ok := allowed[from][to] + return ok +} diff --git a/backend/internal/service/batch_image_billing_hold.go b/backend/internal/service/batch_image_billing_hold.go new file mode 100644 index 0000000000..80f77208ba --- /dev/null +++ b/backend/internal/service/batch_image_billing_hold.go @@ -0,0 +1,116 @@ +package service + +import ( + "context" + "errors" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + batchImageHoldRequestPrefix = "batch_image_hold:" + batchImageCaptureRequestPrefix = "batch_image_capture:" + batchImageReleaseRequestPrefix = "batch_image_release:" +) + +func BatchImageHoldRequestID(batchID string) string { + return batchImageHoldRequestPrefix + strings.TrimSpace(batchID) +} + +func BatchImageCaptureRequestID(batchID string) string { + return batchImageCaptureRequestPrefix + strings.TrimSpace(batchID) +} + +func BatchImageReleaseRequestID(batchID string) string { + return batchImageReleaseRequestPrefix + strings.TrimSpace(batchID) +} + +func buildBatchImageHoldCommand(job *BatchImageJob, requestID string, actualAmount float64, payloadHash string) (*BatchImageBalanceHoldCommand, error) { + if job == nil { + return nil, ErrBatchImageBillingHoldFailed + } + if job.APIKeyID == nil || *job.APIKeyID <= 0 { + return nil, ErrBatchImageSettlementMissingAPIKeyID + } + holdAmount := job.EstimatedCost + if job.HoldAmount != nil { + holdAmount = *job.HoldAmount + } + if holdAmount < 0 { + holdAmount = 0 + } + if actualAmount < 0 { + actualAmount = 0 + } + return &BatchImageBalanceHoldCommand{ + RequestID: requestID, + APIKeyID: *job.APIKeyID, + UserID: job.UserID, + BatchID: job.BatchID, + HoldAmount: holdAmount, + ActualAmount: actualAmount, + RequestPayloadHash: strings.TrimSpace(payloadHash), + }, nil +} + +func reserveBatchImageBalanceHold(ctx context.Context, repo UsageBillingRepository, job *BatchImageJob, payloadHash string) error { + if repo == nil { + return ErrBatchImageBillingHoldFailed.WithCause(errors.New("batch image billing repository is not configured")) + } + cmd, err := buildBatchImageHoldCommand(job, BatchImageHoldRequestID(job.BatchID), 0, payloadHash) + if err != nil { + return err + } + if cmd.HoldAmount <= 0 { + return nil + } + if _, err := repo.ReserveBatchImageBalance(ctx, cmd); err != nil { + if errors.Is(err, ErrBatchImageInsufficientBalance) { + return ErrBatchImageInsufficientBalance + } + return ErrBatchImageBillingHoldFailed.WithCause(err) + } + return nil +} + +func captureBatchImageBalanceHold(ctx context.Context, repo UsageBillingRepository, job *BatchImageJob, actualAmount float64, payloadHash string) error { + if repo == nil { + return ErrBatchImageSettlementBillingFailed.WithCause(errors.New("batch image billing repository is not configured")) + } + cmd, err := buildBatchImageHoldCommand(job, BatchImageCaptureRequestID(job.BatchID), actualAmount, payloadHash) + if err != nil { + return err + } + if _, err := repo.CaptureBatchImageBalance(ctx, cmd); err != nil { + return ErrBatchImageSettlementBillingFailed.WithCause(err) + } + return nil +} + +func releaseBatchImageBalanceHold(ctx context.Context, repo UsageBillingRepository, job *BatchImageJob, payloadHash string) error { + if repo == nil || job == nil { + return nil + } + cmd, err := buildBatchImageHoldCommand(job, BatchImageReleaseRequestID(job.BatchID), 0, payloadHash) + if err != nil { + return err + } + if cmd.HoldAmount <= 0 { + return nil + } + if _, err := repo.ReleaseBatchImageBalance(ctx, cmd); err != nil { + // 同一 release request id 出现指纹冲突,说明此前已有一次携带不同 + // payloadHash 的释放成功提交(资金已归还)。视为幂等成功, + // 避免历史指纹不一致的 job 永远卡在释放失败的毒消息循环里。 + if errors.Is(err, ErrUsageBillingRequestConflict) { + logger.L().Warn("batch_image.release_fingerprint_conflict_treated_as_released", + zap.String("batch_id", job.BatchID), + ) + return nil + } + return ErrBatchImageBillingHoldFailed.WithCause(err) + } + return nil +} diff --git a/backend/internal/service/batch_image_billing_recovery.go b/backend/internal/service/batch_image_billing_recovery.go new file mode 100644 index 0000000000..5987fec48e --- /dev/null +++ b/backend/internal/service/batch_image_billing_recovery.go @@ -0,0 +1,103 @@ +package service + +import ( + "context" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + defaultBatchImageBillingRecoveryStaleAfter = 10 * time.Minute + defaultBatchImageBillingRecoveryLimit = 100 +) + +type BatchImageBillingRecoveryService struct { + Repo BatchImageRepository + Billing UsageBillingRepository + AuthCache APIKeyAuthCacheInvalidator + Queue BatchImageQueue + StaleAfter time.Duration + Limit int +} + +func (s *BatchImageBillingRecoveryService) ReleaseStaleUnsubmittedOnce(ctx context.Context) (int, error) { + if s == nil || s.Repo == nil || s.Billing == nil { + return 0, nil + } + staleAfter := s.StaleAfter + if staleAfter <= 0 { + staleAfter = defaultBatchImageBillingRecoveryStaleAfter + } + limit := s.Limit + if limit <= 0 { + limit = defaultBatchImageBillingRecoveryLimit + } + cutoff := time.Now().Add(-staleAfter) + jobs, err := s.Repo.ListStaleUnsubmittedBatchImageJobs(ctx, cutoff, limit) + if err != nil { + return 0, err + } + released := 0 + var lastErr error + for _, job := range jobs { + if job == nil { + continue + } + if err := ctx.Err(); err != nil { + return released, err + } + msg := "batch image submission did not reach provider before recovery cutoff" + // 原子转 failed 并复核 stale 条件:List 与转态之间 job 可能已被慢提交 + // 心跳续期或提交成功(provider_job_name 已写入),此时绝不能退款, + // 否则上游任务照常产生成本而用户已拿回冻结余额。 + applied, err := s.Repo.FailStaleUnsubmittedBatchImageJob(ctx, job.BatchID, cutoff, "SUBMIT_STALE_BEFORE_PROVIDER", msg) + if err != nil { + // applied=true 时 UPDATE 已提交(仅审计事件写入失败):必须继续释放, + // 否则 job 已转 failed、不再出现在 stale 列表,冻结余额会永久泄漏。 + if !applied { + lastErr = err + continue + } + logger.L().Warn("batch_image.recovery_fail_event_append_failed", + zap.String("batch_id", job.BatchID), + zap.Error(err), + ) + } + if !applied { + continue + } + job.Status = BatchImageJobStatusFailed + if err := releaseBatchImageBalanceHold(ctx, s.Billing, job, batchImageDerefString(job.RequestHash)); err != nil { + // job 已转 failed、不会再进入 stale 列表:必须给释放失败留下 + // 自动重试路径(入队后由 worker 的 releaseTerminalHold 兜底), + // 否则冻结余额永久泄漏。 + logger.L().Warn("batch_image.recovery_release_failed", + zap.String("batch_id", job.BatchID), + zap.Error(err), + ) + s.enqueueReleaseRetry(ctx, job.BatchID) + lastErr = err + continue + } + if s.AuthCache != nil && job.UserID > 0 { + s.AuthCache.InvalidateAuthCacheByUserID(ctx, job.UserID) + } + released++ + } + return released, lastErr +} + +func (s *BatchImageBillingRecoveryService) enqueueReleaseRetry(ctx context.Context, batchID string) { + if s == nil || s.Queue == nil { + return + } + if err := s.Queue.Enqueue(ctx, batchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) { + logger.L().Warn("batch_image.recovery_release_retry_enqueue_failed", + zap.String("batch_id", batchID), + zap.Error(err), + ) + } +} diff --git a/backend/internal/service/batch_image_billing_recovery_test.go b/backend/internal/service/batch_image_billing_recovery_test.go new file mode 100644 index 0000000000..8ddadfff29 --- /dev/null +++ b/backend/internal/service/batch_image_billing_recovery_test.go @@ -0,0 +1,117 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type recordingBatchImageQueue struct { + *fakeBatchImageQueue + enqueued []string +} + +func (q *recordingBatchImageQueue) Enqueue(_ context.Context, batchID string) error { + q.enqueued = append(q.enqueued, batchID) + return nil +} + +func TestBatchImageBillingRecoveryService_ReleasesStaleUnsubmittedHold(t *testing.T) { + repo := newFakeBatchImageRepository() + apiKeyID := int64(22) + holdAmount := 0.5 + stale := &BatchImageJob{ + BatchID: "imgbatch_stale_created", + UserID: 11, + APIKeyID: &apiKeyID, + Status: BatchImageJobStatusCreated, + EstimatedCost: holdAmount, + HoldAmount: &holdAmount, + CreatedAt: time.Now().Add(-time.Hour), + UpdatedAt: time.Now().Add(-time.Hour), + } + activeProviderName := "providers/job" + active := &BatchImageJob{ + BatchID: "imgbatch_has_provider", + UserID: 11, + APIKeyID: &apiKeyID, + Status: BatchImageJobStatusSubmitted, + ProviderJobName: &activeProviderName, + EstimatedCost: holdAmount, + HoldAmount: &holdAmount, + CreatedAt: time.Now().Add(-time.Hour), + UpdatedAt: time.Now().Add(-time.Hour), + } + repo.jobs[stale.BatchID] = stale + repo.jobs[active.BatchID] = active + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, StaleAfter: time.Minute, Limit: 10} + + released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, released) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs[stale.BatchID].Status) + require.Equal(t, "SUBMIT_STALE_BEFORE_PROVIDER", batchImageDerefString(repo.jobs[stale.BatchID].LastErrorCode)) + require.Len(t, billing.releases, 1) + require.Equal(t, BatchImageReleaseRequestID(stale.BatchID), billing.releases[0].RequestID) + require.Equal(t, BatchImageJobStatusSubmitted, repo.jobs[active.BatchID].Status) +} + +func TestBatchImageBillingRecoveryService_SkipsJobRefreshedByHeartbeat(t *testing.T) { + repo := newFakeBatchImageRepository() + apiKeyID := int64(22) + holdAmount := 0.5 + // updated_at 在 cutoff 之后(慢提交心跳持续续期):不得误杀退款。 + fresh := &BatchImageJob{ + BatchID: "imgbatch_fresh_uploading", + UserID: 11, + APIKeyID: &apiKeyID, + Status: BatchImageJobStatusUploading, + EstimatedCost: holdAmount, + HoldAmount: &holdAmount, + CreatedAt: time.Now().Add(-time.Hour), + UpdatedAt: time.Now(), + } + repo.jobs[fresh.BatchID] = fresh + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, StaleAfter: time.Minute, Limit: 10} + + released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background()) + require.NoError(t, err) + require.Equal(t, 0, released) + require.Equal(t, BatchImageJobStatusUploading, repo.jobs[fresh.BatchID].Status) + require.Empty(t, billing.releases) +} + +func TestBatchImageBillingRecoveryService_EnqueuesRetryWhenReleaseFails(t *testing.T) { + repo := newFakeBatchImageRepository() + apiKeyID := int64(22) + holdAmount := 0.5 + stale := &BatchImageJob{ + BatchID: "imgbatch_stale_release_fail", + UserID: 11, + APIKeyID: &apiKeyID, + Status: BatchImageJobStatusCreated, + EstimatedCost: holdAmount, + HoldAmount: &holdAmount, + CreatedAt: time.Now().Add(-time.Hour), + UpdatedAt: time.Now().Add(-time.Hour), + } + repo.jobs[stale.BatchID] = stale + billing := &fakeBatchImageBillingRepo{releaseErr: errors.New("billing db down")} + queue := &recordingBatchImageQueue{fakeBatchImageQueue: newFakeBatchImageQueue("")} + svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, Queue: queue, StaleAfter: time.Minute, Limit: 10} + + released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background()) + // job 已转 failed、不会再出现在 stale 列表:释放失败必须入队重试 + //(由 worker 的 releaseTerminalHold 兜底),否则冻结余额永久泄漏。 + require.Error(t, err) + require.Equal(t, 0, released) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs[stale.BatchID].Status) + require.Equal(t, []string{stale.BatchID}, queue.enqueued) +} diff --git a/backend/internal/service/batch_image_cleanup.go b/backend/internal/service/batch_image_cleanup.go new file mode 100644 index 0000000000..364a7e4795 --- /dev/null +++ b/backend/internal/service/batch_image_cleanup.go @@ -0,0 +1,312 @@ +package service + +import ( + "context" + "errors" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + defaultBatchImageInputRetentionAfterTerminal = 24 * time.Hour + defaultBatchImageOutputRetentionAfterTerminal = 72 * time.Hour + defaultBatchImageCleanupInterval = 30 * time.Minute + defaultBatchImageCleanupBatchSize = 100 +) + +type BatchImageCleanupService struct { + Repo BatchImageRepository + ProviderRegistry *BatchImageProviderRegistry + AccountResolver BatchImageAccountResolver + Config *config.Config + + cancel context.CancelFunc + done chan struct{} + mu sync.Mutex +} + +func NewBatchImageCleanupService(repo BatchImageRepository, accountRepo AccountRepository, cfg *config.Config) *BatchImageCleanupService { + return &BatchImageCleanupService{ + Repo: repo, + ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg), + AccountResolver: &BatchImageAccountRepositoryResolver{Repo: accountRepo}, + Config: cfg, + } +} + +// appendCleanupEvent 追加清理审计事件;事件写入失败不阻断清理流程,但必须留痕。 +func (s *BatchImageCleanupService) appendCleanupEvent(ctx context.Context, batchID, eventType string, payload any) { + if err := s.Repo.AppendBatchImageEvent(ctx, batchID, eventType, payload); err != nil { + logger.L().Warn("batch_image.cleanup_event_failed", + zap.String("batch_id", batchID), + zap.String("event_type", eventType), + zap.Error(err), + ) + } +} + +func (s *BatchImageCleanupService) DeleteOutputsForOwner(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) { + job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + if job.Status == BatchImageJobStatusOutputDeleted || job.OutputDeletedAt != nil { + return BatchImageJobToPublic(job), nil + } + if job.Status != BatchImageJobStatusCompleted { + return nil, ErrBatchImageOutputDeleteNotReady + } + s.appendCleanupEvent(ctx, job.BatchID, "manual_output_delete_requested", map[string]any{ + "batch_id": job.BatchID, + "cleanup_target": "output", + "reason": "manual", + }) + if err := s.cleanupJob(ctx, job, CleanupTargetOutput, "manual"); err != nil { + return nil, err + } + updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + return BatchImageJobToPublic(updated), nil +} + +func (s *BatchImageCleanupService) CleanupInput(ctx context.Context, batchID string) error { + job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID) + if err != nil { + return err + } + return s.cleanupJob(ctx, job, CleanupTargetInput, "ttl") +} + +func (s *BatchImageCleanupService) CleanupOutput(ctx context.Context, batchID string, reason string) error { + job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID) + if err != nil { + return err + } + return s.cleanupJob(ctx, job, CleanupTargetOutput, reason) +} + +func (s *BatchImageCleanupService) RunOnce(ctx context.Context, now time.Time) (BatchImageCleanupRunResult, error) { + if s == nil || s.Repo == nil { + return BatchImageCleanupRunResult{}, ErrBatchImageCleanupFailed + } + if now.IsZero() { + now = time.Now() + } + limit := s.cleanupBatchSize() + result := BatchImageCleanupRunResult{} + inputCutoff := now.Add(-s.inputRetentionAfterTerminal()) + inputJobs, err := s.Repo.ListBatchImageJobsDueForInputCleanup(ctx, inputCutoff, limit) + if err != nil { + return result, err + } + for _, job := range inputJobs { + if job == nil { + continue + } + if err := s.cleanupJob(ctx, job, CleanupTargetInput, "ttl"); err != nil { + result.Failures++ + continue + } + result.InputCleaned++ + } + outputJobs, err := s.Repo.ListBatchImageJobsDueForOutputCleanup(ctx, now, limit) + if err != nil { + return result, err + } + for _, job := range outputJobs { + if job == nil { + continue + } + if err := s.cleanupJob(ctx, job, CleanupTargetOutput, "expired"); err != nil { + result.Failures++ + continue + } + result.OutputCleaned++ + } + return result, nil +} + +func (s *BatchImageCleanupService) Start() { + if s == nil || s.Repo == nil || s.Config == nil || !s.Config.BatchImage.Enabled || s.cleanupInterval() <= 0 { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.cancel != nil { + return + } + ctx, cancel := context.WithCancel(context.Background()) + s.cancel = cancel + s.done = make(chan struct{}) + go func() { + defer close(s.done) + ticker := time.NewTicker(s.cleanupInterval()) + defer ticker.Stop() + for { + _, _ = s.RunOnce(ctx, time.Now()) + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } + }() +} + +func (s *BatchImageCleanupService) Stop() { + if s == nil { + return + } + s.mu.Lock() + cancel := s.cancel + done := s.done + s.cancel = nil + s.done = nil + s.mu.Unlock() + if cancel != nil { + cancel() + } + if done != nil { + <-done + } +} + +func (s *BatchImageCleanupService) cleanupJob(ctx context.Context, job *BatchImageJob, target CleanupTarget, reason string) error { + if job == nil { + return ErrBatchImageJobNotFound + } + switch target { + case CleanupTargetInput: + if job.InputDeletedAt != nil { + return nil + } + if !IsTerminalBatchImageJobStatus(job.Status) { + return ErrBatchImageCleanupFailed + } + s.appendCleanupEvent(ctx, job.BatchID, "input_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil)) + case CleanupTargetOutput: + if job.OutputDeletedAt != nil || job.Status == BatchImageJobStatusOutputDeleted { + return nil + } + if job.Status != BatchImageJobStatusCompleted && job.Status != BatchImageJobStatusFailed && job.Status != BatchImageJobStatusCancelled { + return ErrBatchImageOutputDeleteNotReady + } + s.appendCleanupEvent(ctx, job.BatchID, "output_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil)) + default: + return ErrUnsupportedCleanupTarget + } + + if err := s.callProviderCleanup(ctx, job, target); err != nil { + code := cleanupFailureCode(err) + msg := sanitizeBatchImagePublicMessage(err.Error()) + if recordErr := s.Repo.RecordBatchImageCleanupFailure(ctx, job.BatchID, code, msg); recordErr != nil { + logger.L().Warn("batch_image.cleanup_failure_record_failed", + zap.String("batch_id", job.BatchID), + zap.Error(recordErr), + ) + } + event := string(target) + "_cleanup_failed" + s.appendCleanupEvent(ctx, job.BatchID, event, map[string]any{"batch_id": job.BatchID, "cleanup_target": string(target), "reason": reason, "error_code": code}) + if errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) { + return ErrBatchImageCleanupUnsafePath + } + return ErrBatchImageProviderCleanupFailed + } + + deletedAt := time.Now() + if target == CleanupTargetInput { + return s.Repo.MarkBatchImageInputDeleted(ctx, job.BatchID, deletedAt) + } + return s.Repo.MarkBatchImageOutputDeleted(ctx, job.BatchID, deletedAt) +} + +func (s *BatchImageCleanupService) callProviderCleanup(ctx context.Context, job *BatchImageJob, target CleanupTarget) error { + if s == nil || s.ProviderRegistry == nil || s.AccountResolver == nil { + return ErrBatchImageCleanupFailed + } + provider, ok := s.ProviderRegistry.Get(job.Provider) + if !ok || provider == nil { + return ErrBatchImageUnsupportedProvider + } + if job.AccountID == nil || *job.AccountID <= 0 { + return ErrBatchImageMissingAccountID + } + account, err := s.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID) + if err != nil { + return err + } + if err := provider.Cleanup(ctx, job, account, target); err != nil { + if cleanupErrorIsNotFound(err) { + return nil + } + return err + } + return nil +} + +func (s *BatchImageCleanupService) inputRetentionAfterTerminal() time.Duration { + if s != nil && s.Config != nil && s.Config.BatchImage.InputRetentionAfterTerminalHours > 0 { + return time.Duration(s.Config.BatchImage.InputRetentionAfterTerminalHours) * time.Hour + } + return defaultBatchImageInputRetentionAfterTerminal +} + +func (s *BatchImageCleanupService) cleanupInterval() time.Duration { + if s != nil && s.Config != nil && s.Config.BatchImage.CleanupIntervalMinutes > 0 { + return time.Duration(s.Config.BatchImage.CleanupIntervalMinutes) * time.Minute + } + return defaultBatchImageCleanupInterval +} + +func (s *BatchImageCleanupService) cleanupBatchSize() int { + if s != nil && s.Config != nil && s.Config.BatchImage.CleanupBatchSize > 0 { + return s.Config.BatchImage.CleanupBatchSize + } + return defaultBatchImageCleanupBatchSize +} + +type BatchImageCleanupRunResult struct { + InputCleaned int + OutputCleaned int + Failures int +} + +func cleanupEventPayload(batchID string, target CleanupTarget, reason string, deletedAt *time.Time) map[string]any { + payload := map[string]any{ + "batch_id": batchID, + "cleanup_target": string(target), + "reason": reason, + } + if deletedAt != nil { + payload["deleted_at"] = deletedAt.UTC().Format(time.RFC3339) + } + return payload +} + +func cleanupErrorIsNotFound(err error) bool { + if err == nil { + return false + } + reason := strings.ToUpper(infraerrors.Reason(err)) + msg := strings.ToUpper(err.Error()) + return strings.Contains(reason, "NOT_FOUND") || strings.Contains(msg, "NOT FOUND") || strings.Contains(msg, "404") +} + +func cleanupFailureCode(err error) string { + if errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) { + return "BATCH_IMAGE_CLEANUP_UNSAFE_PATH" + } + reason := strings.TrimSpace(infraerrors.Reason(err)) + if reason != "" { + return reason + } + return "BATCH_IMAGE_PROVIDER_CLEANUP_FAILED" +} diff --git a/backend/internal/service/batch_image_cleanup_test.go b/backend/internal/service/batch_image_cleanup_test.go new file mode 100644 index 0000000000..30d71a9504 --- /dev/null +++ b/backend/internal/service/batch_image_cleanup_test.go @@ -0,0 +1,231 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/stretchr/testify/require" +) + +func TestBatchImageCleanupService_DeleteOutputsForOwner(t *testing.T) { + ctx := context.Background() + + t.Run("deletes completed output and returns public dto", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + + got, err := svc.DeleteOutputsForOwner(ctx, testBatchImageOwner(), "imgbatch_cleanup") + require.NoError(t, err) + require.Equal(t, "output_deleted", got.Status) + require.NotNil(t, got.OutputDeletedAt) + require.Equal(t, []CleanupTarget{CleanupTargetOutput}, provider.cleanupTargets) + require.NotNil(t, repo.jobs["imgbatch_cleanup"].OutputDeletedAt) + require.Equal(t, BatchImageJobStatusOutputDeleted, repo.jobs["imgbatch_cleanup"].Status) + body := mustJSON(t, got) + requireBatchImagePublicJSONHasNoInternals(t, body) + }) + + t.Run("repeated delete is idempotent", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + deletedAt := time.Now() + repo.jobs["imgbatch_cleanup"].Status = BatchImageJobStatusOutputDeleted + repo.jobs["imgbatch_cleanup"].OutputDeletedAt = &deletedAt + + got, err := svc.DeleteOutputsForOwner(ctx, testBatchImageOwner(), "imgbatch_cleanup") + require.NoError(t, err) + require.Equal(t, "output_deleted", got.Status) + require.Empty(t, provider.cleanupTargets) + }) + + t.Run("not completed returns not ready", func(t *testing.T) { + svc, repo, _ := newTestBatchImageCleanupService() + repo.jobs["imgbatch_cleanup"].Status = BatchImageJobStatusRunning + + _, err := svc.DeleteOutputsForOwner(ctx, testBatchImageOwner(), "imgbatch_cleanup") + require.ErrorIs(t, err, ErrBatchImageOutputDeleteNotReady) + }) + + t.Run("non owner returns not found", func(t *testing.T) { + svc, _, _ := newTestBatchImageCleanupService() + _, err := svc.DeleteOutputsForOwner(ctx, BatchImageOwner{UserID: 11, APIKeyID: 999}, "imgbatch_cleanup") + require.ErrorIs(t, err, ErrBatchImageJobNotFound) + }) + + t.Run("provider not found is success", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + provider.cleanupErr = infraerrors.New(404, "PROVIDER_NOT_FOUND", "provider file not found: gs://hidden") + + got, err := svc.DeleteOutputsForOwner(ctx, testBatchImageOwner(), "imgbatch_cleanup") + require.NoError(t, err) + require.Equal(t, "output_deleted", got.Status) + require.NotNil(t, repo.jobs["imgbatch_cleanup"].OutputDeletedAt) + }) + + t.Run("provider transient error is sanitized and records failure", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + provider.cleanupErr = errors.New("temporary cleanup failed for gs://secret-output") + + _, err := svc.DeleteOutputsForOwner(ctx, testBatchImageOwner(), "imgbatch_cleanup") + require.ErrorIs(t, err, ErrBatchImageProviderCleanupFailed) + require.Equal(t, "BATCH_IMAGE_PROVIDER_CLEANUP_FAILED", infraerrors.Reason(err)) + require.NotContains(t, infraerrors.Message(err), "gs://") + require.Equal(t, "BATCH_IMAGE_PROVIDER_CLEANUP_FAILED", batchImageDerefString(repo.jobs["imgbatch_cleanup"].LastErrorCode)) + require.Equal(t, "upstream provider operation failed", batchImageDerefString(repo.jobs["imgbatch_cleanup"].LastErrorMessage)) + }) + + t.Run("unsafe cleanup path is not swallowed", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + provider.cleanupErr = ErrBatchImageProviderUnsafeCleanupPath + + _, err := svc.DeleteOutputsForOwner(ctx, testBatchImageOwner(), "imgbatch_cleanup") + require.ErrorIs(t, err, ErrBatchImageCleanupUnsafePath) + require.Equal(t, "BATCH_IMAGE_CLEANUP_UNSAFE_PATH", batchImageDerefString(repo.jobs["imgbatch_cleanup"].LastErrorCode)) + }) +} + +func TestBatchImageCleanupService_InputOutputAndWorker(t *testing.T) { + ctx := context.Background() + now := time.Now() + + t.Run("input cleanup marks input only", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + + err := svc.CleanupInput(ctx, "imgbatch_cleanup") + require.NoError(t, err) + require.Equal(t, []CleanupTarget{CleanupTargetInput}, provider.cleanupTargets) + require.NotNil(t, repo.jobs["imgbatch_cleanup"].InputDeletedAt) + require.Equal(t, BatchImageJobStatusCompleted, repo.jobs["imgbatch_cleanup"].Status) + + err = svc.CleanupInput(ctx, "imgbatch_cleanup") + require.NoError(t, err) + require.Len(t, provider.cleanupTargets, 1) + }) + + t.Run("output cleanup for failed job keeps status", func(t *testing.T) { + svc, repo, _ := newTestBatchImageCleanupService() + repo.jobs["imgbatch_cleanup"].Status = BatchImageJobStatusFailed + + err := svc.CleanupOutput(ctx, "imgbatch_cleanup", "ttl") + require.NoError(t, err) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs["imgbatch_cleanup"].Status) + require.NotNil(t, repo.jobs["imgbatch_cleanup"].OutputDeletedAt) + }) + + t.Run("worker processes due jobs and continues after failure", func(t *testing.T) { + svc, repo, provider := newTestBatchImageCleanupService() + provider.cleanupErr = nil + old := now.Add(-48 * time.Hour) + expired := now.Add(-time.Minute) + future := now.Add(time.Hour) + repo.jobs["imgbatch_cleanup"].FinishedAt = &old + repo.jobs["imgbatch_cleanup"].OutputExpiresAt = &expired + repo.jobs["imgbatch_running"] = cleanupTestJob("imgbatch_running", BatchImageJobStatusRunning) + repo.jobs["imgbatch_running"].FinishedAt = &old + repo.jobs["imgbatch_running"].OutputExpiresAt = &expired + repo.jobs["imgbatch_future"] = cleanupTestJob("imgbatch_future", BatchImageJobStatusCompleted) + repo.jobs["imgbatch_future"].FinishedAt = &old + repo.jobs["imgbatch_future"].OutputExpiresAt = &future + + result, err := svc.RunOnce(ctx, now) + require.NoError(t, err) + require.Equal(t, 2, result.InputCleaned) + require.Equal(t, 1, result.OutputCleaned) + require.Equal(t, BatchImageJobStatusRunning, repo.jobs["imgbatch_running"].Status) + require.Nil(t, repo.jobs["imgbatch_future"].OutputDeletedAt) + require.NotContains(t, strings.Join(repo.events["imgbatch_running"], ","), "cleanup") + }) +} + +func TestBatchImageSettlementOutputExpiration(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_expire") + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{ + Repo: repo, + BillingRepo: billing, + Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}, + Config: &config.Config{BatchImage: config.BatchImageConfig{OutputRetentionAfterTerminalHours: 5}}, + } + + _, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.NotNil(t, repo.jobs[job.BatchID].OutputExpiresAt) + require.WithinDuration(t, time.Now().Add(5*time.Hour), *repo.jobs[job.BatchID].OutputExpiresAt, time.Minute) + + existing := time.Now().Add(time.Hour) + second := testSettlingBatchImageJob("imgbatch_keep_expire") + second.OutputExpiresAt = &existing + repo.jobs[second.BatchID] = second + _, err = svc.Settle(context.Background(), second.BatchID) + require.NoError(t, err) + require.Equal(t, existing, *repo.jobs[second.BatchID].OutputExpiresAt) +} + +func TestBatchImageDownloadAfterOutputDeletedReturnsGone(t *testing.T) { + svc, repo, _ := newTestBatchImageDownloadService() + now := time.Now() + repo.jobs["imgbatch_download"].Status = BatchImageJobStatusOutputDeleted + repo.jobs["imgbatch_download"].OutputDeletedAt = &now + + stream, err := svc.OpenItemContent(context.Background(), testBatchImageOwner(), "imgbatch_download", "cover/../001", 0) + require.Nil(t, stream) + require.ErrorIs(t, err, ErrBatchImageOutputDeleted) + + var out strings.Builder + result, err := svc.StreamZip(context.Background(), testBatchImageOwner(), "imgbatch_download", BatchImageZipOptions{}, &out) + require.Nil(t, result) + require.ErrorIs(t, err, ErrBatchImageOutputDeleted) +} + +func newTestBatchImageCleanupService() (*BatchImageCleanupService, *fakeBatchImageRepository, *publicBatchImageProvider) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_cleanup"] = cleanupTestJob("imgbatch_cleanup", BatchImageJobStatusCompleted) + provider := &publicBatchImageProvider{name: BatchImageProviderGeminiAPI} + accountID := int64(101) + svc := &BatchImageCleanupService{ + Repo: repo, + ProviderRegistry: NewBatchImageProviderRegistry(provider), + AccountResolver: &fakeBatchImageAccountResolver{account: &Account{ID: accountID, Platform: PlatformGemini, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true}}, + Config: &config.Config{BatchImage: config.BatchImageConfig{CleanupBatchSize: 10, InputRetentionAfterTerminalHours: 24}}, + } + return svc, repo, provider +} + +func cleanupTestJob(batchID, status string) *BatchImageJob { + apiKeyID := int64(22) + accountID := int64(101) + now := time.Now().Add(-48 * time.Hour) + return &BatchImageJob{ + BatchID: batchID, + UserID: 11, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: status, + ProviderJobName: batchImageStringPtr("providers/internal/job"), + ProviderInputRef: batchImageStringPtr("files/internal/input"), + ProviderOutputRef: batchImageStringPtr("files/internal/output"), + ItemCount: 1, + SuccessCount: 1, + CreatedAt: now, + UpdatedAt: now, + FinishedAt: &now, + SettledAt: &now, + } +} + +func mustJSON(t *testing.T, v any) string { + t.Helper() + b, err := json.Marshal(v) + require.NoError(t, err) + return string(b) +} diff --git a/backend/internal/service/batch_image_download.go b/backend/internal/service/batch_image_download.go new file mode 100644 index 0000000000..7f303d22c5 --- /dev/null +++ b/backend/internal/service/batch_image_download.go @@ -0,0 +1,660 @@ +package service + +import ( + "archive/zip" + "bufio" + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "path/filepath" + "sort" + "strings" + "sync" + "time" + "unicode" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +const ( + defaultBatchImageZipMaxItems = 200 + defaultBatchImageZipMaxBytes = 512 * 1024 * 1024 + defaultBatchImageDownloadDuration = 10 * time.Minute + defaultBatchImageDownloadConcurrency = 1 + batchImageDownloadScannerMaxLineBytes = 16 * 1024 * 1024 +) + +var errBatchImageDownloadSizeExceeded = errors.New("batch image download size limit exceeded") + +type BatchImageDownloadLimiter interface { + Acquire(ctx context.Context, userID string, kind string) (BatchImageDownloadPermit, error) +} + +type BatchImageDownloadPermit interface { + Release(ctx context.Context) error +} + +type BatchImageContentStream struct { + Reader io.ReadCloser + ContentType string + Filename string + ContentLength *int64 +} + +type BatchImageZipOptions struct { + Status string + MaxItems int + IncludeManifest bool +} + +type BatchImageZipResult struct { + FileCount int + ErrorCount int +} + +type BatchImageLineImages struct { + CustomID string + Images []BatchImageInlineImage + ErrorCode string + ErrorMessage string +} + +type BatchImageInlineImage struct { + MimeType string + Extension string + Base64Data string +} + +type BatchImageDownloadService struct { + Repo BatchImageRepository + ProviderRegistry *BatchImageProviderRegistry + AccountResolver BatchImageAccountResolver + Limiter BatchImageDownloadLimiter + Config *config.Config +} + +type batchImageDownloadLimitWriter struct { + w io.Writer + limit int64 + written int64 +} + +func (w *batchImageDownloadLimitWriter) Write(p []byte) (int, error) { + if w == nil || w.w == nil { + return 0, io.ErrClosedPipe + } + if w.limit > 0 && w.written+int64(len(p)) > w.limit { + return 0, errBatchImageDownloadSizeExceeded + } + n, err := w.w.Write(p) + w.written += int64(n) + return n, err +} + +func NewBatchImageDownloadService(repo BatchImageRepository, accountRepo AccountRepository, limiter BatchImageDownloadLimiter, cfg *config.Config) *BatchImageDownloadService { + return &BatchImageDownloadService{ + Repo: repo, + ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg), + AccountResolver: &BatchImageAccountRepositoryResolver{Repo: accountRepo}, + Limiter: limiter, + Config: cfg, + } +} + +func (s *BatchImageDownloadService) OpenItemContent(ctx context.Context, owner BatchImageOwner, batchID string, customID string, imageIndex int) (*BatchImageContentStream, error) { + if imageIndex < 0 { + return nil, ErrBatchImageItemImageIndexOutOfRange + } + job, err := s.getCompletedJob(ctx, owner, batchID) + if err != nil { + return nil, err + } + item, err := s.Repo.GetBatchImageItemForDownload(ctx, job.BatchID, customID) + if err != nil { + return nil, err + } + if item.Status != BatchImageItemStatusSuccess { + return nil, ErrBatchImageItemFailed + } + if imageIndex >= item.ImageCount { + return nil, ErrBatchImageItemImageIndexOutOfRange + } + + permit, err := s.acquirePermit(ctx, owner.UserID, "item") + if err != nil { + return nil, err + } + releasePermit := true + defer func() { + if releasePermit && permit != nil { + _ = permit.Release(ctx) + } + }() + + provider, account, err := s.providerAndAccount(ctx, job) + if err != nil { + return nil, err + } + r, _, err := provider.OpenResult(ctx, job, account) + if err != nil { + return nil, ErrBatchImageResultMissing.WithCause(err) + } + defer func() { _ = r.Close() }() + + line, err := findBatchImageLineImages(r, item.CustomID) + if err != nil { + return nil, err + } + if imageIndex >= len(line.Images) { + return nil, ErrBatchImageItemImageIndexOutOfRange + } + image := line.Images[imageIndex] + if strings.TrimSpace(image.Base64Data) == "" { + return nil, ErrBatchImageResultMissing + } + contentType := strings.TrimSpace(image.MimeType) + if contentType == "" { + contentType = "application/octet-stream" + } + extension := strings.TrimSpace(image.Extension) + if extension == "" { + extension = batchImageFileExtension(contentType) + } + if extension == "" { + extension = "bin" + } + + reader := base64.NewDecoder(base64.StdEncoding, strings.NewReader(image.Base64Data)) + releasePermit = false + return &BatchImageContentStream{ + Reader: &batchImagePermitReadCloser{Reader: reader, permit: permit}, + ContentType: contentType, + Filename: BatchImageSafeDownloadFilename(item.CustomID, extension), + }, nil +} + +func (s *BatchImageDownloadService) StreamZip(ctx context.Context, owner BatchImageOwner, batchID string, opts BatchImageZipOptions, w io.Writer) (*BatchImageZipResult, error) { + job, err := s.getCompletedJob(ctx, owner, batchID) + if err != nil { + return nil, err + } + maxItems := opts.MaxItems + if cap := s.maxZipItems(); maxItems <= 0 || maxItems > cap { + // 客户端传入的 max_items 不得放大管理员配置的 ZIP 上限。 + maxItems = cap + } + if job.SuccessCount > maxItems { + return nil, ErrBatchImageZipTooManyItems + } + successItems, err := s.Repo.ListBatchImageItemsForDownload(ctx, job.BatchID, BatchImageItemStatusSuccess, maxItems+1) + if err != nil { + return nil, err + } + if len(successItems) > maxItems { + return nil, ErrBatchImageZipTooManyItems + } + failedItems, err := s.Repo.ListBatchImageItemsForDownload(ctx, job.BatchID, BatchImageItemStatusFailed, maxItems) + if err != nil { + return nil, err + } + + permit, err := s.acquirePermit(ctx, owner.UserID, "zip") + if err != nil { + return nil, err + } + if permit != nil { + defer func() { _ = permit.Release(ctx) }() + } + + provider, account, err := s.providerAndAccount(ctx, job) + if err != nil { + return nil, err + } + r, _, err := provider.OpenResult(ctx, job, account) + if err != nil { + return nil, ErrBatchImageResultMissing.WithCause(err) + } + defer func() { _ = r.Close() }() + + streamCtx := ctx + cancel := func() {} + if d := s.maxDownloadDuration(); d > 0 { + streamCtx, cancel = context.WithTimeout(ctx, d) + } + defer cancel() + + limitedWriter := &batchImageDownloadLimitWriter{w: w, limit: s.maxDownloadBytes()} + zipWriter := zip.NewWriter(limitedWriter) + result, manifestFiles, zipErrors, err := s.writeZipImages(streamCtx, zipWriter, r, successItems) + if err != nil { + _ = zipWriter.Close() + if errors.Is(err, errBatchImageDownloadSizeExceeded) { + return result, ErrBatchImageDownloadTooLarge.WithCause(err) + } + return result, ErrBatchImageDownloadFailed.WithCause(err) + } + zipErrors = append(zipErrors, batchImageZipErrorsFromItems(failedItems)...) + if err := writeBatchImageZipJSON(zipWriter, "manifest.json", batchImageZipManifest{ + BatchID: job.BatchID, + Model: job.Model, + ItemCount: job.ItemCount, + SuccessCount: job.SuccessCount, + FailCount: job.FailCount, + Files: manifestFiles, + }); err != nil { + _ = zipWriter.Close() + if errors.Is(err, errBatchImageDownloadSizeExceeded) { + return result, ErrBatchImageDownloadTooLarge.WithCause(err) + } + return result, ErrBatchImageDownloadFailed.WithCause(err) + } + if err := writeBatchImageZipJSON(zipWriter, "errors.json", zipErrors); err != nil { + _ = zipWriter.Close() + if errors.Is(err, errBatchImageDownloadSizeExceeded) { + return result, ErrBatchImageDownloadTooLarge.WithCause(err) + } + return result, ErrBatchImageDownloadFailed.WithCause(err) + } + result.ErrorCount = len(zipErrors) + if err := zipWriter.Close(); err != nil { + if errors.Is(err, errBatchImageDownloadSizeExceeded) { + return result, ErrBatchImageDownloadTooLarge.WithCause(err) + } + return result, ErrBatchImageDownloadFailed.WithCause(err) + } + return result, nil +} + +func (s *BatchImageDownloadService) writeZipImages(ctx context.Context, zipWriter *zip.Writer, resultReader io.Reader, successItems []*BatchImageItem) (*BatchImageZipResult, []batchImageZipManifestFile, []batchImageZipError, error) { + successByID := make(map[string]*BatchImageItem, len(successItems)) + missing := make(map[string]struct{}, len(successItems)) + for _, item := range successItems { + if item == nil { + continue + } + successByID[item.CustomID] = item + missing[item.CustomID] = struct{}{} + } + scanner := bufio.NewScanner(resultReader) + scanner.Buffer(make([]byte, 0, 64*1024), batchImageDownloadScannerMaxLineBytes) + + result := &BatchImageZipResult{} + var manifestFiles []batchImageZipManifestFile + var zipErrors []batchImageZipError + for scanner.Scan() { + if err := ctx.Err(); err != nil { + return result, manifestFiles, zipErrors, err + } + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + images, err := ExtractBatchImagePartsFromResultLine([]byte(line)) + if err != nil { + return result, manifestFiles, zipErrors, err + } + item := successByID[images.CustomID] + if item == nil { + continue + } + delete(missing, images.CustomID) + if len(images.Images) == 0 { + zipErrors = append(zipErrors, batchImageZipError{CustomID: images.CustomID, Code: "EMPTY_IMAGE_OUTPUT", Message: "provider response contained no image output"}) + continue + } + for idx, image := range images.Images { + extension := image.Extension + if extension == "" { + extension = "bin" + } + filename := batchImageZipImageFilename(item.CustomID, idx, extension) + entry, err := zipWriter.CreateHeader(&zip.FileHeader{Name: filename, Method: zip.Deflate}) + if err != nil { + return result, manifestFiles, zipErrors, err + } + decoder := base64.NewDecoder(base64.StdEncoding, strings.NewReader(image.Base64Data)) + if _, err := io.Copy(entry, decoder); err != nil { + zipErrors = append(zipErrors, batchImageZipError{CustomID: item.CustomID, Code: "IMAGE_DECODE_FAILED", Message: "image data could not be decoded"}) + continue + } + result.FileCount++ + manifestFiles = append(manifestFiles, batchImageZipManifestFile{ + CustomID: item.CustomID, + Filename: filename, + MimeType: image.MimeType, + ImageIndex: idx, + }) + } + } + if err := scanner.Err(); err != nil { + return result, manifestFiles, zipErrors, err + } + missingIDs := make([]string, 0, len(missing)) + for customID := range missing { + missingIDs = append(missingIDs, customID) + } + sort.Strings(missingIDs) + for _, customID := range missingIDs { + zipErrors = append(zipErrors, batchImageZipError{CustomID: customID, Code: "RESULT_MISSING", Message: "provider result was not found for item"}) + } + return result, manifestFiles, zipErrors, nil +} + +func (s *BatchImageDownloadService) getCompletedJob(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImageJob, error) { + if s == nil || s.Repo == nil { + return nil, ErrBatchImageDownloadFailed + } + job, err := s.Repo.GetBatchImageJobForDownload(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + switch job.Status { + case BatchImageJobStatusCompleted: + return job, nil + case BatchImageJobStatusOutputDeleted: + return nil, ErrBatchImageOutputDeleted + default: + return nil, ErrBatchImageNotReady + } +} + +func (s *BatchImageDownloadService) providerAndAccount(ctx context.Context, job *BatchImageJob) (BatchImageProvider, *Account, error) { + if s == nil || s.ProviderRegistry == nil || s.AccountResolver == nil || job == nil { + return nil, nil, ErrBatchImageDownloadFailed + } + provider, ok := s.ProviderRegistry.Get(job.Provider) + if !ok || provider == nil { + return nil, nil, ErrBatchImageUnsupportedProvider + } + if job.AccountID == nil || *job.AccountID <= 0 { + return nil, nil, ErrBatchImageMissingAccountID + } + account, err := s.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID) + if err != nil { + return nil, nil, ErrBatchImageDownloadFailed + } + if !provider.SupportsAccount(account) { + return nil, nil, ErrBatchImageProviderUnsupportedAccount + } + return provider, account, nil +} + +func (s *BatchImageDownloadService) acquirePermit(ctx context.Context, userID int64, kind string) (BatchImageDownloadPermit, error) { + if s == nil || s.Limiter == nil { + return nil, nil + } + permit, err := s.Limiter.Acquire(ctx, fmt.Sprintf("%d", userID), kind) + if err != nil { + if infraerrors.Code(err) == http.StatusTooManyRequests { + return nil, ErrBatchImageDownloadLimited + } + return nil, ErrBatchImageDownloadLimited.WithCause(err) + } + return permit, nil +} + +func (s *BatchImageDownloadService) maxZipItems() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxDownloadItemsZip > 0 { + return s.Config.BatchImage.MaxDownloadItemsZip + } + return defaultBatchImageZipMaxItems +} + +func (s *BatchImageDownloadService) maxDownloadBytes() int64 { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxDownloadBytesPerRequest > 0 { + return s.Config.BatchImage.MaxDownloadBytesPerRequest + } + return defaultBatchImageZipMaxBytes +} + +func (s *BatchImageDownloadService) maxDownloadDuration() time.Duration { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxDownloadDurationSeconds > 0 { + return time.Duration(s.Config.BatchImage.MaxDownloadDurationSeconds) * time.Second + } + return defaultBatchImageDownloadDuration +} + +func ExtractBatchImagePartsFromResultLine(line []byte) (*BatchImageLineImages, error) { + var obj map[string]any + if err := json.Unmarshal(line, &obj); err != nil { + return nil, ErrBatchImageIndexParseFailed.WithCause(err) + } + customID := batchImageFirstNonEmptyString( + batchImageMapString(obj, "key"), + batchImageMapString(obj, "custom_id"), + batchImageMapString(obj, "customId"), + batchImageNestedString(obj, "request", "key"), + ) + if customID == "" { + return nil, ErrBatchImageIndexParseFailed.WithCause(fmt.Errorf("missing custom id")) + } + out := &BatchImageLineImages{CustomID: customID} + out.Images = append(out.Images, extractBatchImageInlineImages(batchImageNestedAny(obj, "response", "candidates"))...) + out.Images = append(out.Images, extractBatchImageInlineImages(obj["candidates"])...) + if len(out.Images) > 0 { + return out, nil + } + if code, message, ok := batchImageFailureFromProviderFields(obj); ok { + out.ErrorCode = code + out.ErrorMessage = truncateBatchImageMessage(message, batchImageMaxErrorMessageLength) + return out, nil + } + if _, hasResponse := obj["response"]; hasResponse || batchImageHasCandidates(obj) { + out.ErrorCode = "EMPTY_IMAGE_OUTPUT" + out.ErrorMessage = "provider response contained no image output" + return out, nil + } + out.ErrorCode = "PROVIDER_ITEM_FAILED" + out.ErrorMessage = "provider result line contained no image output" + return out, nil +} + +func extractBatchImageInlineImages(raw any) []BatchImageInlineImage { + candidates, ok := raw.([]any) + if !ok { + return nil + } + var images []BatchImageInlineImage + for _, candidateRaw := range candidates { + candidate, ok := candidateRaw.(map[string]any) + if !ok { + continue + } + parts, ok := batchImageNestedAny(candidate, "content", "parts").([]any) + if !ok { + continue + } + for _, partRaw := range parts { + part, ok := partRaw.(map[string]any) + if !ok { + continue + } + inline, ok := firstMap(part["inlineData"], part["inline_data"]) + if !ok { + continue + } + data := strings.TrimSpace(batchImageMapString(inline, "data")) + mime := strings.TrimSpace(batchImageFirstNonEmptyString(batchImageMapString(inline, "mimeType"), batchImageMapString(inline, "mime_type"))) + if data == "" || !strings.HasPrefix(strings.ToLower(mime), "image/") { + continue + } + images = append(images, BatchImageInlineImage{ + MimeType: mime, + Extension: batchImageFileExtension(mime), + Base64Data: data, + }) + } + } + return images +} + +func findBatchImageLineImages(r io.Reader, customID string) (*BatchImageLineImages, error) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), batchImageDownloadScannerMaxLineBytes) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + parsed, err := ExtractBatchImagePartsFromResultLine([]byte(line)) + if err != nil { + return nil, err + } + if parsed.CustomID != customID { + continue + } + if len(parsed.Images) == 0 { + if parsed.ErrorCode != "" { + return nil, ErrBatchImageItemFailed + } + return nil, ErrBatchImageResultMissing + } + return parsed, nil + } + if err := scanner.Err(); err != nil { + return nil, ErrBatchImageDownloadFailed.WithCause(err) + } + return nil, ErrBatchImageResultMissing +} + +func BatchImageSafeDownloadFilename(customID, extension string) string { + base := sanitizeBatchImageFilenameBase(customID) + extension = sanitizeBatchImageFilenameExtension(extension) + if extension == "" { + extension = "bin" + } + return base + "." + extension +} + +func BatchImageContentDispositionAttachment(filename string) string { + filename = strings.ReplaceAll(filename, "\\", "_") + filename = strings.ReplaceAll(filename, `"`, "_") + filename = sanitizeBatchImageFilenameBase(strings.TrimSuffix(filename, filepath.Ext(filename))) + filepath.Ext(filename) + return `attachment; filename="` + filename + `"` +} + +func sanitizeBatchImageFilenameBase(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "image" + } + var b strings.Builder + for _, r := range value { + switch { + case r == '/' || r == '\\' || r == ':' || r == 0: + _ = b.WriteByte('_') + case unicode.IsControl(r): + _ = b.WriteByte('_') + case unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '-' || r == '.': + _, _ = b.WriteRune(r) + default: + _ = b.WriteByte('_') + } + } + out := strings.Trim(b.String(), ". ") + for strings.Contains(out, "..") { + out = strings.ReplaceAll(out, "..", "_") + } + out = strings.Trim(out, ". ") + if out == "" { + out = "image" + } + if len(out) > 120 { + out = strings.TrimRight(out[:120], ". ") + } + if out == "" { + out = "image" + } + return out +} + +func sanitizeBatchImageFilenameExtension(extension string) string { + extension = strings.TrimPrefix(strings.TrimSpace(strings.ToLower(extension)), ".") + var b strings.Builder + for _, r := range extension { + if unicode.IsLetter(r) || unicode.IsDigit(r) { + _, _ = b.WriteRune(r) + } + } + out := b.String() + if len(out) > 12 { + out = out[:12] + } + return out +} + +func batchImageZipImageFilename(customID string, imageIndex int, extension string) string { + base := sanitizeBatchImageFilenameBase(customID) + if imageIndex > 0 { + base = fmt.Sprintf("%s_%d", base, imageIndex+1) + } + return "images/" + BatchImageSafeDownloadFilename(base, extension) +} + +func writeBatchImageZipJSON(zipWriter *zip.Writer, name string, value any) error { + entry, err := zipWriter.CreateHeader(&zip.FileHeader{Name: name, Method: zip.Deflate}) + if err != nil { + return err + } + encoder := json.NewEncoder(entry) + encoder.SetIndent("", " ") + return encoder.Encode(value) +} + +type batchImageZipManifest struct { + BatchID string `json:"batch_id"` + Model string `json:"model"` + ItemCount int `json:"item_count"` + SuccessCount int `json:"success_count"` + FailCount int `json:"fail_count"` + Files []batchImageZipManifestFile `json:"files"` +} + +type batchImageZipManifestFile struct { + CustomID string `json:"custom_id"` + Filename string `json:"filename"` + MimeType string `json:"mime_type"` + ImageIndex int `json:"image_index"` +} + +type batchImageZipError struct { + CustomID string `json:"custom_id"` + Code string `json:"code"` + Message string `json:"message"` +} + +func batchImageZipErrorsFromItems(items []*BatchImageItem) []batchImageZipError { + out := make([]batchImageZipError, 0, len(items)) + for _, item := range items { + if item == nil { + continue + } + out = append(out, batchImageZipError{ + CustomID: item.CustomID, + Code: batchImageDerefString(item.ErrorCode), + Message: sanitizeBatchImagePublicMessage(batchImageDerefString(item.ErrorMessage)), + }) + } + return out +} + +type batchImagePermitReadCloser struct { + io.Reader + permit BatchImageDownloadPermit + once sync.Once + err error +} + +func (r *batchImagePermitReadCloser) Close() error { + r.once.Do(func() { + if r.permit != nil { + r.err = r.permit.Release(context.Background()) + } + }) + return r.err +} diff --git a/backend/internal/service/batch_image_download_test.go b/backend/internal/service/batch_image_download_test.go new file mode 100644 index 0000000000..cc10dca699 --- /dev/null +++ b/backend/internal/service/batch_image_download_test.go @@ -0,0 +1,299 @@ +//go:build unit + +package service + +import ( + "archive/zip" + "bytes" + "context" + "encoding/json" + "io" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestBatchImageDownloadService_OpenItemContent(t *testing.T) { + ctx := context.Background() + + t.Run("streams image bytes with safe headers data", func(t *testing.T) { + svc, _, limiter := newTestBatchImageDownloadService() + + stream, err := svc.OpenItemContent(ctx, testBatchImageOwner(), "imgbatch_download", "cover/../001", 1) + require.NoError(t, err) + defer stream.Reader.Close() + + body, err := io.ReadAll(stream.Reader) + require.NoError(t, err) + require.Equal(t, []byte("second"), body) + require.Equal(t, "image/jpeg", stream.ContentType) + require.Equal(t, "cover___001.jpg", stream.Filename) + require.Equal(t, 1, limiter.acquireCount) + require.Zero(t, limiter.releaseCount) + require.NoError(t, stream.Reader.Close()) + require.Equal(t, 1, limiter.releaseCount) + }) + + tests := []struct { + name string + mutate func(*fakeBatchImageRepository) + id string + item string + index int + want error + }{ + {name: "non_owner", id: "imgbatch_download", item: "cover/../001", mutate: func(r *fakeBatchImageRepository) { + v := int64(999) + r.jobs["imgbatch_download"].APIKeyID = &v + }, want: ErrBatchImageJobNotFound}, + {name: "not_completed", id: "imgbatch_download", item: "cover/../001", mutate: func(r *fakeBatchImageRepository) { + r.jobs["imgbatch_download"].Status = BatchImageJobStatusRunning + }, want: ErrBatchImageNotReady}, + {name: "output_deleted", id: "imgbatch_download", item: "cover/../001", mutate: func(r *fakeBatchImageRepository) { + r.jobs["imgbatch_download"].Status = BatchImageJobStatusOutputDeleted + }, want: ErrBatchImageOutputDeleted}, + {name: "missing_item", id: "imgbatch_download", item: "missing", want: ErrBatchImageItemNotFound}, + {name: "failed_item", id: "imgbatch_download", item: "bad", want: ErrBatchImageItemFailed}, + {name: "out_of_range", id: "imgbatch_download", item: "cover/../001", index: 2, want: ErrBatchImageItemImageIndexOutOfRange}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + svc, repo, _ := newTestBatchImageDownloadService() + if tt.mutate != nil { + tt.mutate(repo) + } + + got, err := svc.OpenItemContent(ctx, testBatchImageOwner(), tt.id, tt.item, tt.index) + require.Nil(t, got) + require.ErrorIs(t, err, tt.want) + require.NotContains(t, err.Error(), batchImageDownloadTestBase64) + require.NotContains(t, err.Error(), "providers/") + require.NotContains(t, err.Error(), "gs://") + }) + } +} + +func TestBatchImageDownloadService_StreamZip(t *testing.T) { + ctx := context.Background() + + t.Run("streams zip with images manifest and errors", func(t *testing.T) { + svc, _, limiter := newTestBatchImageDownloadService() + var buf bytes.Buffer + + result, err := svc.StreamZip(ctx, testBatchImageOwner(), "imgbatch_download", BatchImageZipOptions{}, &buf) + require.NoError(t, err) + require.Equal(t, 3, result.FileCount) + require.Equal(t, 1, limiter.acquireCount) + require.Equal(t, 1, limiter.releaseCount) + + files := readZipFiles(t, buf.Bytes()) + require.Equal(t, []byte("first"), files["images/cover___001.png"]) + require.Equal(t, []byte("second"), files["images/cover___001_2.jpg"]) + require.Equal(t, []byte("third"), files["images/ok_2.webp"]) + require.Contains(t, files, "manifest.json") + require.Contains(t, files, "errors.json") + + zipText := string(bytes.Join(mapValues(files), []byte("\n"))) + require.NotContains(t, zipText, batchImageDownloadTestBase64) + require.NotContains(t, zipText, "provider_job_name") + require.NotContains(t, zipText, "provider_input_ref") + require.NotContains(t, zipText, "gcs_output_uri") + require.NotContains(t, zipText, "account_id") + require.NotContains(t, zipText, "providers/") + require.NotContains(t, zipText, "gs://") + + var manifest struct { + Files []struct { + CustomID string `json:"custom_id"` + Filename string `json:"filename"` + MimeType string `json:"mime_type"` + ImageIndex int `json:"image_index"` + } `json:"files"` + } + require.NoError(t, json.Unmarshal(files["manifest.json"], &manifest)) + require.Len(t, manifest.Files, 3) + require.Equal(t, "images/cover___001_2.jpg", manifest.Files[1].Filename) + require.Equal(t, 1, manifest.Files[1].ImageIndex) + + var errorsJSON []map[string]string + require.NoError(t, json.Unmarshal(files["errors.json"], &errorsJSON)) + require.Len(t, errorsJSON, 1) + require.Equal(t, "bad", errorsJSON[0]["custom_id"]) + require.Equal(t, "SAFETY_BLOCKED", errorsJSON[0]["code"]) + }) + + t.Run("limiter denial returns public limit error", func(t *testing.T) { + svc, _, limiter := newTestBatchImageDownloadService() + limiter.deny = true + var buf bytes.Buffer + + result, err := svc.StreamZip(ctx, testBatchImageOwner(), "imgbatch_download", BatchImageZipOptions{}, &buf) + require.Nil(t, result) + require.ErrorIs(t, err, ErrBatchImageDownloadLimited) + require.Empty(t, buf.Bytes()) + }) + + t.Run("rejects too many zip items before opening output", func(t *testing.T) { + svc, repo, _ := newTestBatchImageDownloadService() + repo.jobs["imgbatch_download"].SuccessCount = 3 + svc.Config.BatchImage.MaxDownloadItemsZip = 1 + var buf bytes.Buffer + + result, err := svc.StreamZip(ctx, testBatchImageOwner(), "imgbatch_download", BatchImageZipOptions{}, &buf) + require.Nil(t, result) + require.ErrorIs(t, err, ErrBatchImageZipTooManyItems) + require.Empty(t, buf.Bytes()) + }) +} + +func TestExtractBatchImagePartsFromResultLine(t *testing.T) { + tests := []struct { + name string + line string + wantID string + wantMime string + wantError string + }{ + {name: "inlineData_mimeType_response", line: `{"key":"a","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageDownloadTestBase64 + `"}}]}}]}}`, wantID: "a", wantMime: "image/png"}, + {name: "inline_data_mime_type_top_level", line: `{"custom_id":"b","candidates":[{"content":{"parts":[{"inline_data":{"mime_type":"image/jpeg","data":"` + batchImageDownloadTestBase64 + `"}}]}}]}`, wantID: "b", wantMime: "image/jpeg"}, + {name: "status_failure", line: `{"key":"c","status":{"code":"INVALID_ARGUMENT","message":"bad prompt"}}`, wantID: "c", wantError: "INVALID_ARGUMENT"}, + {name: "error_failure", line: `{"key":"d","error":{"code":"SAFETY","message":"blocked"}}`, wantID: "d", wantError: "SAFETY_BLOCKED"}, + {name: "empty_output", line: `{"key":"e","response":{"candidates":[{"content":{"parts":[{"text":"none"}]}}]}}`, wantID: "e", wantError: "EMPTY_IMAGE_OUTPUT"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ExtractBatchImagePartsFromResultLine([]byte(tt.line)) + require.NoError(t, err) + require.Equal(t, tt.wantID, got.CustomID) + if tt.wantMime != "" { + require.Len(t, got.Images, 1) + require.Equal(t, tt.wantMime, got.Images[0].MimeType) + require.NotEmpty(t, got.Images[0].Base64Data) + } + if tt.wantError != "" { + require.Equal(t, tt.wantError, got.ErrorCode) + } + }) + } + + _, err := ExtractBatchImagePartsFromResultLine([]byte(`{"response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageDownloadTestBase64 + `"}}]}}]}}`)) + require.Error(t, err) + require.NotContains(t, err.Error(), batchImageDownloadTestBase64) +} + +func TestBatchImageDownloadFilenames(t *testing.T) { + require.Equal(t, "___secret_name.png", BatchImageSafeDownloadFilename("../../secret\nname", "png")) + require.Equal(t, `attachment; filename="cover_001.png"`, BatchImageContentDispositionAttachment(`cover"001.png`)) +} + +func newTestBatchImageDownloadService() (*BatchImageDownloadService, *fakeBatchImageRepository, *fakeBatchImageDownloadLimiter) { + repo := newFakeBatchImageRepository() + apiKeyID := int64(22) + accountID := int64(101) + repo.jobs["imgbatch_download"] = &BatchImageJob{ + BatchID: "imgbatch_download", + UserID: 11, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: BatchImageJobStatusCompleted, + ProviderJobName: batchImageStringPtr("providers/internal/job"), + ProviderOutputRef: batchImageStringPtr("gs://bucket/internal/output.jsonl"), + ItemCount: 3, + SuccessCount: 2, + FailCount: 1, + CreatedAt: time.Now(), + } + mime := "image/png" + ext := "png" + webp := "image/webp" + webpExt := "webp" + code := "SAFETY_BLOCKED" + msg := "blocked in gs://bucket/internal/output.jsonl" + repo.items["imgbatch_download"] = []CreateBatchImageItemParams{ + {JobID: "imgbatch_download", CustomID: "cover/../001", Status: BatchImageItemStatusSuccess, MimeType: &mime, FileExtension: &ext, ImageCount: 2}, + {JobID: "imgbatch_download", CustomID: "bad", Status: BatchImageItemStatusFailed, ErrorCode: &code, ErrorMessage: &msg}, + {JobID: "imgbatch_download", CustomID: "ok_2", Status: BatchImageItemStatusSuccess, MimeType: &webp, FileExtension: &webpExt, ImageCount: 1}, + } + provider := &publicBatchImageProvider{name: BatchImageProviderGeminiAPI, result: batchImageDownloadResultJSONL()} + limiter := &fakeBatchImageDownloadLimiter{} + svc := &BatchImageDownloadService{ + Repo: repo, + ProviderRegistry: NewBatchImageProviderRegistry(provider), + AccountResolver: &fakeBatchImageAccountResolver{account: &Account{ID: accountID, Platform: PlatformGemini, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true}}, + Limiter: limiter, + Config: &config.Config{BatchImage: config.BatchImageConfig{MaxDownloadItemsZip: 10, MaxDownloadDurationSeconds: 60}}, + } + return svc, repo, limiter +} + +const batchImageDownloadTestBase64 = "Zmlyc3Q=" + +func batchImageDownloadResultJSONL() string { + return strings.Join([]string{ + `{"key":"cover/../001","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"Zmlyc3Q="}},{"inlineData":{"mimeType":"image/jpeg","data":"c2Vjb25k"}}]}}]}}`, + `{"key":"bad","error":{"code":"SAFETY","message":"blocked"}}`, + `{"key":"ok_2","candidates":[{"content":{"parts":[{"inline_data":{"mime_type":"image/webp","data":"dGhpcmQ="}}]}}]}`, + }, "\n") + "\n" +} + +func readZipFiles(t *testing.T, data []byte) map[string][]byte { + t.Helper() + reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data))) + require.NoError(t, err) + out := make(map[string][]byte, len(reader.File)) + for _, file := range reader.File { + rc, err := file.Open() + require.NoError(t, err) + body, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + out[file.Name] = body + } + return out +} + +func mapValues(in map[string][]byte) [][]byte { + out := make([][]byte, 0, len(in)) + for _, value := range in { + out = append(out, value) + } + return out +} + +type fakeBatchImageDownloadLimiter struct { + acquireCount int + releaseCount int + deny bool +} + +func (l *fakeBatchImageDownloadLimiter) Acquire(context.Context, string, string) (BatchImageDownloadPermit, error) { + l.acquireCount++ + if l.deny { + return nil, ErrBatchImageDownloadLimited + } + return &fakeBatchImageDownloadPermit{release: func() { l.releaseCount++ }}, nil +} + +type fakeBatchImageDownloadPermit struct { + once bool + release func() +} + +func (p *fakeBatchImageDownloadPermit) Release(context.Context) error { + if p.once { + return nil + } + p.once = true + if p.release != nil { + p.release() + } + return nil +} + +var _ BatchImageDownloadLimiter = (*fakeBatchImageDownloadLimiter)(nil) +var _ BatchImageDownloadPermit = (*fakeBatchImageDownloadPermit)(nil) diff --git a/backend/internal/service/batch_image_mvp_smoke_test.go b/backend/internal/service/batch_image_mvp_smoke_test.go new file mode 100644 index 0000000000..b9daa8f457 --- /dev/null +++ b/backend/internal/service/batch_image_mvp_smoke_test.go @@ -0,0 +1,265 @@ +//go:build unit + +package service + +import ( + "archive/zip" + "bytes" + "context" + "encoding/json" + "io" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestBatchImageMVPFlow(t *testing.T) { + ctx := context.Background() + repo := newFakeBatchImageRepository() + queue := &publicBatchImageQueue{} + provider := &batchImageSmokeProvider{ + name: BatchImageProviderGeminiAPI, + states: []BatchProviderInternalState{ + BatchProviderStateRunning, + BatchProviderStateSucceeded, + }, + result: batchImageSmokeResultJSONL(), + } + accountID := int64(101) + accountRepo := &publicBatchImageAccountRepo{accounts: []Account{testBatchImageAccount(accountID, AccountTypeAPIKey)}} + cfg := &config.Config{BatchImage: config.BatchImageConfig{ + Enabled: true, + MaxItemsPerJobDefault: 10, + MaxPromptCharsPerItem: 8000, + DefaultResponseMimeType: "image/png", + DefaultImageSize: "1K", + MaxDownloadItemsZip: 10, + MaxDownloadDurationSeconds: 60, + OutputRetentionAfterTerminalHours: 72, + }} + registry := NewBatchImageProviderRegistry(provider) + billing := &fakeBatchImageBillingRepo{} + pricing := &fakeBatchImagePricingResolver{unitPrice: 0.25} + owner := testBatchImageOwner() + + publicSvc := &BatchImagePublicService{ + Repo: repo, + AccountRepo: accountRepo, + Queue: queue, + ProviderRegistry: registry, + Pricing: pricing, + BillingRepo: billing, + Config: cfg, + } + processor := &BatchImagePipelineProcessor{ + ProviderProcessor: &BatchImageProviderProcessor{ + Repo: repo, + ProviderRegistry: registry, + AccountResolver: &fakeBatchImageAccountResolver{account: &accountRepo.accounts[0]}, + BillingRepo: billing, + }, + SettlementService: &BatchImageSettlementService{ + Repo: repo, + BillingRepo: billing, + Pricing: pricing, + Config: cfg, + }, + } + downloadSvc := &BatchImageDownloadService{ + Repo: repo, + ProviderRegistry: registry, + AccountResolver: &fakeBatchImageAccountResolver{account: &accountRepo.accounts[0]}, + Limiter: &fakeBatchImageDownloadLimiter{}, + Config: cfg, + } + cleanupSvc := &BatchImageCleanupService{ + Repo: repo, + ProviderRegistry: registry, + AccountResolver: &fakeBatchImageAccountResolver{account: &accountRepo.accounts[0]}, + Config: cfg, + } + + submitted, err := publicSvc.Submit(ctx, owner, validBatchImageSubmitRequest(), "") + require.NoError(t, err) + require.Equal(t, "image.batch", submitted.Object) + require.True(t, strings.HasPrefix(submitted.ID, "imgbatch_")) + require.Equal(t, "queued", submitted.Status) + require.Equal(t, 2, submitted.ItemCount) + require.Equal(t, []string{submitted.ID}, queue.enqueued) + require.Len(t, provider.submits, 1) + require.Len(t, billing.reserves, 1) + require.Equal(t, BatchImageHoldRequestID(submitted.ID), billing.reserves[0].RequestID) + require.InDelta(t, 0.3, billing.reserves[0].HoldAmount, 1e-12) + requireBatchImagePublicJSONHasNoInternals(t, mustMarshalBatchImageSmokeJSON(t, submitted)) + + firstProcess, err := processor.Process(ctx, submitted.ID) + require.NoError(t, err) + require.False(t, firstProcess.Terminal) + require.Equal(t, BatchImageJobStatusRunning, repo.jobs[submitted.ID].Status) + + indexProcess, err := processor.Process(ctx, submitted.ID) + require.NoError(t, err) + require.False(t, indexProcess.Terminal) + require.Equal(t, time.Millisecond, indexProcess.RequeueAfter) + require.Equal(t, BatchImageJobStatusSettling, repo.jobs[submitted.ID].Status) + require.Equal(t, BatchImageCounts{SuccessCount: 1, FailCount: 1}, repo.counts[submitted.ID]) + + settleProcess, err := processor.Process(ctx, submitted.ID) + require.NoError(t, err) + require.True(t, settleProcess.Terminal) + job := repo.jobs[submitted.ID] + require.Equal(t, BatchImageJobStatusCompleted, job.Status) + require.NotNil(t, job.OutputExpiresAt) + require.Equal(t, 1, job.SuccessCount) + require.Equal(t, 1, job.FailCount) + require.Len(t, billing.captures, 1) + require.Equal(t, BatchImageCaptureRequestID(submitted.ID), billing.captures[0].RequestID) + require.InDelta(t, 0.3, billing.captures[0].HoldAmount, 1e-12) + require.InDelta(t, 0.125, billing.captures[0].ActualAmount, 1e-12) + + secondSettlement, err := processor.SettlementService.Settle(ctx, submitted.ID) + require.NoError(t, err) + require.True(t, secondSettlement.AlreadySettled) + require.Len(t, billing.captures, 1) + + status, err := publicSvc.Get(ctx, owner, submitted.ID) + require.NoError(t, err) + require.Equal(t, "completed", status.Status) + require.Equal(t, 1, status.SuccessCount) + require.Equal(t, 1, status.FailCount) + require.NotNil(t, status.ActualCost) + requireBatchImagePublicJSONHasNoInternals(t, mustMarshalBatchImageSmokeJSON(t, status)) + + items, err := publicSvc.ListItems(ctx, owner, submitted.ID, BatchImageItemsQuery{Limit: 100}) + require.NoError(t, err) + require.False(t, items.HasMore) + require.Len(t, items.Data, 2) + require.Equal(t, "cover_001", items.Data[0].CustomID) + require.Equal(t, "succeeded", items.Data[0].Status) + require.Equal(t, "cover_002", items.Data[1].CustomID) + require.Equal(t, "failed", items.Data[1].Status) + require.NotNil(t, items.Data[1].Error) + require.Nil(t, repo.items[submitted.ID][1].BilledAmount) + requireBatchImagePublicJSONHasNoInternals(t, mustMarshalBatchImageSmokeJSON(t, items)) + + stream, err := downloadSvc.OpenItemContent(ctx, owner, submitted.ID, "cover_001", 0) + require.NoError(t, err) + body, err := io.ReadAll(stream.Reader) + require.NoError(t, err) + require.NoError(t, stream.Reader.Close()) + require.Equal(t, []byte("smoke-png"), body) + require.Equal(t, "image/png", stream.ContentType) + require.Equal(t, "cover_001.png", stream.Filename) + + var zipBuf bytes.Buffer + zipResult, err := downloadSvc.StreamZip(ctx, owner, submitted.ID, BatchImageZipOptions{}, &zipBuf) + require.NoError(t, err) + require.Equal(t, 1, zipResult.FileCount) + require.Equal(t, 1, zipResult.ErrorCount) + zipFiles := readZipFiles(t, zipBuf.Bytes()) + require.Equal(t, []byte("smoke-png"), zipFiles["images/cover_001.png"]) + require.Contains(t, zipFiles, "manifest.json") + require.Contains(t, zipFiles, "errors.json") + requireBatchImagePublicJSONHasNoInternals(t, string(bytes.Join(mapValues(zipFiles), []byte("\n")))) + + zipReader, err := zip.NewReader(bytes.NewReader(zipBuf.Bytes()), int64(zipBuf.Len())) + require.NoError(t, err) + require.ElementsMatch(t, []string{"images/cover_001.png", "manifest.json", "errors.json"}, batchImageSmokeZipNames(zipReader)) + + deleted, err := cleanupSvc.DeleteOutputsForOwner(ctx, owner, submitted.ID) + require.NoError(t, err) + require.Equal(t, "output_deleted", deleted.Status) + require.Equal(t, []CleanupTarget{CleanupTargetOutput}, provider.cleanupTargets) + requireBatchImagePublicJSONHasNoInternals(t, mustMarshalBatchImageSmokeJSON(t, deleted)) + + deletedAgain, err := cleanupSvc.DeleteOutputsForOwner(ctx, owner, submitted.ID) + require.NoError(t, err) + require.Equal(t, "output_deleted", deletedAgain.Status) + require.Equal(t, []CleanupTarget{CleanupTargetOutput}, provider.cleanupTargets) + + stream, err = downloadSvc.OpenItemContent(ctx, owner, submitted.ID, "cover_001", 0) + require.Nil(t, stream) + require.ErrorIs(t, err, ErrBatchImageOutputDeleted) + var afterDelete bytes.Buffer + zipResult, err = downloadSvc.StreamZip(ctx, owner, submitted.ID, BatchImageZipOptions{}, &afterDelete) + require.Nil(t, zipResult) + require.ErrorIs(t, err, ErrBatchImageOutputDeleted) + require.Empty(t, afterDelete.Bytes()) +} + +func batchImageSmokeResultJSONL() string { + return strings.Join([]string{ + `{"key":"cover_001","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"c21va2UtcG5n"}}]}}]}}`, + `{"key":"cover_002","status":{"code":3,"message":"blocked by safety policy"}}`, + }, "\n") + "\n" +} + +func mustMarshalBatchImageSmokeJSON(t *testing.T, value any) string { + t.Helper() + body, err := json.Marshal(value) + require.NoError(t, err) + return string(body) +} + +func batchImageSmokeZipNames(reader *zip.Reader) []string { + names := make([]string, 0, len(reader.File)) + for _, file := range reader.File { + names = append(names, file.Name) + } + return names +} + +type batchImageSmokeProvider struct { + name string + states []BatchProviderInternalState + submits []BatchImageInput + result string + cleanupTargets []CleanupTarget +} + +func (p *batchImageSmokeProvider) Name() string { return p.name } + +func (p *batchImageSmokeProvider) SupportsAccount(account *Account) bool { + return account != nil && account.IsSchedulable() +} + +func (p *batchImageSmokeProvider) Submit(_ context.Context, _ *BatchImageJob, _ *Account, input BatchImageInput) (*BatchProviderJob, error) { + p.submits = append(p.submits, input) + return &BatchProviderJob{ + ProviderJobName: "providers/fake-provider-job/raw-id", + ProviderInputRef: "files/fake-provider-job/input.jsonl", + ProviderOutputRef: "files/fake-provider-job/output.jsonl", + }, nil +} + +func (p *batchImageSmokeProvider) Get(context.Context, *BatchImageJob, *Account) (*BatchProviderStatus, error) { + state := BatchProviderStateSucceeded + if len(p.states) > 0 { + state = p.states[0] + p.states = p.states[1:] + } + return &BatchProviderStatus{ + RawState: strings.ToUpper(string(state)), + InternalState: state, + Done: state == BatchProviderStateSucceeded, + ProviderOutputRef: "files/fake-provider-job/output.jsonl", + }, nil +} + +func (p *batchImageSmokeProvider) Cancel(context.Context, *BatchImageJob, *Account) error { + return nil +} + +func (p *batchImageSmokeProvider) OpenResult(context.Context, *BatchImageJob, *Account) (io.ReadCloser, string, error) { + return io.NopCloser(strings.NewReader(p.result)), "application/jsonl", nil +} + +func (p *batchImageSmokeProvider) Cleanup(_ context.Context, _ *BatchImageJob, _ *Account, target CleanupTarget) error { + p.cleanupTargets = append(p.cleanupTargets, target) + return nil +} + +var _ BatchImageProvider = (*batchImageSmokeProvider)(nil) diff --git a/backend/internal/service/batch_image_processor.go b/backend/internal/service/batch_image_processor.go new file mode 100644 index 0000000000..414890bd88 --- /dev/null +++ b/backend/internal/service/batch_image_processor.go @@ -0,0 +1,682 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sort" + "strconv" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + BatchImageParsedStatusSucceeded = "succeeded" + BatchImageParsedStatusFailed = "failed" + + defaultBatchImageProcessorRequeue = 30 * time.Second + batchImageProviderErrorRequeue = time.Minute + batchImageMaxErrorMessageLength = 1000 +) + +type BatchImageAccountResolver interface { + ResolveBatchImageAccount(ctx context.Context, accountID int64) (*Account, error) +} + +type BatchImageAccountLookup interface { + GetByID(ctx context.Context, id int64) (*Account, error) +} + +type BatchImageAccountRepositoryResolver struct { + Repo BatchImageAccountLookup +} + +func (r *BatchImageAccountRepositoryResolver) ResolveBatchImageAccount(ctx context.Context, accountID int64) (*Account, error) { + if r == nil || r.Repo == nil { + return nil, ErrAccountNotFound + } + return r.Repo.GetByID(ctx, accountID) +} + +type BatchImageProviderProcessor struct { + Repo BatchImageRepository + ProviderRegistry *BatchImageProviderRegistry + AccountResolver BatchImageAccountResolver + Indexer *BatchImageResultIndexer + BillingRepo UsageBillingRepository + AuthCache APIKeyAuthCacheInvalidator + DefaultRequeue time.Duration +} + +func (p *BatchImageProviderProcessor) Process(ctx context.Context, batchID string) (BatchImageProcessResult, error) { + if p == nil || p.Repo == nil || p.ProviderRegistry == nil || p.AccountResolver == nil { + return BatchImageProcessResult{}, infraerrors.New(http.StatusInternalServerError, "BATCH_IMAGE_PROCESSOR_NOT_CONFIGURED", "batch image processor is not configured") + } + + job, err := p.Repo.GetBatchImageJobByBatchID(ctx, batchID) + if err != nil { + return BatchImageProcessResult{}, err + } + if isBatchImageProcessorDoneStatus(job.Status) { + if err := p.releaseTerminalHold(ctx, job); err != nil { + return BatchImageProcessResult{}, err + } + return BatchImageProcessResult{Terminal: true}, nil + } + + provider, ok := p.ProviderRegistry.Get(job.Provider) + if !ok || provider == nil { + return BatchImageProcessResult{}, ErrBatchImageUnsupportedProvider + } + if job.AccountID == nil || *job.AccountID <= 0 { + return BatchImageProcessResult{}, ErrBatchImageMissingAccountID + } + account, err := p.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID) + if err != nil { + return BatchImageProcessResult{}, err + } + if !provider.SupportsAccount(account) { + return BatchImageProcessResult{}, ErrBatchImageProviderUnsupportedAccount + } + if strings.TrimSpace(batchImageDerefString(job.ProviderJobName)) == "" { + return BatchImageProcessResult{}, ErrBatchImageMissingProviderJobName + } + + if job.Status == BatchImageJobStatusIndexing { + return p.indexAndSettle(ctx, job, provider, account) + } + + status, err := provider.Get(ctx, job, account) + if err != nil { + logger.L().Warn("batch_image.provider_status_check_failed", + zap.String("batch_id", job.BatchID), + zap.String("provider", job.Provider), + zap.String("provider_job_name", batchImageDerefString(job.ProviderJobName)), + zap.Error(err), + ) + return BatchImageProcessResult{RequeueAfter: batchImageProviderErrorRequeue}, nil + } + if status == nil { + return BatchImageProcessResult{RequeueAfter: p.requeueDelay(0)}, nil + } + if err := p.persistProviderOutputRef(ctx, job, status.ProviderOutputRef); err != nil { + return BatchImageProcessResult{}, err + } + + switch status.InternalState { + case BatchProviderStateQueued: + return BatchImageProcessResult{RequeueAfter: p.requeueDelay(status.SuggestedRequeueAfter)}, nil + case BatchProviderStateRunning: + if job.Status != BatchImageJobStatusRunning { + if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusRunning, BatchImageTransitionOptions{ + EventType: "provider_status_checked", + EventPayload: map[string]any{"provider_state": status.RawState}, + }); err != nil { + return BatchImageProcessResult{}, err + } + job.Status = BatchImageJobStatusRunning + } + return BatchImageProcessResult{RequeueAfter: p.requeueDelay(status.SuggestedRequeueAfter)}, nil + case BatchProviderStateSucceeded: + if job.Status != BatchImageJobStatusIndexing { + if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusIndexing, BatchImageTransitionOptions{ + EventType: "indexing_started", + EventPayload: map[string]any{"provider_state": status.RawState}, + }); err != nil { + return BatchImageProcessResult{}, err + } + job.Status = BatchImageJobStatusIndexing + } + return p.indexAndSettle(ctx, job, provider, account) + case BatchProviderStateFailed, BatchProviderStateExpired: + code := strings.TrimSpace(status.ErrorCode) + if code == "" && status.InternalState == BatchProviderStateExpired { + code = "PROVIDER_BATCH_EXPIRED" + } + if code == "" { + code = "PROVIDER_BATCH_FAILED" + } + msg := truncateBatchImageMessage(status.ErrorMessage, batchImageMaxErrorMessageLength) + if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{ + EventType: "job_failed", + EventPayload: map[string]any{"provider_state": status.RawState, "error_code": code}, + ErrorCode: batchImageStringPtr(code), + ErrorMessage: batchImageOptionalStringPtr(msg), + }); err != nil { + return BatchImageProcessResult{}, err + } + job.Status = BatchImageJobStatusFailed + if err := p.releaseTerminalHold(ctx, job); err != nil { + return BatchImageProcessResult{}, err + } + return BatchImageProcessResult{Terminal: true}, nil + case BatchProviderStateCancelled: + if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusCancelled, BatchImageTransitionOptions{ + EventType: "job_failed", + EventPayload: map[string]any{"provider_state": status.RawState, "error_code": "PROVIDER_BATCH_CANCELLED"}, + }); err != nil { + return BatchImageProcessResult{}, err + } + job.Status = BatchImageJobStatusCancelled + if err := p.releaseTerminalHold(ctx, job); err != nil { + return BatchImageProcessResult{}, err + } + return BatchImageProcessResult{Terminal: true}, nil + default: + return BatchImageProcessResult{RequeueAfter: p.requeueDelay(status.SuggestedRequeueAfter)}, nil + } +} + +func (p *BatchImageProviderProcessor) indexAndSettle(ctx context.Context, job *BatchImageJob, provider BatchImageProvider, account *Account) (BatchImageProcessResult, error) { + indexer := p.Indexer + if indexer == nil { + indexer = &BatchImageResultIndexer{Repo: p.Repo} + } + if indexer.Repo == nil { + indexer.Repo = p.Repo + } + + result, err := indexer.Index(ctx, job, provider, account) + if err != nil { + if errors.Is(err, ErrBatchImageIndexOutputMissing) { + return BatchImageProcessResult{}, err + } + // job 状态已被并发方推进(如已进入 settling/终态):不是索引数据问题, + // 短延迟 requeue 让下一轮按最新状态处理,不能误转 failed。 + if errors.Is(err, ErrBatchImageIndexStateConflict) { + return BatchImageProcessResult{RequeueAfter: time.Millisecond}, nil + } + code := "INDEX_PARSE_FAILED" + if errors.Is(err, ErrBatchImageDuplicateCustomID) { + code = "DUPLICATE_CUSTOM_ID_IN_OUTPUT" + } + msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength) + transitionErr := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{ + EventType: "indexing_failed", + EventPayload: map[string]any{"error_code": code}, + ErrorCode: batchImageStringPtr(code), + ErrorMessage: batchImageOptionalStringPtr(msg), + }) + if transitionErr != nil { + return BatchImageProcessResult{}, transitionErr + } + job.Status = BatchImageJobStatusFailed + if err := p.releaseTerminalHold(ctx, job); err != nil { + return BatchImageProcessResult{}, err + } + return BatchImageProcessResult{Terminal: true}, nil + } + + if err := p.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusSettling, BatchImageTransitionOptions{ + EventType: "indexing_completed", + EventPayload: map[string]any{ + "success_count": result.SuccessCount, + "fail_count": result.FailCount, + "total_count": result.TotalCount, + }, + }); err != nil { + return BatchImageProcessResult{}, err + } + return BatchImageProcessResult{RequeueAfter: time.Millisecond}, nil +} + +func (p *BatchImageProviderProcessor) releaseTerminalHold(ctx context.Context, job *BatchImageJob) error { + if p == nil || job == nil { + return nil + } + if job.Status != BatchImageJobStatusFailed && job.Status != BatchImageJobStatusCancelled { + return nil + } + if err := releaseBatchImageBalanceHold(ctx, p.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil { + return err + } + if p.AuthCache != nil && job.UserID > 0 { + p.AuthCache.InvalidateAuthCacheByUserID(ctx, job.UserID) + } + return nil +} + +func (p *BatchImageProviderProcessor) persistProviderOutputRef(ctx context.Context, job *BatchImageJob, ref string) error { + ref = strings.TrimSpace(ref) + if ref == "" || job == nil || batchImageDerefString(job.ProviderOutputRef) == ref { + return nil + } + if err := p.Repo.UpdateBatchImageJobProviderOutputRef(ctx, job.BatchID, ref); err != nil { + return err + } + job.ProviderOutputRef = &ref + return nil +} + +func (p *BatchImageProviderProcessor) requeueDelay(suggested time.Duration) time.Duration { + if suggested > 0 { + return suggested + } + if p != nil && p.DefaultRequeue > 0 { + return p.DefaultRequeue + } + return defaultBatchImageProcessorRequeue +} + +func isBatchImageProcessorDoneStatus(status string) bool { + if status == BatchImageJobStatusSettling { + return true + } + return IsTerminalBatchImageJobStatus(status) +} + +type BatchImageIndexResult struct { + SuccessCount int + FailCount int + TotalCount int +} + +type BatchImageResultIndexer struct { + Repo BatchImageRepository +} + +func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob, provider BatchImageProvider, account *Account) (*BatchImageIndexResult, error) { + if i == nil || i.Repo == nil || job == nil || provider == nil { + return nil, ErrBatchImageIndexOutputMissing + } + expected, err := i.listExpectedCustomIDs(ctx, job.BatchID) + if err != nil { + return nil, err + } + + r, _, err := provider.OpenResult(ctx, job, account) + if err != nil { + return nil, ErrBatchImageIndexOutputMissing.WithCause(err) + } + defer func() { _ = r.Close() }() + + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + + seen := make(map[string]int) + unknownCount := 0 + var items []CreateBatchImageItemParams + result := &BatchImageIndexResult{} + lineNumber := 0 + now := time.Now() + sourceObject := batchImageDerefString(job.ProviderOutputRef) + if sourceObject == "" { + sourceObject = batchImageDerefString(job.ProviderJobName) + } + + for scanner.Scan() { + lineNumber++ + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + parsed, err := ParseBatchImageResultLine([]byte(line), lineNumber) + if err != nil { + return nil, err + } + // 与提交时的 custom_id 集对账:provider 输出中未知/多余的行不能进入 item 表, + // 否则 success+fail > item_count 会让结算永远校验失败。 + if len(expected) > 0 { + if _, ok := expected[parsed.CustomID]; !ok { + unknownCount++ + continue + } + } + if firstLine, ok := seen[parsed.CustomID]; ok { + return nil, ErrBatchImageDuplicateCustomID.WithCause(fmt.Errorf("custom id %q duplicated at lines %d and %d", parsed.CustomID, firstLine, lineNumber)) + } + seen[parsed.CustomID] = lineNumber + + lineNo := parsed.SourceLineNumber + item := CreateBatchImageItemParams{ + JobID: job.BatchID, + CustomID: parsed.CustomID, + Status: BatchImageItemStatusFailed, + ProviderSourceObject: batchImageOptionalStringPtr(sourceObject), + SourceLineNumber: &lineNo, + ImageCount: parsed.ImageCount, + IndexedAt: &now, + } + if parsed.Status == BatchImageParsedStatusSucceeded { + item.Status = BatchImageItemStatusSuccess + item.MimeType = batchImageOptionalStringPtr(parsed.MimeType) + item.FileExtension = batchImageOptionalStringPtr(parsed.FileExtension) + result.SuccessCount++ + } else { + item.ErrorCode = batchImageOptionalStringPtr(parsed.ErrorCode) + item.ErrorMessage = batchImageOptionalStringPtr(parsed.ErrorMessage) + result.FailCount++ + } + items = append(items, item) + result.TotalCount++ + } + if err := scanner.Err(); err != nil { + if errors.Is(err, io.ErrUnexpectedEOF) { + return nil, ErrBatchImageIndexParseFailed.WithCause(err) + } + return nil, err + } + // 输出中漏掉的已提交项必须补失败记录,而不是静默消失: + // 否则用户看不到该项,且只按成功数计费会掩盖 provider 的丢单。 + missingCount := 0 + if len(expected) > 0 { + missingIDs := make([]string, 0) + for customID := range expected { + if _, ok := seen[customID]; !ok { + missingIDs = append(missingIDs, customID) + } + } + sort.Strings(missingIDs) + for _, customID := range missingIDs { + items = append(items, CreateBatchImageItemParams{ + JobID: job.BatchID, + CustomID: customID, + Status: BatchImageItemStatusFailed, + ProviderSourceObject: batchImageOptionalStringPtr(sourceObject), + ErrorCode: batchImageStringPtr("PROVIDER_RESULT_MISSING"), + ErrorMessage: batchImageStringPtr("provider output did not include a result for this item"), + IndexedAt: &now, + }) + result.FailCount++ + result.TotalCount++ + } + missingCount = len(missingIDs) + } + if result.TotalCount == 0 { + return nil, ErrBatchImageIndexNoResultLines + } + if unknownCount > 0 || missingCount > 0 { + logger.L().Warn("batch_image.index_reconciled", + zap.String("batch_id", job.BatchID), + zap.Int("unknown_custom_ids", unknownCount), + zap.Int("missing_custom_ids", missingCount), + ) + if err := i.Repo.AppendBatchImageEvent(ctx, job.BatchID, "index_reconciled", map[string]any{ + "batch_id": job.BatchID, + "unknown_custom_ids": unknownCount, + "missing_custom_ids": missingCount, + }); err != nil { + logger.L().Warn("batch_image.index_reconcile_event_failed", + zap.String("batch_id", job.BatchID), + zap.Error(err), + ) + } + } + if err := i.Repo.ReplaceBatchImageItemsForJob(ctx, job.BatchID, items, BatchImageCounts{ + SuccessCount: result.SuccessCount, + FailCount: result.FailCount, + }); err != nil { + return nil, err + } + return result, nil +} + +// listExpectedCustomIDs 返回该 job 当前 item 表中的全部 custom_id 集合, +// 即提交时预创建(或上一轮索引重建)的完整条目清单,用于与 provider 输出对账。 +func (i *BatchImageResultIndexer) listExpectedCustomIDs(ctx context.Context, batchID string) (map[string]struct{}, error) { + const pageSize = 500 + expected := make(map[string]struct{}) + offset := 0 + for { + page, err := i.Repo.ListBatchImageItems(ctx, batchID, BatchImageItemFilter{Limit: pageSize, Offset: offset}) + if err != nil { + return nil, err + } + for _, item := range page { + if item != nil { + expected[item.CustomID] = struct{}{} + } + } + if len(page) < pageSize { + return expected, nil + } + offset += len(page) + } +} + +type ParsedBatchImageResult struct { + CustomID string + Status string + MimeType string + FileExtension string + ImageCount int + + ErrorCode string + ErrorMessage string + + SourceLineNumber int +} + +func ParseBatchImageResultLine(line []byte, lineNumber int) (*ParsedBatchImageResult, error) { + var obj map[string]any + if err := json.Unmarshal(line, &obj); err != nil { + return nil, ErrBatchImageIndexParseFailed.WithCause(fmt.Errorf("line %d: %w", lineNumber, err)) + } + + customID := batchImageFirstNonEmptyString( + batchImageMapString(obj, "key"), + batchImageMapString(obj, "custom_id"), + batchImageMapString(obj, "customId"), + batchImageNestedString(obj, "request", "key"), + ) + if customID == "" { + return nil, ErrBatchImageIndexParseFailed.WithCause(fmt.Errorf("line %d: missing custom id", lineNumber)) + } + + parsed := &ParsedBatchImageResult{ + CustomID: customID, + SourceLineNumber: lineNumber, + } + imageCount, mimeType := batchImageFindImageParts(obj) + if imageCount > 0 { + parsed.Status = BatchImageParsedStatusSucceeded + parsed.ImageCount = imageCount + parsed.MimeType = mimeType + parsed.FileExtension = batchImageFileExtension(mimeType) + return parsed, nil + } + + if code, message, ok := batchImageFailureFromProviderFields(obj); ok { + parsed.Status = BatchImageParsedStatusFailed + parsed.ErrorCode = code + parsed.ErrorMessage = truncateBatchImageMessage(message, batchImageMaxErrorMessageLength) + return parsed, nil + } + + if _, hasResponse := obj["response"]; hasResponse || batchImageHasCandidates(obj) { + parsed.Status = BatchImageParsedStatusFailed + parsed.ErrorCode = "EMPTY_IMAGE_OUTPUT" + parsed.ErrorMessage = "provider response contained no image output" + return parsed, nil + } + + parsed.Status = BatchImageParsedStatusFailed + parsed.ErrorCode = "PROVIDER_ITEM_FAILED" + parsed.ErrorMessage = "provider result line contained no image output" + return parsed, nil +} + +func batchImageFindImageParts(obj map[string]any) (int, string) { + count, mimeType := batchImageFindImagePartsInCandidates(batchImageNestedAny(obj, "response", "candidates")) + if count > 0 { + return count, mimeType + } + return batchImageFindImagePartsInCandidates(obj["candidates"]) +} + +func batchImageFindImagePartsInCandidates(raw any) (int, string) { + candidates, ok := raw.([]any) + if !ok { + return 0, "" + } + count := 0 + firstMime := "" + for _, candidateRaw := range candidates { + candidate, ok := candidateRaw.(map[string]any) + if !ok { + continue + } + partsRaw := batchImageNestedAny(candidate, "content", "parts") + parts, ok := partsRaw.([]any) + if !ok { + continue + } + for _, partRaw := range parts { + part, ok := partRaw.(map[string]any) + if !ok { + continue + } + inline, ok := firstMap(part["inlineData"], part["inline_data"]) + if !ok { + continue + } + data := strings.TrimSpace(batchImageMapString(inline, "data")) + mime := batchImageFirstNonEmptyString(batchImageMapString(inline, "mimeType"), batchImageMapString(inline, "mime_type")) + if data == "" || !strings.HasPrefix(strings.ToLower(strings.TrimSpace(mime)), "image/") { + continue + } + count++ + if firstMime == "" { + firstMime = strings.TrimSpace(mime) + } + } + } + return count, firstMime +} + +func batchImageFailureFromProviderFields(obj map[string]any) (string, string, bool) { + if status, ok := obj["status"].(map[string]any); ok { + message := batchImageFirstNonEmptyString(batchImageMapString(status, "message"), batchImageMapString(status, "details")) + code := batchImageFirstNonEmptyString(batchImageMapString(status, "code"), batchImageMapString(status, "status")) + return batchImageMapFailureCode(code, message), message, true + } + if errObj, ok := obj["error"].(map[string]any); ok { + message := batchImageFirstNonEmptyString(batchImageMapString(errObj, "message"), batchImageMapString(errObj, "details")) + code := batchImageFirstNonEmptyString(batchImageMapString(errObj, "code"), batchImageMapString(errObj, "status")) + return batchImageMapFailureCode(code, message), message, true + } + return "", "", false +} + +func batchImageMapFailureCode(code, message string) string { + text := strings.ToLower(strings.TrimSpace(code + " " + message)) + switch { + case strings.Contains(text, "safety"), strings.Contains(text, "policy"), strings.Contains(text, "blocked"), strings.Contains(text, "prohibited"): + return "SAFETY_BLOCKED" + case strings.Contains(text, "invalid_argument"), strings.Contains(text, "invalid argument"), strings.Contains(text, "bad request"): + return "INVALID_ARGUMENT" + case strings.Contains(text, "quota"), strings.Contains(text, "rate"), strings.Contains(text, "resource_exhausted"), strings.Contains(text, "too many requests"): + return "PROVIDER_RATE_LIMITED" + default: + return "PROVIDER_ITEM_FAILED" + } +} + +func batchImageFileExtension(mimeType string) string { + switch strings.ToLower(strings.TrimSpace(mimeType)) { + case "image/png": + return "png" + case "image/jpeg", "image/jpg": + return "jpg" + case "image/webp": + return "webp" + default: + return "" + } +} + +func batchImageHasCandidates(obj map[string]any) bool { + if _, ok := obj["candidates"]; ok { + return true + } + _, ok := batchImageNestedAny(obj, "response", "candidates").([]any) + return ok +} + +func batchImageMapString(m map[string]any, key string) string { + if m == nil { + return "" + } + switch v := m[key].(type) { + case string: + return strings.TrimSpace(v) + case json.Number: + return v.String() + case float64: + return strconv.FormatInt(int64(v), 10) + default: + return "" + } +} + +func batchImageNestedString(m map[string]any, keys ...string) string { + if nested, ok := batchImageNestedAny(m, keys...).(string); ok { + return strings.TrimSpace(nested) + } + return "" +} + +func batchImageNestedAny(m map[string]any, keys ...string) any { + var current any = m + for _, key := range keys { + cm, ok := current.(map[string]any) + if !ok { + return nil + } + current = cm[key] + } + return current +} + +func firstMap(values ...any) (map[string]any, bool) { + for _, value := range values { + if m, ok := value.(map[string]any); ok { + return m, true + } + } + return nil, false +} + +func batchImageFirstNonEmptyString(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} + +func batchImageDerefString(v *string) string { + if v == nil { + return "" + } + return strings.TrimSpace(*v) +} + +func batchImageStringPtr(v string) *string { + return &v +} + +func batchImageOptionalStringPtr(v string) *string { + v = strings.TrimSpace(v) + if v == "" { + return nil + } + return &v +} + +func truncateBatchImageMessage(message string, limit int) string { + message = strings.TrimSpace(message) + if limit <= 0 || len(message) <= limit { + return message + } + return message[:limit] +} diff --git a/backend/internal/service/batch_image_processor_test.go b/backend/internal/service/batch_image_processor_test.go new file mode 100644 index 0000000000..f54a1a99ec --- /dev/null +++ b/backend/internal/service/batch_image_processor_test.go @@ -0,0 +1,929 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + "testing" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/stretchr/testify/require" +) + +const batchImageTestData = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ" + +func TestParseBatchImageResultLine_SuccessShapes(t *testing.T) { + tests := []struct { + name string + line string + wantID string + wantMime string + wantExt string + wantCount int + }{ + { + name: "gemini_inlineData", + line: `{"key":"cover_001","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageTestData + `"}}]}}]}}`, + wantID: "cover_001", wantMime: "image/png", wantExt: "png", wantCount: 1, + }, + { + name: "snake_case_inline_data", + line: `{"custom_id":"cover_002","response":{"candidates":[{"content":{"parts":[{"inline_data":{"mime_type":"image/jpeg","data":"` + batchImageTestData + `"}}]}}]}}`, + wantID: "cover_002", wantMime: "image/jpeg", wantExt: "jpg", wantCount: 1, + }, + { + name: "vertex_top_level_response", + line: `{"customId":"cover_003","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/webp","data":"` + batchImageTestData + `"}}]}}]}}`, + wantID: "cover_003", wantMime: "image/webp", wantExt: "webp", wantCount: 1, + }, + { + name: "top_level_candidates", + line: `{"request":{"key":"cover_004"},"candidates":[{"content":{"parts":[{"inline_data":{"mime_type":"image/png","data":"` + batchImageTestData + `"}},{"inlineData":{"mimeType":"image/png","data":"` + batchImageTestData + `"}}]}}]}`, + wantID: "cover_004", wantMime: "image/png", wantExt: "png", wantCount: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ParseBatchImageResultLine([]byte(tt.line), 7) + require.NoError(t, err) + require.Equal(t, tt.wantID, got.CustomID) + require.Equal(t, BatchImageParsedStatusSucceeded, got.Status) + require.Equal(t, tt.wantMime, got.MimeType) + require.Equal(t, tt.wantExt, got.FileExtension) + require.Equal(t, tt.wantCount, got.ImageCount) + require.Equal(t, 7, got.SourceLineNumber) + require.NotContains(t, fmt.Sprintf("%+v", got), batchImageTestData) + }) + } +} + +func TestParseBatchImageResultLine_FailureShapes(t *testing.T) { + tests := []struct { + name string + line string + wantCode string + }{ + {name: "status_row", line: `{"key":"cover_001","status":{"code":3,"message":"invalid argument: bad prompt"}}`, wantCode: "INVALID_ARGUMENT"}, + {name: "error_row", line: `{"key":"cover_002","error":{"code":"SAFETY","message":"blocked by safety policy"}}`, wantCode: "SAFETY_BLOCKED"}, + {name: "quota_row", line: `{"key":"cover_003","error":{"code":"RESOURCE_EXHAUSTED","message":"quota exceeded"}}`, wantCode: "PROVIDER_RATE_LIMITED"}, + {name: "empty_image_output", line: `{"key":"cover_004","response":{"candidates":[{"content":{"parts":[{"text":"no image"}]}}]}}`, wantCode: "EMPTY_IMAGE_OUTPUT"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ParseBatchImageResultLine([]byte(tt.line), 1) + require.NoError(t, err) + require.Equal(t, BatchImageParsedStatusFailed, got.Status) + require.Equal(t, tt.wantCode, got.ErrorCode) + }) + } +} + +func TestParseBatchImageResultLine_RejectsMissingCustomIDAndDoesNotLeakData(t *testing.T) { + _, err := ParseBatchImageResultLine([]byte(`{"response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"`+batchImageTestData+`"}}]}}]}}`), 3) + require.ErrorIs(t, err, ErrBatchImageIndexParseFailed) + require.NotContains(t, err.Error(), batchImageTestData) +} + +func TestBatchImageResultIndexer_WritesCountsAndReplacesItems(t *testing.T) { + output := strings.Join([]string{ + `{"key":"ok","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageTestData + `"}}]}}]}}`, + `{"key":"bad","error":{"code":"SAFETY","message":"blocked by safety policy"}}`, + }, "\n") + "\n" + repo := newFakeBatchImageRepository() + outputRef := "files/output" + job := &BatchImageJob{BatchID: "imgbatch_index", ProviderOutputRef: &outputRef} + provider := &fakeProcessorProvider{result: output} + + result, err := (&BatchImageResultIndexer{Repo: repo}).Index(context.Background(), job, provider, &Account{}) + require.NoError(t, err) + require.True(t, provider.openResultCalled) + require.Equal(t, 1, result.SuccessCount) + require.Equal(t, 1, result.FailCount) + require.Equal(t, 2, result.TotalCount) + require.Equal(t, 1, repo.replaceCalls) + require.Len(t, repo.items[job.BatchID], 2) + require.Equal(t, BatchImageItemStatusSuccess, repo.items[job.BatchID][0].Status) + require.Equal(t, BatchImageItemStatusFailed, repo.items[job.BatchID][1].Status) + require.Equal(t, BatchImageCounts{SuccessCount: 1, FailCount: 1}, repo.counts[job.BatchID]) + require.NotContains(t, fmt.Sprintf("%+v", repo.items[job.BatchID]), batchImageTestData) + + // 重新索引时与现有 custom_id 集对账:未知的 "ok2" 被丢弃, + // 输出中缺失的 ok/bad 补为 PROVIDER_RESULT_MISSING 失败记录。 + provider.result = `{"key":"ok2","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/webp","data":"` + batchImageTestData + `"}}]}}]}}` + "\n" + result, err = (&BatchImageResultIndexer{Repo: repo}).Index(context.Background(), job, provider, &Account{}) + require.NoError(t, err) + require.Equal(t, 2, result.TotalCount) + require.Equal(t, 0, result.SuccessCount) + require.Equal(t, 2, result.FailCount) + require.Len(t, repo.items[job.BatchID], 2) + gotIDs := []string{repo.items[job.BatchID][0].CustomID, repo.items[job.BatchID][1].CustomID} + require.ElementsMatch(t, []string{"ok", "bad"}, gotIDs) + for _, item := range repo.items[job.BatchID] { + require.Equal(t, BatchImageItemStatusFailed, item.Status) + require.Equal(t, "PROVIDER_RESULT_MISSING", batchImageDerefString(item.ErrorCode)) + } +} + +func TestBatchImageResultIndexer_ReconcilesMissingAndUnknownCustomIDs(t *testing.T) { + repo := newFakeBatchImageRepository() + outputRef := "files/output" + job := &BatchImageJob{BatchID: "imgbatch_reconcile", ProviderOutputRef: &outputRef, ItemCount: 3} + // 预创建提交时的 pending 条目(提交流程的行为)。 + require.NoError(t, repo.BulkCreateBatchImageItems(context.Background(), []CreateBatchImageItemParams{ + {JobID: job.BatchID, CustomID: "a", Status: BatchImageItemStatusPending}, + {JobID: job.BatchID, CustomID: "b", Status: BatchImageItemStatusPending}, + {JobID: job.BatchID, CustomID: "c", Status: BatchImageItemStatusPending}, + })) + // provider 输出:a 成功,b 失败,c 漏掉,多出未知的 x。 + output := strings.Join([]string{ + `{"key":"a","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageTestData + `"}}]}}]}}`, + `{"key":"b","error":{"code":"SAFETY","message":"blocked"}}`, + `{"key":"x","error":{"code":"UNKNOWN","message":"not ours"}}`, + }, "\n") + "\n" + provider := &fakeProcessorProvider{result: output} + + result, err := (&BatchImageResultIndexer{Repo: repo}).Index(context.Background(), job, provider, &Account{}) + require.NoError(t, err) + require.Equal(t, 3, result.TotalCount) + require.Equal(t, 1, result.SuccessCount) + require.Equal(t, 2, result.FailCount) + require.Len(t, repo.items[job.BatchID], 3) + byID := make(map[string]CreateBatchImageItemParams) + for _, item := range repo.items[job.BatchID] { + byID[item.CustomID] = item + } + require.NotContains(t, byID, "x") + require.Equal(t, BatchImageItemStatusSuccess, byID["a"].Status) + require.Equal(t, BatchImageItemStatusFailed, byID["b"].Status) + require.Equal(t, BatchImageItemStatusFailed, byID["c"].Status) + require.Equal(t, "PROVIDER_RESULT_MISSING", batchImageDerefString(byID["c"].ErrorCode)) + // 对账后 success+fail == item_count,结算计数校验可通过。 + require.Equal(t, job.ItemCount, result.SuccessCount+result.FailCount) +} + +func TestBatchImageResultIndexer_EmptyInvalidAndDuplicateOutput(t *testing.T) { + tests := []struct { + name string + body string + want error + }{ + {name: "empty", body: "\n", want: ErrBatchImageIndexNoResultLines}, + {name: "invalid_json", body: "{bad-json}\n", want: ErrBatchImageIndexParseFailed}, + {name: "duplicate_custom_id", body: `{"key":"dup","error":{"message":"one"}}` + "\n" + `{"key":"dup","error":{"message":"two"}}` + "\n", want: ErrBatchImageDuplicateCustomID}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := newFakeBatchImageRepository() + _, err := (&BatchImageResultIndexer{Repo: repo}).Index(context.Background(), &BatchImageJob{BatchID: "imgbatch_bad"}, &fakeProcessorProvider{result: tt.body}, &Account{}) + require.ErrorIs(t, err, tt.want) + require.Empty(t, repo.items["imgbatch_bad"]) + }) + } +} + +func TestBatchImageProviderProcessor_ValidationAndTerminalCases(t *testing.T) { + ctx := context.Background() + accountID := int64(10) + providerJob := "providers/job" + + t.Run("terminal job returns without provider call", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_done"] = &BatchImageJob{BatchID: "imgbatch_done", Status: BatchImageJobStatusFailed} + provider := &fakeProcessorProvider{} + got, err := (&BatchImageProviderProcessor{ + Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(provider), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}, + }).Process(ctx, "imgbatch_done") + require.NoError(t, err) + require.True(t, got.Terminal) + require.False(t, provider.getCalled) + }) + + t.Run("missing provider", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_missing_provider"] = &BatchImageJob{BatchID: "imgbatch_missing_provider", Status: BatchImageJobStatusSubmitted, Provider: "missing", AccountID: &accountID, ProviderJobName: &providerJob} + _, err := (&BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}}).Process(ctx, "imgbatch_missing_provider") + require.ErrorIs(t, err, ErrBatchImageUnsupportedProvider) + }) + + t.Run("missing account id", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_missing_account"] = &BatchImageJob{BatchID: "imgbatch_missing_account", Status: BatchImageJobStatusSubmitted, Provider: "fake", ProviderJobName: &providerJob} + _, err := (&BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}}).Process(ctx, "imgbatch_missing_account") + require.ErrorIs(t, err, ErrBatchImageMissingAccountID) + }) + + t.Run("missing provider job name", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_missing_name"] = &BatchImageJob{BatchID: "imgbatch_missing_name", Status: BatchImageJobStatusSubmitted, Provider: "fake", AccountID: &accountID} + _, err := (&BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}}).Process(ctx, "imgbatch_missing_name") + require.ErrorIs(t, err, ErrBatchImageMissingProviderJobName) + }) +} + +func TestBatchImageProviderProcessor_StatusFlow(t *testing.T) { + ctx := context.Background() + accountID := int64(10) + providerJob := "providers/job" + newJob := func(status string) *BatchImageJob { + return &BatchImageJob{BatchID: "imgbatch_flow", Status: status, Provider: "fake", AccountID: &accountID, ProviderJobName: &providerJob} + } + + t.Run("running status updates and requeues", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_flow"] = newJob(BatchImageJobStatusSubmitted) + provider := &fakeProcessorProvider{status: &BatchProviderStatus{InternalState: BatchProviderStateRunning, RawState: "RUNNING", SuggestedRequeueAfter: 12 * time.Second}} + got, err := newTestBatchImageProcessor(repo, provider).Process(ctx, "imgbatch_flow") + require.NoError(t, err) + require.False(t, got.Terminal) + require.Equal(t, 12*time.Second, got.RequeueAfter) + require.Equal(t, BatchImageJobStatusRunning, repo.jobs["imgbatch_flow"].Status) + }) + + t.Run("queued status requeues", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_flow"] = newJob(BatchImageJobStatusSubmitted) + provider := &fakeProcessorProvider{status: &BatchProviderStatus{InternalState: BatchProviderStateQueued}} + got, err := newTestBatchImageProcessor(repo, provider).Process(ctx, "imgbatch_flow") + require.NoError(t, err) + require.False(t, got.Terminal) + require.Equal(t, defaultBatchImageProcessorRequeue, got.RequeueAfter) + require.Equal(t, BatchImageJobStatusSubmitted, repo.jobs["imgbatch_flow"].Status) + }) + + t.Run("transient provider get error requeues", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_flow"] = newJob(BatchImageJobStatusSubmitted) + provider := &fakeProcessorProvider{getErr: errors.New("temporary upstream failure")} + got, err := newTestBatchImageProcessor(repo, provider).Process(ctx, "imgbatch_flow") + require.NoError(t, err) + require.False(t, got.Terminal) + require.Equal(t, time.Minute, got.RequeueAfter) + }) + + t.Run("succeeded indexes and settles from submitted", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_flow"] = newJob(BatchImageJobStatusSubmitted) + provider := &fakeProcessorProvider{ + status: &BatchProviderStatus{InternalState: BatchProviderStateSucceeded, RawState: "SUCCEEDED", ProviderOutputRef: "files/output"}, + result: `{"key":"ok","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageTestData + `"}}]}}]}}` + "\n", + } + got, err := newTestBatchImageProcessor(repo, provider).Process(ctx, "imgbatch_flow") + require.NoError(t, err) + require.False(t, got.Terminal) + require.Equal(t, time.Millisecond, got.RequeueAfter) + require.Equal(t, BatchImageJobStatusSettling, repo.jobs["imgbatch_flow"].Status) + require.Equal(t, "files/output", batchImageDerefString(repo.jobs["imgbatch_flow"].ProviderOutputRef)) + require.Equal(t, []string{BatchImageJobStatusIndexing, BatchImageJobStatusSettling}, repo.transitions["imgbatch_flow"]) + require.Equal(t, BatchImageCounts{SuccessCount: 1}, repo.counts["imgbatch_flow"]) + }) + + t.Run("failed provider marks job failed", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_flow"] = newJob(BatchImageJobStatusRunning) + provider := &fakeProcessorProvider{status: &BatchProviderStatus{InternalState: BatchProviderStateFailed, RawState: "FAILED", ErrorCode: "BAD_PROMPT", ErrorMessage: "bad prompt"}} + got, err := newTestBatchImageProcessor(repo, provider).Process(ctx, "imgbatch_flow") + require.NoError(t, err) + require.True(t, got.Terminal) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs["imgbatch_flow"].Status) + require.Equal(t, "BAD_PROMPT", batchImageDerefString(repo.jobs["imgbatch_flow"].LastErrorCode)) + }) + + t.Run("cancelled provider marks job cancelled", func(t *testing.T) { + repo := newFakeBatchImageRepository() + repo.jobs["imgbatch_flow"] = newJob(BatchImageJobStatusRunning) + apiKeyID := int64(22) + holdAmount := 0.5 + repo.jobs["imgbatch_flow"].UserID = 11 + repo.jobs["imgbatch_flow"].APIKeyID = &apiKeyID + repo.jobs["imgbatch_flow"].EstimatedCost = holdAmount + repo.jobs["imgbatch_flow"].HoldAmount = &holdAmount + provider := &fakeProcessorProvider{status: &BatchProviderStatus{InternalState: BatchProviderStateCancelled, RawState: "CANCELLED"}} + processor := newTestBatchImageProcessor(repo, provider) + billing := &fakeBatchImageBillingRepo{} + processor.BillingRepo = billing + got, err := processor.Process(ctx, "imgbatch_flow") + require.NoError(t, err) + require.True(t, got.Terminal) + require.Equal(t, BatchImageJobStatusCancelled, repo.jobs["imgbatch_flow"].Status) + require.Len(t, billing.releases, 1) + require.Equal(t, BatchImageReleaseRequestID("imgbatch_flow"), billing.releases[0].RequestID) + }) +} + +func TestCanTransitionBatchImageJob_PR5DirectIndexing(t *testing.T) { + require.True(t, CanTransitionBatchImageJob(BatchImageJobStatusSubmitted, BatchImageJobStatusIndexing)) + require.True(t, CanTransitionBatchImageJob(BatchImageJobStatusSubmitted, BatchImageJobStatusFailed)) + require.True(t, CanTransitionBatchImageJob(BatchImageJobStatusIndexing, BatchImageJobStatusFailed)) +} + +func newTestBatchImageProcessor(repo *fakeBatchImageRepository, provider *fakeProcessorProvider) *BatchImageProviderProcessor { + return &BatchImageProviderProcessor{ + Repo: repo, + ProviderRegistry: NewBatchImageProviderRegistry(provider), + AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}, + Indexer: &BatchImageResultIndexer{Repo: repo}, + } +} + +type fakeBatchImageAccountResolver struct { + account *Account + err error +} + +func (r *fakeBatchImageAccountResolver) ResolveBatchImageAccount(context.Context, int64) (*Account, error) { + if r.err != nil { + return nil, r.err + } + return r.account, nil +} + +type fakeProcessorProvider struct { + status *BatchProviderStatus + getErr error + result string + + getCalled bool + openResultCalled bool +} + +func (p *fakeProcessorProvider) Name() string { return "fake" } +func (p *fakeProcessorProvider) SupportsAccount(*Account) bool { + return true +} +func (p *fakeProcessorProvider) Submit(context.Context, *BatchImageJob, *Account, BatchImageInput) (*BatchProviderJob, error) { + panic("Submit must not be called by PR5 processor") +} +func (p *fakeProcessorProvider) Get(context.Context, *BatchImageJob, *Account) (*BatchProviderStatus, error) { + p.getCalled = true + if p.getErr != nil { + return nil, p.getErr + } + if p.status == nil { + return &BatchProviderStatus{InternalState: BatchProviderStateQueued}, nil + } + return p.status, nil +} +func (p *fakeProcessorProvider) Cancel(context.Context, *BatchImageJob, *Account) error { return nil } +func (p *fakeProcessorProvider) OpenResult(context.Context, *BatchImageJob, *Account) (io.ReadCloser, string, error) { + p.openResultCalled = true + return io.NopCloser(strings.NewReader(p.result)), "application/jsonl", nil +} +func (p *fakeProcessorProvider) Cleanup(context.Context, *BatchImageJob, *Account, CleanupTarget) error { + return nil +} + +type fakeBatchImageRepository struct { + jobs map[string]*BatchImageJob + items map[string][]CreateBatchImageItemParams + counts map[string]BatchImageCounts + transitions map[string][]string + events map[string][]string + transitionErr error + replaceCalls int +} + +func newFakeBatchImageRepository() *fakeBatchImageRepository { + return &fakeBatchImageRepository{ + jobs: make(map[string]*BatchImageJob), + items: make(map[string][]CreateBatchImageItemParams), + counts: make(map[string]BatchImageCounts), + transitions: make(map[string][]string), + events: make(map[string][]string), + } +} + +func (r *fakeBatchImageRepository) CreateBatchImageJob(_ context.Context, params CreateBatchImageJobParams) (*BatchImageJob, error) { + job := &BatchImageJob{ + BatchID: params.BatchID, + UserID: params.UserID, + APIKeyID: params.APIKeyID, + AccountID: params.AccountID, + Status: params.Status, + Provider: params.Provider, + Model: params.Model, + TaskName: params.TaskName, + ProviderJobName: params.ProviderJobName, + ItemCount: params.ItemCount, + EstimatedCost: params.EstimatedCost, + HoldAmount: params.HoldAmount, + HoldID: params.HoldID, + BaseUnitPrice: params.BaseUnitPrice, + GroupRateMultiplier: params.GroupRateMultiplier, + AccountRateMultiplier: params.AccountRateMultiplier, + BatchDiscountMultiplier: params.BatchDiscountMultiplier, + HoldMultiplier: params.HoldMultiplier, + BillableUnitPrice: params.BillableUnitPrice, + HoldUnitPrice: params.HoldUnitPrice, + PricingSnapshotVersion: params.PricingSnapshotVersion, + Currency: params.Currency, + IdempotencyKey: params.IdempotencyKey, + RequestHash: params.RequestHash, + CreatedAt: time.Now(), + } + r.jobs[job.BatchID] = job + return job, nil +} + +func (r *fakeBatchImageRepository) GetBatchImageJobByBatchID(_ context.Context, batchID string) (*BatchImageJob, error) { + job, ok := r.jobs[batchID] + if !ok { + return nil, ErrBatchImageJobNotFound + } + return job, nil +} + +func (r *fakeBatchImageRepository) GetBatchImageJobByIdempotencyKey(_ context.Context, userID, apiKeyID int64, key string) (*BatchImageJob, error) { + for _, job := range r.jobs { + if job.UserID == userID && job.APIKeyID != nil && *job.APIKeyID == apiKeyID && batchImageDerefString(job.IdempotencyKey) == key { + return job, nil + } + } + return nil, ErrBatchImageJobNotFound +} + +func (r *fakeBatchImageRepository) GetBatchImageJobByBatchIDForOwner(_ context.Context, userID, apiKeyID int64, batchID string) (*BatchImageJob, error) { + job, ok := r.jobs[batchID] + if !ok || job.UserID != userID || job.APIKeyID == nil || *job.APIKeyID != apiKeyID { + return nil, ErrBatchImageJobNotFound + } + return job, nil +} + +func (r *fakeBatchImageRepository) ListBatchImageJobsForOwner(_ context.Context, userID, apiKeyID int64, filter BatchImageJobFilter) ([]*BatchImageJob, error) { + limit := filter.Limit + if limit <= 0 || limit > 100 { + limit = 20 + } + offset := filter.Offset + if offset < 0 { + offset = 0 + } + var jobs []*BatchImageJob + for _, job := range r.jobs { + if job.UserID != userID || job.APIKeyID == nil || *job.APIKeyID != apiKeyID { + continue + } + if filter.Status != "" && job.Status != filter.Status { + continue + } + if filter.TaskNameLike != "" && !strings.Contains(strings.ToLower(job.TaskName), strings.ToLower(filter.TaskNameLike)) { + continue + } + if filter.ExcludeDeleted && job.UserDeletedAt != nil { + continue + } + if filter.Downloaded != nil { + downloaded := job.DownloadedAt != nil + if downloaded != *filter.Downloaded { + continue + } + } + if filter.CreatedAfter != nil && job.CreatedAt.Before(*filter.CreatedAfter) { + continue + } + if filter.CreatedBefore != nil && !job.CreatedAt.Before(*filter.CreatedBefore) { + continue + } + if offset > 0 { + offset-- + continue + } + jobs = append(jobs, job) + if len(jobs) >= limit { + break + } + } + return jobs, nil +} + +func (r *fakeBatchImageRepository) GetBatchImageJobByID(_ context.Context, id int64) (*BatchImageJob, error) { + for _, job := range r.jobs { + if job.ID == id { + return job, nil + } + } + return nil, ErrBatchImageJobNotFound +} + +func (r *fakeBatchImageRepository) TransitionBatchImageJobStatus(_ context.Context, batchID, toStatus string, opts BatchImageTransitionOptions) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if !CanTransitionBatchImageJob(job.Status, toStatus) { + return ErrBatchImageInvalidTransition + } + if r.transitionErr != nil { + return r.transitionErr + } + job.Status = toStatus + job.LastErrorCode = opts.ErrorCode + job.LastErrorMessage = opts.ErrorMessage + r.transitions[batchID] = append(r.transitions[batchID], toStatus) + if opts.EventType != "" { + r.events[batchID] = append(r.events[batchID], opts.EventType) + } + return nil +} + +func (r *fakeBatchImageRepository) TouchBatchImageJobSubmitting(_ context.Context, batchID string) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if job.Status == BatchImageJobStatusCreated || job.Status == BatchImageJobStatusUploading { + job.UpdatedAt = time.Now() + } + return nil +} + +func (r *fakeBatchImageRepository) FailStaleUnsubmittedBatchImageJob(_ context.Context, batchID string, cutoff time.Time, code, message string) (bool, error) { + job, ok := r.jobs[batchID] + if !ok { + return false, ErrBatchImageJobNotFound + } + if job.Status != BatchImageJobStatusCreated && job.Status != BatchImageJobStatusUploading { + return false, nil + } + if batchImageDerefString(job.ProviderJobName) != "" || job.UpdatedAt.After(cutoff) { + return false, nil + } + job.Status = BatchImageJobStatusFailed + job.LastErrorCode = batchImageStringPtr(code) + job.LastErrorMessage = batchImageStringPtr(message) + job.UpdatedAt = time.Now() + r.transitions[batchID] = append(r.transitions[batchID], BatchImageJobStatusFailed) + r.events[batchID] = append(r.events[batchID], "billing_hold_recovery_failed_unsubmitted") + return true, nil +} + +func (r *fakeBatchImageRepository) UpdateBatchImageJobProviderOutputRef(_ context.Context, batchID, providerOutputRef string) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + job.ProviderOutputRef = &providerOutputRef + return nil +} + +func (r *fakeBatchImageRepository) UpdateBatchImageJobProviderSubmit(_ context.Context, params UpdateBatchImageJobProviderSubmitParams) error { + job, ok := r.jobs[params.BatchID] + if !ok { + return ErrBatchImageJobNotFound + } + if !CanTransitionBatchImageJob(job.Status, BatchImageJobStatusSubmitted) { + return ErrBatchImageInvalidTransition + } + job.Status = BatchImageJobStatusSubmitted + job.ProviderJobName = batchImageOptionalStringPtr(params.ProviderJobName) + job.ProviderInputRef = batchImageOptionalStringPtr(params.ProviderInputRef) + job.ProviderOutputRef = batchImageOptionalStringPtr(params.ProviderOutputRef) + job.GCSInputURI = batchImageOptionalStringPtr(params.GCSInputURI) + job.GCSOutputURI = batchImageOptionalStringPtr(params.GCSOutputURI) + now := time.Now() + job.SubmittedAt = &now + r.transitions[params.BatchID] = append(r.transitions[params.BatchID], BatchImageJobStatusSubmitted) + r.events[params.BatchID] = append(r.events[params.BatchID], "provider_submitted") + return nil +} + +func (r *fakeBatchImageRepository) RecordBatchImageJobSubmitFailure(_ context.Context, batchID, code, message string, markFailed bool) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if markFailed { + job.Status = BatchImageJobStatusFailed + } + job.LastErrorCode = batchImageOptionalStringPtr(code) + job.LastErrorMessage = batchImageOptionalStringPtr(message) + eventType := "submit_failed" + if !markFailed { + eventType = "queue_failed" + } + r.events[batchID] = append(r.events[batchID], eventType) + return nil +} + +func (r *fakeBatchImageRepository) MarkBatchImageJobSettled(_ context.Context, params MarkBatchImageJobSettledParams) error { + job, ok := r.jobs[params.BatchID] + if !ok { + return ErrBatchImageJobNotFound + } + if job.Status != BatchImageJobStatusSettling { + if job.Status == BatchImageJobStatusCompleted { + return ErrBatchImageAlreadySettled + } + return ErrBatchImageSettlementInvalidStatus + } + if batchImageDerefString(job.ManifestHash) != "" && batchImageDerefString(job.ManifestHash) != params.ManifestHash { + return ErrBatchImageSettlementManifestConflict + } + now := time.Now() + job.Status = BatchImageJobStatusCompleted + job.ActualCost = ¶ms.ActualCost + job.ManifestHash = ¶ms.ManifestHash + job.SettledAt = &now + if job.OutputExpiresAt == nil && params.OutputExpiresAt != nil { + job.OutputExpiresAt = params.OutputExpiresAt + } + r.transitions[params.BatchID] = append(r.transitions[params.BatchID], BatchImageJobStatusCompleted) + r.events[params.BatchID] = append(r.events[params.BatchID], "settlement_completed") + return nil +} + +func (r *fakeBatchImageRepository) SetBatchImageJobSettlementFailed(_ context.Context, batchID, code, message string) (int, error) { + job, ok := r.jobs[batchID] + if !ok { + return 0, ErrBatchImageJobNotFound + } + job.LastErrorCode = batchImageStringPtr(code) + job.LastErrorMessage = batchImageOptionalStringPtr(message) + job.RetryCount++ + r.events[batchID] = append(r.events[batchID], "settlement_failed") + return job.RetryCount, nil +} + +func (r *fakeBatchImageRepository) CreateBatchImageItem(_ context.Context, params CreateBatchImageItemParams) (*BatchImageItem, error) { + r.items[params.JobID] = append(r.items[params.JobID], params) + return &BatchImageItem{JobID: params.JobID, CustomID: params.CustomID, Status: params.Status}, nil +} + +func (r *fakeBatchImageRepository) BulkCreateBatchImageItems(ctx context.Context, params []CreateBatchImageItemParams) error { + for _, param := range params { + if _, err := r.CreateBatchImageItem(ctx, param); err != nil { + return err + } + } + return nil +} + +func (r *fakeBatchImageRepository) ReplaceBatchImageItemsForJob(_ context.Context, batchID string, items []CreateBatchImageItemParams, counts BatchImageCounts) error { + // 与真实实现一致:仅 indexing 状态允许重建 item 表(未注册的 job 保持宽松, + // 供直接构造 job 的单测使用)。 + if job, ok := r.jobs[batchID]; ok && job.Status != BatchImageJobStatusIndexing { + return ErrBatchImageIndexStateConflict + } + r.replaceCalls++ + copied := append([]CreateBatchImageItemParams(nil), items...) + for idx := range copied { + copied[idx].JobID = batchID + } + r.items[batchID] = copied + r.counts[batchID] = counts + if job, ok := r.jobs[batchID]; ok { + job.SuccessCount = counts.SuccessCount + job.FailCount = counts.FailCount + job.ItemCount = len(copied) + } + return nil +} + +func (r *fakeBatchImageRepository) ListBatchImageItems(_ context.Context, batchID string, filter BatchImageItemFilter) ([]*BatchImageItem, error) { + limit := filter.Limit + if limit <= 0 || limit > 500 { + limit = 100 + } + offset := filter.Offset + if offset < 0 { + offset = 0 + } + var result []*BatchImageItem + for _, item := range r.items[batchID] { + if filter.Status != "" && item.Status != filter.Status { + continue + } + if offset > 0 { + offset-- + continue + } + result = append(result, &BatchImageItem{ + JobID: item.JobID, + CustomID: item.CustomID, + Status: item.Status, + RequestHash: item.RequestHash, + PromptPreview: item.PromptPreview, + ProviderSourceObject: item.ProviderSourceObject, + SourceLineNumber: item.SourceLineNumber, + SourceByteOffset: item.SourceByteOffset, + SourceByteLength: item.SourceByteLength, + MimeType: item.MimeType, + FileExtension: item.FileExtension, + ImageCount: item.ImageCount, + ErrorCode: item.ErrorCode, + ErrorMessage: item.ErrorMessage, + BilledAmount: item.BilledAmount, + IndexedAt: item.IndexedAt, + }) + if len(result) >= limit { + break + } + } + return result, nil +} + +func (r *fakeBatchImageRepository) ListBatchImageItemsForOwner(ctx context.Context, userID, apiKeyID int64, batchID string, filter BatchImageItemFilter) ([]*BatchImageItem, error) { + if _, err := r.GetBatchImageJobByBatchIDForOwner(ctx, userID, apiKeyID, batchID); err != nil { + return nil, err + } + return r.ListBatchImageItems(ctx, batchID, filter) +} + +func (r *fakeBatchImageRepository) GetBatchImageJobForDownload(ctx context.Context, userID, apiKeyID int64, batchID string) (*BatchImageJob, error) { + return r.GetBatchImageJobByBatchIDForOwner(ctx, userID, apiKeyID, batchID) +} + +func (r *fakeBatchImageRepository) GetBatchImageItemForDownload(_ context.Context, batchID, customID string) (*BatchImageItem, error) { + for _, item := range r.items[batchID] { + if item.CustomID != customID { + continue + } + return &BatchImageItem{ + JobID: item.JobID, + CustomID: item.CustomID, + Status: item.Status, + RequestHash: item.RequestHash, + PromptPreview: item.PromptPreview, + ProviderSourceObject: item.ProviderSourceObject, + SourceLineNumber: item.SourceLineNumber, + SourceByteOffset: item.SourceByteOffset, + SourceByteLength: item.SourceByteLength, + MimeType: item.MimeType, + FileExtension: item.FileExtension, + ImageCount: item.ImageCount, + ErrorCode: item.ErrorCode, + ErrorMessage: item.ErrorMessage, + BilledAmount: item.BilledAmount, + IndexedAt: item.IndexedAt, + }, nil + } + return nil, ErrBatchImageItemNotFound +} + +func (r *fakeBatchImageRepository) ListBatchImageItemsForDownload(ctx context.Context, batchID string, status string, limit int) ([]*BatchImageItem, error) { + return r.ListBatchImageItems(ctx, batchID, BatchImageItemFilter{Status: status, Limit: limit}) +} + +func (r *fakeBatchImageRepository) ListBatchImageJobsDueForInputCleanup(_ context.Context, cutoff time.Time, limit int) ([]*BatchImageJob, error) { + if limit <= 0 { + limit = 100 + } + var jobs []*BatchImageJob + for _, job := range r.jobs { + if job.InputDeletedAt != nil || batchImageDerefString(job.ProviderInputRef) == "" || !IsTerminalBatchImageJobStatus(job.Status) { + continue + } + at := job.FinishedAt + if at == nil { + at = job.SettledAt + } + if at == nil { + at = &job.UpdatedAt + } + if at != nil && at.After(cutoff) { + continue + } + jobs = append(jobs, job) + if len(jobs) >= limit { + break + } + } + return jobs, nil +} + +func (r *fakeBatchImageRepository) ListBatchImageJobsDueForOutputCleanup(_ context.Context, now time.Time, limit int) ([]*BatchImageJob, error) { + if limit <= 0 { + limit = 100 + } + var jobs []*BatchImageJob + for _, job := range r.jobs { + if job.OutputDeletedAt != nil || batchImageDerefString(job.ProviderOutputRef) == "" || job.Status != BatchImageJobStatusCompleted || job.OutputExpiresAt == nil || job.OutputExpiresAt.After(now) { + continue + } + jobs = append(jobs, job) + if len(jobs) >= limit { + break + } + } + return jobs, nil +} + +func (r *fakeBatchImageRepository) ListStaleUnsubmittedBatchImageJobs(_ context.Context, cutoff time.Time, limit int) ([]*BatchImageJob, error) { + if limit <= 0 { + limit = 100 + } + jobs := make([]*BatchImageJob, 0, limit) + for _, job := range r.jobs { + if len(jobs) >= limit { + break + } + if job.Status != BatchImageJobStatusCreated && job.Status != BatchImageJobStatusUploading { + continue + } + if batchImageDerefString(job.ProviderJobName) != "" { + continue + } + holdAmount := job.EstimatedCost + if job.HoldAmount != nil { + holdAmount = *job.HoldAmount + } + if holdAmount <= 0 || job.UpdatedAt.After(cutoff) { + continue + } + jobs = append(jobs, job) + } + return jobs, nil +} + +func (r *fakeBatchImageRepository) MarkBatchImageInputDeleted(_ context.Context, batchID string, deletedAt time.Time) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if job.InputDeletedAt == nil { + job.InputDeletedAt = &deletedAt + } + r.events[batchID] = append(r.events[batchID], "input_cleanup_completed") + return nil +} + +func (r *fakeBatchImageRepository) MarkBatchImageOutputDeleted(_ context.Context, batchID string, deletedAt time.Time) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if job.OutputDeletedAt == nil { + job.OutputDeletedAt = &deletedAt + } + if job.Status == BatchImageJobStatusCompleted { + job.Status = BatchImageJobStatusOutputDeleted + } + r.events[batchID] = append(r.events[batchID], "output_cleanup_completed") + return nil +} + +func (r *fakeBatchImageRepository) MarkBatchImageDownloaded(_ context.Context, batchID string, downloadedAt time.Time) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if job.DownloadedAt == nil { + job.DownloadedAt = &downloadedAt + } + r.events[batchID] = append(r.events[batchID], "download_completed") + return nil +} + +func (r *fakeBatchImageRepository) MarkBatchImageJobUserDeleted(_ context.Context, userID, apiKeyID int64, batchID string, deletedAt time.Time) error { + job, ok := r.jobs[batchID] + if !ok || job.UserID != userID || job.APIKeyID == nil || *job.APIKeyID != apiKeyID { + return ErrBatchImageJobNotFound + } + if !isBatchImageProcessorDoneStatus(job.Status) { + return ErrBatchImageRecordDeleteNotReady + } + if job.UserDeletedAt == nil { + job.UserDeletedAt = &deletedAt + } + r.events[batchID] = append(r.events[batchID], "user_record_deleted") + return nil +} + +func (r *fakeBatchImageRepository) SetBatchImageOutputExpiresAt(_ context.Context, batchID string, expiresAt time.Time) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + if job.OutputExpiresAt == nil { + job.OutputExpiresAt = &expiresAt + } + return nil +} + +func (r *fakeBatchImageRepository) RecordBatchImageCleanupFailure(_ context.Context, batchID, code, message string) error { + job, ok := r.jobs[batchID] + if !ok { + return ErrBatchImageJobNotFound + } + job.LastErrorCode = batchImageStringPtr(code) + job.LastErrorMessage = batchImageOptionalStringPtr(message) + r.events[batchID] = append(r.events[batchID], "output_cleanup_failed") + return nil +} + +func (r *fakeBatchImageRepository) AppendBatchImageEvent(_ context.Context, batchID, eventType string, _ any) error { + r.events[batchID] = append(r.events[batchID], eventType) + return nil +} + +var _ BatchImageRepository = (*fakeBatchImageRepository)(nil) +var _ BatchImageProvider = (*fakeProcessorProvider)(nil) +var _ BatchImageAccountResolver = (*fakeBatchImageAccountResolver)(nil) +var _ = infraerrors.Reason diff --git a/backend/internal/service/batch_image_provider.go b/backend/internal/service/batch_image_provider.go new file mode 100644 index 0000000000..9f4bcca6e4 --- /dev/null +++ b/backend/internal/service/batch_image_provider.go @@ -0,0 +1,180 @@ +package service + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +type BatchImageProvider interface { + Name() string + SupportsAccount(account *Account) bool + Submit(ctx context.Context, job *BatchImageJob, account *Account, input BatchImageInput) (*BatchProviderJob, error) + Get(ctx context.Context, job *BatchImageJob, account *Account) (*BatchProviderStatus, error) + Cancel(ctx context.Context, job *BatchImageJob, account *Account) error + OpenResult(ctx context.Context, job *BatchImageJob, account *Account) (io.ReadCloser, string, error) + Cleanup(ctx context.Context, job *BatchImageJob, account *Account, target CleanupTarget) error +} + +type BatchImageProviderRegistry struct { + providers map[string]BatchImageProvider +} + +func NewBatchImageProviderRegistry(providers ...BatchImageProvider) *BatchImageProviderRegistry { + r := &BatchImageProviderRegistry{providers: make(map[string]BatchImageProvider, len(providers))} + for _, provider := range providers { + if provider == nil || strings.TrimSpace(provider.Name()) == "" { + continue + } + r.providers[provider.Name()] = provider + } + return r +} + +func NewDefaultBatchImageProviderRegistry() *BatchImageProviderRegistry { + return NewBatchImageProviderRegistry( + NewGeminiAPIBatchImageProvider(nil), + NewVertexBatchImageProvider(VertexBatchImageProviderOptions{}, nil, nil, nil), + ) +} + +func NewBatchImageProviderRegistryFromConfig(cfg *config.Config) *BatchImageProviderRegistry { + return NewBatchImageProviderRegistry( + NewGeminiAPIBatchImageProvider(nil), + NewVertexBatchImageProviderFromConfig(cfg, nil, nil, nil), + ) +} + +func (r *BatchImageProviderRegistry) Get(provider string) (BatchImageProvider, bool) { + if r == nil { + return nil, false + } + p, ok := r.providers[provider] + return p, ok +} + +func (r *BatchImageProviderRegistry) MustGet(provider string) (BatchImageProvider, error) { + p, ok := r.Get(provider) + if !ok { + return nil, ErrBatchImageInvalidProvider + } + return p, nil +} + +type BatchImageInput struct { + BatchID string + Model string + DisplayName string + Items []BatchImageInputItem + + ResponseMimeType string + AspectRatio string + ImageSize string + + Metadata map[string]string +} + +type BatchImageInputItem struct { + CustomID string + Prompt string + + ReferenceImages []BatchImageReference +} + +type BatchImageReference struct { + ID string + Type string + MimeType string + Data []byte + FileURI string +} + +type BatchProviderJob struct { + ProviderJobName string + ProviderInputRef string + ProviderOutputRef string + RawState string +} + +type BatchProviderInternalState string + +const ( + BatchProviderStateQueued BatchProviderInternalState = "queued" + BatchProviderStateRunning BatchProviderInternalState = "running" + BatchProviderStateSucceeded BatchProviderInternalState = "succeeded" + BatchProviderStateFailed BatchProviderInternalState = "failed" + BatchProviderStateCancelled BatchProviderInternalState = "cancelled" + BatchProviderStateExpired BatchProviderInternalState = "expired" +) + +type BatchProviderStatus struct { + RawState string + + InternalState BatchProviderInternalState + Done bool + + ProviderOutputRef string + + ErrorCode string + ErrorMessage string + + SuggestedRequeueAfter time.Duration +} + +type CleanupTarget string + +const ( + CleanupTargetInput CleanupTarget = "input" + CleanupTargetOutput CleanupTarget = "output" + CleanupTargetAll CleanupTarget = "all" +) + +var ( + ErrBatchImageProviderUnsupportedAccount = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_UNSUPPORTED_ACCOUNT", "batch image provider does not support this account") + ErrBatchImageProviderMissingAPIKey = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_MISSING_API_KEY", "batch image provider account is missing api key") + ErrBatchImageProviderMissingServiceAccount = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_MISSING_SERVICE_ACCOUNT", "batch image provider account is missing service account credentials") + ErrBatchImageProviderMissingJobName = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_MISSING_JOB_NAME", "batch image provider job name is missing") + ErrBatchImageProviderMissingResultRef = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_MISSING_RESULT_REF", "batch image provider result reference is missing") + ErrBatchImageProviderInlineResultUnsupported = infraerrors.New(http.StatusBadRequest, "GEMINI_INLINE_BATCH_RESULT_UNSUPPORTED", "Gemini inline batch result is not supported") + ErrBatchImageProviderInvalidInput = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_INVALID_INPUT", "invalid batch image provider input") + ErrBatchImageProviderUnsafeCleanupPath = infraerrors.New(http.StatusBadRequest, "VERTEX_UNSAFE_CLEANUP_PATH", "unsafe batch image cleanup path") + ErrUnsupportedCleanupTarget = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_PROVIDER_UNSUPPORTED_CLEANUP_TARGET", "unsupported batch image cleanup target") +) + +func batchImageProviderJobName(job *BatchImageJob) string { + if job == nil || job.ProviderJobName == nil { + return "" + } + return strings.TrimSpace(*job.ProviderJobName) +} + +func batchImageProviderInputRef(job *BatchImageJob) string { + if job == nil || job.ProviderInputRef == nil { + return "" + } + return strings.TrimSpace(*job.ProviderInputRef) +} + +func batchImageProviderOutputRef(job *BatchImageJob) string { + if job == nil || job.ProviderOutputRef == nil { + return "" + } + return strings.TrimSpace(*job.ProviderOutputRef) +} + +func batchImageProviderAPIKey(account *Account) string { + if account == nil { + return "" + } + return strings.TrimSpace(account.GetCredential("api_key")) +} + +func batchImageProviderInputError(format string, args ...any) error { + return ErrBatchImageProviderInvalidInput.WithCause(fmt.Errorf(format, args...)) +} diff --git a/backend/internal/service/batch_image_provider_gemini.go b/backend/internal/service/batch_image_provider_gemini.go new file mode 100644 index 0000000000..64efb0b35d --- /dev/null +++ b/backend/internal/service/batch_image_provider_gemini.go @@ -0,0 +1,719 @@ +package service + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "mime/multipart" + "net/http" + "net/textproto" + "net/url" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient" +) + +const defaultGeminiBatchRequeueAfter = 30 * time.Second + +type GeminiBatchClient interface { + UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error) + CreateBatch(ctx context.Context, apiKey string, model string, fileName string, displayName string) (*GeminiBatchJob, error) + GetBatch(ctx context.Context, apiKey string, batchName string) (*GeminiBatchJob, error) + CancelBatch(ctx context.Context, apiKey string, batchName string) error + DownloadFile(ctx context.Context, apiKey string, fileName string) (io.ReadCloser, string, error) + DeleteFile(ctx context.Context, apiKey string, fileName string) error +} + +type GeminiUploadedFile struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + URI string `json:"uri"` + MimeType string `json:"mimeType"` +} + +type GeminiBatchJob struct { + Name string `json:"name"` + State string `json:"state"` + Dest *GeminiBatchDest `json:"dest"` + Response *GeminiBatchResponse `json:"response"` + Error *GeminiBatchError `json:"error"` + Raw map[string]any `json:"-"` +} + +type GeminiBatchDest struct { + FileName string `json:"fileName"` + FileNameSnake string `json:"file_name"` +} + +type GeminiBatchResponse struct { + ResponsesFile string `json:"responsesFile"` + ResponsesFileSnake string `json:"responses_file"` + InlinedResponses []any `json:"inlinedResponses"` + InlinedResponsesAlt []any `json:"inlined_responses"` +} + +type GeminiBatchError struct { + Code string `json:"code"` + Message string `json:"message"` + Status string `json:"status"` +} + +type GeminiAPIBatchImageProvider struct { + client GeminiBatchClient +} + +func NewGeminiAPIBatchImageProvider(client GeminiBatchClient) *GeminiAPIBatchImageProvider { + if client == nil { + client = NewGeminiBatchHTTPClient("", nil) + } + return &GeminiAPIBatchImageProvider{client: client} +} + +func (p *GeminiAPIBatchImageProvider) Name() string { + return BatchImageProviderGeminiAPI +} + +func (p *GeminiAPIBatchImageProvider) SupportsAccount(account *Account) bool { + return account != nil && + account.Platform == PlatformGemini && + account.Type == AccountTypeAPIKey && + batchImageProviderAPIKey(account) != "" +} + +func (p *GeminiAPIBatchImageProvider) Submit(ctx context.Context, job *BatchImageJob, account *Account, input BatchImageInput) (*BatchProviderJob, error) { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey { + return nil, ErrBatchImageProviderUnsupportedAccount + } + apiKey := batchImageProviderAPIKey(account) + if apiKey == "" { + return nil, ErrBatchImageProviderMissingAPIKey + } + if input.BatchID == "" && job != nil { + input.BatchID = job.BatchID + } + if input.Model == "" && job != nil { + input.Model = job.Model + } + + jsonl, err := BuildGeminiBatchJSONL(input) + if err != nil { + return nil, err + } + + displayName := strings.TrimSpace(input.DisplayName) + if displayName == "" { + displayName = strings.TrimSpace(input.BatchID) + } + + uploaded, err := p.client.UploadJSONL(ctx, apiKey, displayName, bytes.NewReader(jsonl)) + if err != nil { + return nil, mapGeminiClientError(err) + } + if uploaded == nil || strings.TrimSpace(uploaded.Name) == "" { + return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini upload response is missing file name", nil) + } + + batch, err := p.client.CreateBatch(ctx, apiKey, input.Model, uploaded.Name, displayName) + if err != nil { + return nil, mapGeminiClientError(err) + } + if batch == nil || strings.TrimSpace(batch.Name) == "" { + return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini batch response is missing job name", nil) + } + + return &BatchProviderJob{ + ProviderJobName: batch.Name, + ProviderInputRef: uploaded.Name, + RawState: batch.State, + }, nil +} + +func (p *GeminiAPIBatchImageProvider) Get(ctx context.Context, job *BatchImageJob, account *Account) (*BatchProviderStatus, error) { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey { + return nil, ErrBatchImageProviderUnsupportedAccount + } + apiKey := batchImageProviderAPIKey(account) + if apiKey == "" { + return nil, ErrBatchImageProviderMissingAPIKey + } + jobName := batchImageProviderJobName(job) + if jobName == "" { + return nil, ErrBatchImageProviderMissingJobName + } + + batch, err := p.client.GetBatch(ctx, apiKey, jobName) + if err != nil { + return nil, mapGeminiClientError(err) + } + if batch == nil { + return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini batch response is empty", nil) + } + + status := mapGeminiBatchState(batch) + if status.InternalState == BatchProviderStateSucceeded { + if geminiBatchHasInlineResults(batch) { + return nil, ErrBatchImageProviderInlineResultUnsupported + } + outputRef := geminiBatchOutputRef(batch) + if outputRef == "" { + status.InternalState = BatchProviderStateFailed + status.Done = true + status.ErrorCode = "GEMINI_RESULT_FILE_MISSING" + status.ErrorMessage = "Gemini batch succeeded without a result file reference" + } + status.ProviderOutputRef = outputRef + } + return status, nil +} + +func (p *GeminiAPIBatchImageProvider) Cancel(ctx context.Context, job *BatchImageJob, account *Account) error { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey { + return ErrBatchImageProviderUnsupportedAccount + } + apiKey := batchImageProviderAPIKey(account) + if apiKey == "" { + return ErrBatchImageProviderMissingAPIKey + } + jobName := batchImageProviderJobName(job) + if jobName == "" { + return ErrBatchImageProviderMissingJobName + } + return mapGeminiClientError(p.client.CancelBatch(ctx, apiKey, jobName)) +} + +func (p *GeminiAPIBatchImageProvider) OpenResult(ctx context.Context, job *BatchImageJob, account *Account) (io.ReadCloser, string, error) { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey { + return nil, "", ErrBatchImageProviderUnsupportedAccount + } + apiKey := batchImageProviderAPIKey(account) + if apiKey == "" { + return nil, "", ErrBatchImageProviderMissingAPIKey + } + outputRef := batchImageProviderOutputRef(job) + if outputRef == "" { + return nil, "", ErrBatchImageProviderMissingResultRef + } + r, contentType, err := p.client.DownloadFile(ctx, apiKey, outputRef) + return r, contentType, mapGeminiClientError(err) +} + +func (p *GeminiAPIBatchImageProvider) Cleanup(ctx context.Context, job *BatchImageJob, account *Account, target CleanupTarget) error { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey { + return ErrBatchImageProviderUnsupportedAccount + } + apiKey := batchImageProviderAPIKey(account) + if apiKey == "" { + return ErrBatchImageProviderMissingAPIKey + } + + switch target { + case CleanupTargetInput: + return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderInputRef(job)) + case CleanupTargetOutput: + return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderOutputRef(job)) + case CleanupTargetAll: + if err := p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderInputRef(job)); err != nil { + return err + } + return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderOutputRef(job)) + default: + return ErrUnsupportedCleanupTarget + } +} + +func (p *GeminiAPIBatchImageProvider) deleteGeminiFileIfPresent(ctx context.Context, apiKey, fileName string) error { + if strings.TrimSpace(fileName) == "" { + return nil + } + return mapGeminiClientError(p.client.DeleteFile(ctx, apiKey, fileName)) +} + +type geminiJSONLLine struct { + Key string `json:"key"` + Request geminiGenerateRequest `json:"request"` +} + +type geminiGenerateRequest struct { + Contents []geminiContent `json:"contents"` + GenerationConfig geminiGenerationConfig `json:"generationConfig"` +} + +type geminiContent struct { + Parts []geminiPart `json:"parts"` +} + +type geminiPart struct { + Text string `json:"text,omitempty"` + InlineData *geminiInlineData `json:"inlineData,omitempty"` + FileData *geminiFileData `json:"fileData,omitempty"` +} + +type geminiInlineData struct { + MimeType string `json:"mimeType"` + Data string `json:"data"` +} + +type geminiFileData struct { + MimeType string `json:"mimeType"` + FileURI string `json:"fileUri"` +} + +type geminiGenerationConfig struct { + ResponseModalities []string `json:"responseModalities"` +} + +func BuildGeminiBatchJSONL(input BatchImageInput) ([]byte, error) { + if strings.TrimSpace(input.Model) == "" { + return nil, batchImageProviderInputError("model is required") + } + if len(input.Items) == 0 { + return nil, batchImageProviderInputError("at least one item is required") + } + + seen := make(map[string]struct{}, len(input.Items)) + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + for _, item := range input.Items { + customID := strings.TrimSpace(item.CustomID) + if customID == "" { + return nil, batchImageProviderInputError("custom_id is required") + } + if _, ok := seen[customID]; ok { + return nil, batchImageProviderInputError("duplicate custom_id %q", customID) + } + seen[customID] = struct{}{} + + prompt := strings.TrimSpace(item.Prompt) + if prompt == "" { + return nil, batchImageProviderInputError("prompt is required for custom_id %q", customID) + } + parts, err := batchImageGeminiParts(prompt, item.ReferenceImages) + if err != nil { + return nil, err + } + + // TODO(batch-image): add response_mime_type/aspect_ratio/image_size once the + // Gemini batch image REST shape is stabilized for those options. + line := geminiJSONLLine{ + Key: customID, + Request: geminiGenerateRequest{ + Contents: []geminiContent{{ + Parts: parts, + }}, + GenerationConfig: geminiGenerationConfig{ + ResponseModalities: []string{"TEXT", "IMAGE"}, + }, + }, + } + if err := enc.Encode(line); err != nil { + return nil, err + } + } + return buf.Bytes(), nil +} + +func batchImageGeminiParts(prompt string, refs []BatchImageReference) ([]geminiPart, error) { + parts := []geminiPart{{Text: prompt}} + for _, ref := range refs { + mimeType := normalizeBatchImageReferenceMimeType(ref.MimeType) + if mimeType == "" { + return nil, batchImageProviderInputError("reference image mime_type is required") + } + fileURI := strings.TrimSpace(ref.FileURI) + switch { + case len(ref.Data) > 0 && fileURI == "": + parts = append(parts, geminiPart{InlineData: &geminiInlineData{ + MimeType: mimeType, + Data: base64.StdEncoding.EncodeToString(ref.Data), + }}) + case len(ref.Data) == 0 && fileURI != "": + parts = append(parts, geminiPart{FileData: &geminiFileData{ + MimeType: mimeType, + FileURI: fileURI, + }}) + default: + return nil, batchImageProviderInputError("reference image must contain exactly one of data or file_uri") + } + } + return parts, nil +} + +func mapGeminiBatchState(batch *GeminiBatchJob) *BatchProviderStatus { + state := strings.TrimSpace(batch.State) + normalized := strings.ToUpper(state) + status := &BatchProviderStatus{ + RawState: state, + InternalState: BatchProviderStateRunning, + SuggestedRequeueAfter: defaultGeminiBatchRequeueAfter, + } + + switch normalized { + case "JOB_STATE_PENDING", "JOB_STATE_QUEUED": + status.InternalState = BatchProviderStateQueued + case "JOB_STATE_RUNNING": + status.InternalState = BatchProviderStateRunning + case "JOB_STATE_SUCCEEDED": + status.InternalState = BatchProviderStateSucceeded + status.Done = true + case "JOB_STATE_FAILED": + status.InternalState = BatchProviderStateFailed + status.Done = true + status.ErrorCode = "GEMINI_BATCH_FAILED" + case "JOB_STATE_CANCELLED": + status.InternalState = BatchProviderStateCancelled + status.Done = true + status.ErrorCode = "GEMINI_BATCH_CANCELLED" + case "JOB_STATE_EXPIRED": + status.InternalState = BatchProviderStateExpired + status.Done = true + status.ErrorCode = "GEMINI_BATCH_EXPIRED" + default: + if batch.Error != nil && (strings.TrimSpace(batch.Error.Message) != "" || strings.TrimSpace(batch.Error.Code) != "") { + status.InternalState = BatchProviderStateFailed + status.Done = true + status.ErrorCode = "GEMINI_BATCH_FAILED" + } + } + + if batch.Error != nil { + if code := strings.TrimSpace(batch.Error.Code); code != "" { + status.ErrorCode = code + } else if status.ErrorCode == "" && strings.TrimSpace(batch.Error.Status) != "" { + status.ErrorCode = strings.TrimSpace(batch.Error.Status) + } + status.ErrorMessage = strings.TrimSpace(batch.Error.Message) + } + return status +} + +func geminiBatchOutputRef(batch *GeminiBatchJob) string { + if batch == nil { + return "" + } + if batch.Dest != nil { + if v := strings.TrimSpace(batch.Dest.FileName); v != "" { + return v + } + if v := strings.TrimSpace(batch.Dest.FileNameSnake); v != "" { + return v + } + } + if batch.Response != nil { + if v := strings.TrimSpace(batch.Response.ResponsesFile); v != "" { + return v + } + if v := strings.TrimSpace(batch.Response.ResponsesFileSnake); v != "" { + return v + } + } + return "" +} + +func geminiBatchHasInlineResults(batch *GeminiBatchJob) bool { + return batch != nil && + batch.Response != nil && + (len(batch.Response.InlinedResponses) > 0 || len(batch.Response.InlinedResponsesAlt) > 0) +} + +func geminiProviderError(reason, message string, cause error) error { + err := infraerrors.New(http.StatusBadGateway, reason, message) + if cause != nil { + return err.WithCause(cause) + } + return err +} + +func mapGeminiClientError(err error) error { + if err == nil { + return nil + } + var apiErr *GeminiAPIError + if errors.As(err, &apiErr) { + switch apiErr.StatusCode { + case http.StatusUnauthorized, http.StatusForbidden: + return geminiProviderError("GEMINI_AUTH_FAILED", "Gemini authentication failed", nil) + case http.StatusTooManyRequests: + return geminiProviderError("GEMINI_RATE_LIMITED", "Gemini rate limit exceeded", nil) + case http.StatusNotFound: + return geminiProviderError("GEMINI_BATCH_NOT_FOUND", "Gemini batch resource was not found", nil) + default: + return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini API request failed", nil) + } + } + return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini API request failed", nil) +} + +type GeminiBatchHTTPClient struct { + baseURL string + client *http.Client +} + +func NewGeminiBatchHTTPClient(baseURL string, client *http.Client) *GeminiBatchHTTPClient { + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" { + baseURL = geminicli.AIStudioBaseURL + } + if client == nil { + client = batchImageDefaultHTTPClient() + } + return &GeminiBatchHTTPClient{baseURL: baseURL, client: client} +} + +// batchImageDefaultHTTPClient 返回带连接/握手/响应头超时的共享客户端。 +// 不设整体 Timeout:大文件上传与结果流式下载耗时不可预估, +// 但拨号、TLS、等待响应头必须有界,否则挂死的连接会无限占用提交路径。 +func batchImageDefaultHTTPClient() *http.Client { + client, err := httpclient.GetClient(httpclient.Options{ + ResponseHeaderTimeout: 60 * time.Second, + }) + if err != nil { + return http.DefaultClient + } + return client +} + +func (c *GeminiBatchHTTPClient) UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + metadataHeader := textproto.MIMEHeader{} + metadataHeader.Set("Content-Disposition", `form-data; name="metadata"`) + metadataHeader.Set("Content-Type", "application/json; charset=utf-8") + metadataPart, err := writer.CreatePart(metadataHeader) + if err != nil { + return nil, err + } + metadata := map[string]any{"file": map[string]any{"displayName": displayName, "mimeType": "application/jsonl"}} + if err := json.NewEncoder(metadataPart).Encode(metadata); err != nil { + return nil, err + } + fileHeader := textproto.MIMEHeader{} + fileHeader.Set("Content-Disposition", `form-data; name="file"; filename="batch.jsonl"`) + fileHeader.Set("Content-Type", "application/jsonl") + filePart, err := writer.CreatePart(fileHeader) + if err != nil { + return nil, err + } + if _, err := io.Copy(filePart, r); err != nil { + return nil, err + } + if err := writer.Close(); err != nil { + return nil, err + } + + req, err := c.newRequest(ctx, http.MethodPost, "/upload/v1beta/files?uploadType=multipart", apiKey, &body) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", writer.FormDataContentType()) + + var resp struct { + File *GeminiUploadedFile `json:"file"` + *GeminiUploadedFile + } + if err := c.doJSON(req, &resp); err != nil { + return nil, err + } + if resp.File != nil { + return resp.File, nil + } + return resp.GeminiUploadedFile, nil +} + +func (c *GeminiBatchHTTPClient) CreateBatch(ctx context.Context, apiKey string, model string, fileName string, displayName string) (*GeminiBatchJob, error) { + body := map[string]any{ + "batch": map[string]any{ + "displayName": displayName, + "inputConfig": map[string]any{ + "fileName": fileName, + }, + }, + } + payload, _ := json.Marshal(body) + path := fmt.Sprintf("/v1beta/models/%s:batchGenerateContent", url.PathEscape(strings.TrimSpace(model))) + req, err := c.newRequest(ctx, http.MethodPost, path, apiKey, bytes.NewReader(payload)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + return c.doBatchJob(req) +} + +func (c *GeminiBatchHTTPClient) GetBatch(ctx context.Context, apiKey string, batchName string) (*GeminiBatchJob, error) { + req, err := c.newRequest(ctx, http.MethodGet, "/v1beta/"+strings.TrimLeft(batchName, "/"), apiKey, nil) + if err != nil { + return nil, err + } + return c.doBatchJob(req) +} + +func (c *GeminiBatchHTTPClient) CancelBatch(ctx context.Context, apiKey string, batchName string) error { + req, err := c.newRequest(ctx, http.MethodPost, "/v1beta/"+strings.TrimLeft(batchName, "/")+":cancel", apiKey, nil) + if err != nil { + return err + } + return c.doNoBody(req) +} + +func (c *GeminiBatchHTTPClient) DownloadFile(ctx context.Context, apiKey string, fileName string) (io.ReadCloser, string, error) { + metaReq, err := c.newRequest(ctx, http.MethodGet, "/v1beta/"+strings.TrimLeft(fileName, "/"), apiKey, nil) + if err != nil { + return nil, "", err + } + var metadata struct { + DownloadURI string `json:"downloadUri"` + DownloadURL string `json:"download_url"` + MimeType string `json:"mimeType"` + } + if err := c.doJSON(metaReq, &metadata); err != nil { + return nil, "", err + } + downloadURL := strings.TrimSpace(metadata.DownloadURI) + if downloadURL == "" { + downloadURL = strings.TrimSpace(metadata.DownloadURL) + } + if downloadURL == "" { + downloadURL = c.baseURL + "/v1beta/" + strings.TrimLeft(fileName, "/") + ":download" + } + // 纵深加固:downloadUri 来自上游响应,跟随前校验目标 host, + // 防止异常/被劫持的响应把带 api key 的请求带到任意主机。 + if err := validateGeminiDownloadHost(downloadURL, c.baseURL); err != nil { + return nil, "", err + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil) + if err != nil { + return nil, "", err + } + req.Header.Set("x-goog-api-key", apiKey) + resp, err := c.client.Do(req) + if err != nil { + return nil, "", err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + defer func() { _ = resp.Body.Close() }() + return nil, "", readGeminiAPIError(resp) + } + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + contentType = metadata.MimeType + } + if contentType == "" { + contentType = "application/octet-stream" + } + return resp.Body, contentType, nil +} + +func (c *GeminiBatchHTTPClient) DeleteFile(ctx context.Context, apiKey string, fileName string) error { + req, err := c.newRequest(ctx, http.MethodDelete, "/v1beta/"+strings.TrimLeft(fileName, "/"), apiKey, nil) + if err != nil { + return err + } + return c.doNoBody(req) +} + +func (c *GeminiBatchHTTPClient) doBatchJob(req *http.Request) (*GeminiBatchJob, error) { + var job GeminiBatchJob + if err := c.doJSON(req, &job); err != nil { + return nil, err + } + job.Raw = map[string]any{} + return &job, nil +} + +func (c *GeminiBatchHTTPClient) doNoBody(req *http.Request) error { + resp, err := c.client.Do(req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return readGeminiAPIError(resp) + } + return nil +} + +func (c *GeminiBatchHTTPClient) doJSON(req *http.Request, out any) error { + resp, err := c.client.Do(req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return readGeminiAPIError(resp) + } + return json.NewDecoder(resp.Body).Decode(out) +} + +func (c *GeminiBatchHTTPClient) newRequest(ctx context.Context, method, path, apiKey string, body io.Reader) (*http.Request, error) { + if strings.TrimSpace(apiKey) == "" { + return nil, ErrBatchImageProviderMissingAPIKey + } + req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, body) + if err != nil { + return nil, err + } + req.Header.Set("x-goog-api-key", apiKey) + return req, nil +} + +// validateGeminiDownloadHost 只允许跟随到 googleapis.com(含子域) +// 或与配置的 baseURL 同 host 的下载地址。 +func validateGeminiDownloadHost(downloadURL, baseURL string) error { + parsed, err := url.Parse(downloadURL) + if err != nil { + return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri is invalid", err) + } + if parsed.Scheme != "https" { + return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri must use https", nil) + } + host := strings.ToLower(parsed.Hostname()) + if host == "googleapis.com" || strings.HasSuffix(host, ".googleapis.com") { + return nil + } + if base, err := url.Parse(baseURL); err == nil && strings.EqualFold(base.Hostname(), host) { + return nil + } + return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri host is not allowed", nil) +} + +type GeminiAPIError struct { + StatusCode int + Code string + Message string +} + +func (e *GeminiAPIError) Error() string { + if e == nil { + return "" + } + if e.Code != "" { + return fmt.Sprintf("gemini api error: status=%d code=%s message=%s", e.StatusCode, e.Code, e.Message) + } + return fmt.Sprintf("gemini api error: status=%d message=%s", e.StatusCode, e.Message) +} + +func readGeminiAPIError(resp *http.Response) error { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 8192)) + message := string(body) + var parsed struct { + Error struct { + Code any `json:"code"` + Message string `json:"message"` + Status string `json:"status"` + } `json:"error"` + } + if err := json.Unmarshal(body, &parsed); err == nil && parsed.Error.Message != "" { + message = parsed.Error.Message + return &GeminiAPIError{StatusCode: resp.StatusCode, Code: parsed.Error.Status, Message: message} + } + return &GeminiAPIError{StatusCode: resp.StatusCode, Message: message} +} + +var _ BatchImageProvider = (*GeminiAPIBatchImageProvider)(nil) +var _ GeminiBatchClient = (*GeminiBatchHTTPClient)(nil) diff --git a/backend/internal/service/batch_image_provider_gemini_test.go b/backend/internal/service/batch_image_provider_gemini_test.go new file mode 100644 index 0000000000..101c188b2d --- /dev/null +++ b/backend/internal/service/batch_image_provider_gemini_test.go @@ -0,0 +1,360 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/stretchr/testify/require" +) + +func TestBatchImageProviderRegistry_ReturnsGeminiAPI(t *testing.T) { + registry := NewDefaultBatchImageProviderRegistry() + provider, ok := registry.Get(BatchImageProviderGeminiAPI) + require.True(t, ok) + require.Equal(t, BatchImageProviderGeminiAPI, provider.Name()) + + must, err := registry.MustGet(BatchImageProviderGeminiAPI) + require.NoError(t, err) + require.Same(t, provider, must) + + _, err = registry.MustGet("unknown_provider") + require.ErrorIs(t, err, ErrBatchImageInvalidProvider) +} + +func TestGeminiProvider_SupportsOnlyGeminiAPIKeyWithSecret(t *testing.T) { + provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{}) + + require.True(t, provider.SupportsAccount(geminiAPIKeyAccount("sk-gemini"))) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{}})) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"api_key": "sk"}})) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk"}})) + require.False(t, provider.SupportsAccount(nil)) +} + +func TestGeminiProvider_MissingAPIKeyRejected(t *testing.T) { + provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{}) + _, err := provider.Submit(context.Background(), nil, &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey}, validGeminiBatchInput()) + require.ErrorIs(t, err, ErrBatchImageProviderMissingAPIKey) +} + +func TestBuildGeminiBatchJSONL_WritesValidLinesAndPreservesCustomID(t *testing.T) { + input := validGeminiBatchInput() + input.Items = append(input.Items, BatchImageInputItem{CustomID: "cover_002", Prompt: "Second prompt"}) + + jsonl, err := BuildGeminiBatchJSONL(input) + require.NoError(t, err) + + lines := strings.Split(strings.TrimSpace(string(jsonl)), "\n") + require.Len(t, lines, 2) + requireJSONLLine(t, lines[0], "cover_001", "A clean product hero image") + requireJSONLLine(t, lines[1], "cover_002", "Second prompt") +} + +func TestBuildGeminiBatchJSONL_RejectsDuplicateCustomIDs(t *testing.T) { + input := validGeminiBatchInput() + input.Items = append(input.Items, BatchImageInputItem{CustomID: "cover_001", Prompt: "Duplicate"}) + + _, err := BuildGeminiBatchJSONL(input) + require.ErrorIs(t, err, ErrBatchImageProviderInvalidInput) +} + +func TestBuildGeminiBatchJSONL_RejectsEmptyPrompt(t *testing.T) { + input := validGeminiBatchInput() + input.Items[0].Prompt = " " + + _, err := BuildGeminiBatchJSONL(input) + require.ErrorIs(t, err, ErrBatchImageProviderInvalidInput) +} + +func TestBuildGeminiBatchJSONL_WritesReferenceImages(t *testing.T) { + input := validGeminiBatchInput() + input.Items[0].ReferenceImages = []BatchImageReference{ + {MimeType: "image/webp", Data: []byte("webp-bytes")}, + {MimeType: "image/jpeg", FileURI: "gs://bucket/refs/style.jpg"}, + } + + jsonl, err := BuildGeminiBatchJSONL(input) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(jsonl)), "\n") + require.Len(t, lines, 1) + + var got map[string]any + require.NoError(t, json.Unmarshal([]byte(lines[0]), &got)) + request := got["request"].(map[string]any) + contents := request["contents"].([]any) + parts := contents[0].(map[string]any)["parts"].([]any) + require.Len(t, parts, 3) + require.Equal(t, "A clean product hero image", parts[0].(map[string]any)["text"]) + inlineData := parts[1].(map[string]any)["inlineData"].(map[string]any) + require.Equal(t, "image/webp", inlineData["mimeType"]) + require.Equal(t, "d2VicC1ieXRlcw==", inlineData["data"]) + fileData := parts[2].(map[string]any)["fileData"].(map[string]any) + require.Equal(t, "image/jpeg", fileData["mimeType"]) + require.Equal(t, "gs://bucket/refs/style.jpg", fileData["fileUri"]) +} + +func TestGeminiProvider_SubmitUploadsJSONLThenCreatesBatch(t *testing.T) { + client := &fakeGeminiBatchClient{ + uploaded: &GeminiUploadedFile{Name: "files/input-jsonl"}, + created: &GeminiBatchJob{Name: "batches/job-123", State: "JOB_STATE_PENDING"}, + } + provider := NewGeminiAPIBatchImageProvider(client) + + got, err := provider.Submit(context.Background(), &BatchImageJob{BatchID: "imgbatch_123", Model: "gemini-3.1-flash-image"}, geminiAPIKeyAccount("sk-secret"), validGeminiBatchInput()) + require.NoError(t, err) + require.Equal(t, []string{"upload", "create"}, client.calls) + require.Equal(t, "files/input-jsonl", got.ProviderInputRef) + require.Equal(t, "batches/job-123", got.ProviderJobName) + require.Empty(t, got.ProviderOutputRef) + require.NotContains(t, got.ProviderInputRef, "A clean product hero image") + require.NotContains(t, string(client.uploadedJSONL), "sk-secret") +} + +func TestGeminiProvider_GetMapsStates(t *testing.T) { + tests := []struct { + name string + job *GeminiBatchJob + wantState BatchProviderInternalState + wantDone bool + wantRef string + wantCode string + }{ + {name: "running", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_RUNNING"}, wantState: BatchProviderStateRunning}, + {name: "succeeded_dest_fileName", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_SUCCEEDED", Dest: &GeminiBatchDest{FileName: "files/out"}}, wantState: BatchProviderStateSucceeded, wantDone: true, wantRef: "files/out"}, + {name: "failed", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_FAILED", Error: &GeminiBatchError{Code: "BAD_PROMPT", Message: "bad prompt"}}, wantState: BatchProviderStateFailed, wantDone: true, wantCode: "BAD_PROMPT"}, + {name: "cancelled", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_CANCELLED"}, wantState: BatchProviderStateCancelled, wantDone: true, wantCode: "GEMINI_BATCH_CANCELLED"}, + {name: "expired", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_EXPIRED"}, wantState: BatchProviderStateExpired, wantDone: true, wantCode: "GEMINI_BATCH_EXPIRED"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{got: tt.job}) + got, err := provider.Get(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret")) + require.NoError(t, err) + require.Equal(t, tt.wantState, got.InternalState) + require.Equal(t, tt.wantDone, got.Done) + require.Equal(t, tt.wantRef, got.ProviderOutputRef) + require.Equal(t, tt.wantCode, got.ErrorCode) + require.NotContains(t, got.ErrorMessage, "sk-secret") + }) + } +} + +func TestGeminiProvider_GetExtractsResponsesFileReference(t *testing.T) { + provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{ + got: &GeminiBatchJob{ + Name: "batches/1", + State: "JOB_STATE_SUCCEEDED", + Response: &GeminiBatchResponse{ResponsesFile: "files/responses-jsonl"}, + }, + }) + + got, err := provider.Get(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret")) + require.NoError(t, err) + require.Equal(t, BatchProviderStateSucceeded, got.InternalState) + require.Equal(t, "files/responses-jsonl", got.ProviderOutputRef) +} + +func TestGeminiProvider_GetRejectsInlineResultShape(t *testing.T) { + provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{ + got: &GeminiBatchJob{ + Name: "batches/1", + State: "JOB_STATE_SUCCEEDED", + Response: &GeminiBatchResponse{InlinedResponses: []any{map[string]any{"response": "large"}}}, + }, + }) + + _, err := provider.Get(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret")) + require.ErrorIs(t, err, ErrBatchImageProviderInlineResultUnsupported) +} + +func TestGeminiProvider_OpenResultStreamsResultFile(t *testing.T) { + client := &fakeGeminiBatchClient{downloadBody: "line1\n", downloadContentType: "application/jsonl"} + provider := NewGeminiAPIBatchImageProvider(client) + + outputRef := "files/output-jsonl" + r, contentType, err := provider.OpenResult(context.Background(), &BatchImageJob{ProviderOutputRef: &outputRef}, geminiAPIKeyAccount("sk-secret")) + require.NoError(t, err) + defer r.Close() + + body, err := io.ReadAll(r) + require.NoError(t, err) + require.Equal(t, "line1\n", string(body)) + require.Equal(t, "application/jsonl", contentType) + require.Equal(t, "files/output-jsonl", client.downloadedFile) +} + +func TestGeminiProvider_CancelCallsClient(t *testing.T) { + client := &fakeGeminiBatchClient{} + provider := NewGeminiAPIBatchImageProvider(client) + + require.NoError(t, provider.Cancel(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret"))) + require.Equal(t, "batches/1", client.cancelledBatch) +} + +func TestGeminiProvider_CleanupDeletesRefsOnlyWhenPresent(t *testing.T) { + inputRef := "files/input" + outputRef := "files/output" + client := &fakeGeminiBatchClient{} + provider := NewGeminiAPIBatchImageProvider(client) + + err := provider.Cleanup(context.Background(), &BatchImageJob{ProviderInputRef: &inputRef, ProviderOutputRef: &outputRef}, geminiAPIKeyAccount("sk-secret"), CleanupTargetAll) + require.NoError(t, err) + require.Equal(t, []string{"files/input", "files/output"}, client.deletedFiles) + + err = provider.Cleanup(context.Background(), &BatchImageJob{}, geminiAPIKeyAccount("sk-secret"), CleanupTargetAll) + require.NoError(t, err) + require.Equal(t, []string{"files/input", "files/output"}, client.deletedFiles) +} + +func TestGeminiProvider_ErrorsDoNotExposeAPIKey(t *testing.T) { + apiKey := "sk-top-secret" + client := &fakeGeminiBatchClient{uploadErr: &GeminiAPIError{StatusCode: 401, Message: "upstream body should be hidden " + apiKey}} + provider := NewGeminiAPIBatchImageProvider(client) + + _, err := provider.Submit(context.Background(), nil, geminiAPIKeyAccount(apiKey), validGeminiBatchInput()) + require.Error(t, err) + require.Equal(t, "GEMINI_AUTH_FAILED", infraerrors.Reason(err)) + require.NotContains(t, err.Error(), apiKey) +} + +func TestGeminiProvider_MetadataDoesNotStoreImageBytesOrBase64(t *testing.T) { + client := &fakeGeminiBatchClient{ + uploaded: &GeminiUploadedFile{Name: "files/input-jsonl"}, + created: &GeminiBatchJob{Name: "batches/job-123", State: "JOB_STATE_PENDING"}, + } + provider := NewGeminiAPIBatchImageProvider(client) + + got, err := provider.Submit(context.Background(), nil, geminiAPIKeyAccount("sk-secret"), validGeminiBatchInput()) + require.NoError(t, err) + require.NotContains(t, got.ProviderJobName, "base64") + require.NotContains(t, got.ProviderInputRef, "base64") + require.NotContains(t, got.ProviderOutputRef, "base64") + require.NotContains(t, got.ProviderJobName+got.ProviderInputRef+got.ProviderOutputRef, "iVBOR") + require.NotContains(t, got.ProviderJobName+got.ProviderInputRef+got.ProviderOutputRef, "A clean product hero image") +} + +func requireJSONLLine(t *testing.T, line, wantKey, wantPrompt string) { + t.Helper() + var got map[string]any + require.NoError(t, json.Unmarshal([]byte(line), &got)) + require.Equal(t, wantKey, got["key"]) + request := got["request"].(map[string]any) + config := request["generationConfig"].(map[string]any) + require.Equal(t, []any{"TEXT", "IMAGE"}, config["responseModalities"]) + contents := request["contents"].([]any) + parts := contents[0].(map[string]any)["parts"].([]any) + require.Equal(t, wantPrompt, parts[0].(map[string]any)["text"]) +} + +func validGeminiBatchInput() BatchImageInput { + return BatchImageInput{ + BatchID: "imgbatch_123", + Model: "gemini-3.1-flash-image", + DisplayName: "test batch", + Items: []BatchImageInputItem{{ + CustomID: "cover_001", + Prompt: "A clean product hero image", + }}, + } +} + +func geminiAPIKeyAccount(apiKey string) *Account { + return &Account{ + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": apiKey}, + } +} + +func jobWithProviderName(name string) *BatchImageJob { + return &BatchImageJob{ProviderJobName: &name} +} + +type fakeGeminiBatchClient struct { + calls []string + uploaded *GeminiUploadedFile + created *GeminiBatchJob + got *GeminiBatchJob + uploadErr error + createErr error + getErr error + cancelErr error + downloadErr error + deleteErr error + uploadedJSONL []byte + createdFile string + cancelledBatch string + downloadedFile string + downloadBody string + downloadContentType string + deletedFiles []string +} + +func (f *fakeGeminiBatchClient) UploadJSONL(_ context.Context, apiKey string, _ string, r io.Reader) (*GeminiUploadedFile, error) { + if strings.TrimSpace(apiKey) == "" { + return nil, errors.New("missing api key") + } + f.calls = append(f.calls, "upload") + f.uploadedJSONL, _ = io.ReadAll(r) + if f.uploadErr != nil { + return nil, f.uploadErr + } + if f.uploaded != nil { + return f.uploaded, nil + } + return &GeminiUploadedFile{Name: "files/input-jsonl"}, nil +} + +func (f *fakeGeminiBatchClient) CreateBatch(_ context.Context, _ string, _ string, fileName string, _ string) (*GeminiBatchJob, error) { + f.calls = append(f.calls, "create") + f.createdFile = fileName + if f.createErr != nil { + return nil, f.createErr + } + if f.created != nil { + return f.created, nil + } + return &GeminiBatchJob{Name: "batches/job-123", State: "JOB_STATE_PENDING"}, nil +} + +func (f *fakeGeminiBatchClient) GetBatch(_ context.Context, _ string, _ string) (*GeminiBatchJob, error) { + f.calls = append(f.calls, "get") + if f.getErr != nil { + return nil, f.getErr + } + return f.got, nil +} + +func (f *fakeGeminiBatchClient) CancelBatch(_ context.Context, _ string, batchName string) error { + f.calls = append(f.calls, "cancel") + f.cancelledBatch = batchName + return f.cancelErr +} + +func (f *fakeGeminiBatchClient) DownloadFile(_ context.Context, _ string, fileName string) (io.ReadCloser, string, error) { + f.calls = append(f.calls, "download") + f.downloadedFile = fileName + if f.downloadErr != nil { + return nil, "", f.downloadErr + } + contentType := f.downloadContentType + if contentType == "" { + contentType = "application/octet-stream" + } + return io.NopCloser(bytes.NewBufferString(f.downloadBody)), contentType, nil +} + +func (f *fakeGeminiBatchClient) DeleteFile(_ context.Context, _ string, fileName string) error { + f.calls = append(f.calls, "delete") + f.deletedFiles = append(f.deletedFiles, fileName) + return f.deleteErr +} diff --git a/backend/internal/service/batch_image_provider_vertex.go b/backend/internal/service/batch_image_provider_vertex.go new file mode 100644 index 0000000000..bb3363dfe3 --- /dev/null +++ b/backend/internal/service/batch_image_provider_vertex.go @@ -0,0 +1,997 @@ +package service + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "sort" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +const ( + defaultVertexBatchRequeueAfter = 30 * time.Second + defaultVertexBatchLocation = "global" + defaultVertexManagedGCSPrefix = "batch-image/{env}/{batch_id}" +) + +type VertexBatchImageProviderOptions struct { + Enabled bool + ProjectID string + Location string + ManagedGCSBucket string + ManagedGCSPrefix string + Environment string + InputRetentionHours int + OutputRetentionHours int + BatchPredictionBaseURL string + GCSBaseURL string +} + +func NewVertexBatchImageProviderOptionsFromConfig(cfg *config.Config) VertexBatchImageProviderOptions { + if cfg == nil { + return VertexBatchImageProviderOptions{} + } + return VertexBatchImageProviderOptions{ + Enabled: cfg.BatchImage.VertexEnabled, + ProjectID: cfg.BatchImage.VertexProjectID, + Location: cfg.BatchImage.VertexLocation, + ManagedGCSBucket: cfg.BatchImage.VertexManagedGCSBucket, + ManagedGCSPrefix: cfg.BatchImage.VertexManagedGCSPrefix, + Environment: cfg.Log.Environment, + InputRetentionHours: cfg.BatchImage.VertexInputRetentionHours, + OutputRetentionHours: cfg.BatchImage.VertexOutputRetentionHours, + BatchPredictionBaseURL: cfg.BatchImage.VertexBatchPredictionBaseURL, + GCSBaseURL: cfg.BatchImage.VertexGCSBaseURL, + } +} + +type VertexBatchClient interface { + CreateBatchPredictionJob(ctx context.Context, accessToken string, req VertexCreateBatchPredictionJobRequest) (*VertexBatchPredictionJob, error) + GetBatchPredictionJob(ctx context.Context, accessToken string, name string) (*VertexBatchPredictionJob, error) + CancelBatchPredictionJob(ctx context.Context, accessToken string, name string) error +} + +type VertexBatchObjectStore interface { + UploadJSONL(ctx context.Context, accessToken string, uri string, r io.Reader) error + ListJSONLObjects(ctx context.Context, accessToken string, prefixURI string) ([]string, error) + OpenObject(ctx context.Context, accessToken string, uri string) (io.ReadCloser, string, error) + DeleteObject(ctx context.Context, accessToken string, uri string) error + DeletePrefix(ctx context.Context, accessToken string, prefixURI string) error +} + +type VertexCreateBatchPredictionJobRequest struct { + ProjectID string `json:"-"` + Location string `json:"-"` + DisplayName string `json:"displayName"` + Model string `json:"model"` + InputConfig VertexBatchInputConfig `json:"inputConfig"` + OutputConfig VertexBatchOutputConfig `json:"outputConfig"` + InstanceConfig *VertexBatchInstanceConfig `json:"instanceConfig,omitempty"` +} + +type VertexBatchInputConfig struct { + InstancesFormat string `json:"instancesFormat"` + GCSSource VertexBatchGCSSource `json:"gcsSource"` +} + +type VertexBatchGCSSource struct { + URIs []string `json:"uris"` +} + +type VertexBatchOutputConfig struct { + PredictionsFormat string `json:"predictionsFormat"` + GCSDestination VertexBatchGCSDestination `json:"gcsDestination"` +} + +type VertexBatchGCSDestination struct { + OutputURIPrefix string `json:"outputUriPrefix"` +} + +type VertexBatchInstanceConfig struct { + KeyField string `json:"keyField"` +} + +type VertexBatchPredictionJob struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + State string `json:"state"` + OutputConfig VertexBatchOutputConfig `json:"outputConfig"` + Error *VertexBatchJobError `json:"error"` +} + +type VertexBatchJobError struct { + Code any `json:"code"` + Message string `json:"message"` + Status string `json:"status"` +} + +type VertexBatchImageProvider struct { + opts VertexBatchImageProviderOptions + client VertexBatchClient + objectStore VertexBatchObjectStore + tokenCache GeminiTokenCache +} + +func NewVertexBatchImageProvider(opts VertexBatchImageProviderOptions, client VertexBatchClient, objectStore VertexBatchObjectStore, tokenCache GeminiTokenCache) *VertexBatchImageProvider { + opts = normalizeVertexBatchImageProviderOptions(opts) + if client == nil { + client = NewVertexBatchHTTPClient(opts.BatchPredictionBaseURL, nil) + } + if objectStore == nil { + objectStore = NewVertexGCSObjectStore(opts.GCSBaseURL, nil) + } + return &VertexBatchImageProvider{ + opts: opts, + client: client, + objectStore: objectStore, + tokenCache: tokenCache, + } +} + +func NewVertexBatchImageProviderFromConfig(cfg *config.Config, client VertexBatchClient, objectStore VertexBatchObjectStore, tokenCache GeminiTokenCache) *VertexBatchImageProvider { + return NewVertexBatchImageProvider(NewVertexBatchImageProviderOptionsFromConfig(cfg), client, objectStore, tokenCache) +} + +func normalizeVertexBatchImageProviderOptions(opts VertexBatchImageProviderOptions) VertexBatchImageProviderOptions { + opts.ProjectID = strings.TrimSpace(opts.ProjectID) + opts.Location = strings.TrimSpace(opts.Location) + if opts.Location == "" { + opts.Location = defaultVertexBatchLocation + } + opts.ManagedGCSBucket = strings.Trim(strings.TrimSpace(opts.ManagedGCSBucket), "/") + opts.ManagedGCSPrefix = strings.Trim(strings.TrimSpace(opts.ManagedGCSPrefix), "/") + if opts.ManagedGCSPrefix == "" { + opts.ManagedGCSPrefix = defaultVertexManagedGCSPrefix + } + opts.Environment = strings.TrimSpace(opts.Environment) + if opts.Environment == "" { + opts.Environment = "default" + } + opts.BatchPredictionBaseURL = strings.TrimRight(strings.TrimSpace(opts.BatchPredictionBaseURL), "/") + opts.GCSBaseURL = strings.TrimRight(strings.TrimSpace(opts.GCSBaseURL), "/") + return opts +} + +func (p *VertexBatchImageProvider) Name() string { + return BatchImageProviderVertex +} + +func (p *VertexBatchImageProvider) SupportsAccount(account *Account) bool { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeServiceAccount { + return false + } + _, err := parseVertexServiceAccountKey(account) + return err == nil +} + +func (p *VertexBatchImageProvider) Submit(ctx context.Context, job *BatchImageJob, account *Account, input BatchImageInput) (*BatchProviderJob, error) { + if err := p.validateAccount(account); err != nil { + return nil, err + } + if strings.TrimSpace(p.opts.ManagedGCSBucket) == "" { + return nil, vertexProviderError("VERTEX_MANAGED_GCS_BUCKET_MISSING", "Vertex managed GCS bucket is not configured", nil) + } + if input.BatchID == "" && job != nil { + input.BatchID = job.BatchID + } + if input.Model == "" && job != nil { + input.Model = job.Model + } + + jsonl, err := BuildVertexBatchJSONL(input) + if err != nil { + return nil, err + } + refs, err := p.managedRefs(input.BatchID) + if err != nil { + return nil, err + } + + accessToken, err := p.accessToken(ctx, account) + if err != nil { + return nil, mapVertexClientError(err) + } + if err := p.objectStore.UploadJSONL(ctx, accessToken, refs.InputURI, bytes.NewReader(jsonl)); err != nil { + return nil, vertexProviderError("VERTEX_GCS_UPLOAD_FAILED", "Vertex managed GCS upload failed", nil) + } + + projectID := strings.TrimSpace(p.opts.ProjectID) + if projectID == "" { + projectID = account.VertexProjectID() + } + if projectID == "" { + return nil, vertexProviderError("VERTEX_PROJECT_ID_MISSING", "Vertex project id is not configured", nil) + } + location := strings.TrimSpace(p.opts.Location) + if location == "" { + location = account.VertexLocation(input.Model) + } + + req := VertexCreateBatchPredictionJobRequest{ + ProjectID: projectID, + Location: location, + DisplayName: vertexBatchDisplayName(input), + Model: NormalizeVertexBatchModelPath(input.Model), + InputConfig: VertexBatchInputConfig{InstancesFormat: "jsonl", GCSSource: VertexBatchGCSSource{URIs: []string{refs.InputURI}}}, + OutputConfig: VertexBatchOutputConfig{PredictionsFormat: "jsonl", GCSDestination: VertexBatchGCSDestination{OutputURIPrefix: refs.OutputPrefixURI}}, + InstanceConfig: &VertexBatchInstanceConfig{KeyField: "key"}, + } + created, err := p.client.CreateBatchPredictionJob(ctx, accessToken, req) + if err != nil { + return nil, mapVertexClientError(err) + } + if created == nil || strings.TrimSpace(created.Name) == "" { + return nil, vertexProviderError("VERTEX_INVALID_RESPONSE", "Vertex batch response is missing job name", nil) + } + return &BatchProviderJob{ + ProviderJobName: created.Name, + ProviderInputRef: refs.InputURI, + ProviderOutputRef: refs.OutputPrefixURI, + RawState: created.State, + }, nil +} + +func (p *VertexBatchImageProvider) Get(ctx context.Context, job *BatchImageJob, account *Account) (*BatchProviderStatus, error) { + if err := p.validateAccount(account); err != nil { + return nil, err + } + jobName := batchImageProviderJobName(job) + if jobName == "" { + return nil, ErrBatchImageProviderMissingJobName + } + accessToken, err := p.accessToken(ctx, account) + if err != nil { + return nil, mapVertexClientError(err) + } + vertexJob, err := p.client.GetBatchPredictionJob(ctx, accessToken, jobName) + if err != nil { + return nil, mapVertexClientError(err) + } + if vertexJob == nil { + return nil, vertexProviderError("VERTEX_INVALID_RESPONSE", "Vertex batch response is empty", nil) + } + status := mapVertexBatchState(vertexJob) + outputRef := strings.TrimSpace(vertexJob.OutputConfig.GCSDestination.OutputURIPrefix) + if outputRef == "" { + outputRef = batchImageProviderOutputRef(job) + } + if outputRef == "" && job != nil && job.GCSOutputURI != nil { + outputRef = strings.TrimSpace(*job.GCSOutputURI) + } + status.ProviderOutputRef = outputRef + return status, nil +} + +func (p *VertexBatchImageProvider) Cancel(ctx context.Context, job *BatchImageJob, account *Account) error { + if err := p.validateAccount(account); err != nil { + return err + } + jobName := batchImageProviderJobName(job) + if jobName == "" { + return ErrBatchImageProviderMissingJobName + } + accessToken, err := p.accessToken(ctx, account) + if err != nil { + return mapVertexClientError(err) + } + return mapVertexClientError(p.client.CancelBatchPredictionJob(ctx, accessToken, jobName)) +} + +func (p *VertexBatchImageProvider) OpenResult(ctx context.Context, job *BatchImageJob, account *Account) (io.ReadCloser, string, error) { + if err := p.validateAccount(account); err != nil { + return nil, "", err + } + outputRef := batchImageProviderOutputRef(job) + if outputRef == "" && job != nil && job.GCSOutputURI != nil { + outputRef = strings.TrimSpace(*job.GCSOutputURI) + } + if outputRef == "" { + return nil, "", ErrBatchImageProviderMissingResultRef + } + accessToken, err := p.accessToken(ctx, account) + if err != nil { + return nil, "", mapVertexClientError(err) + } + objects, err := p.objectStore.ListJSONLObjects(ctx, accessToken, outputRef) + if err != nil { + return nil, "", vertexProviderError("VERTEX_GCS_LIST_FAILED", "Vertex managed GCS list failed", nil) + } + sort.Strings(objects) + if len(objects) == 0 { + return nil, "", vertexProviderError("VERTEX_RESULT_OBJECTS_MISSING", "Vertex result objects are missing", nil) + } + return &vertexCombinedJSONLReadCloser{ + ctx: ctx, + accessToken: accessToken, + objects: objects, + store: p.objectStore, + }, "application/jsonl", nil +} + +func (p *VertexBatchImageProvider) Cleanup(ctx context.Context, job *BatchImageJob, account *Account, target CleanupTarget) error { + if err := p.validateAccount(account); err != nil { + return err + } + accessToken, err := p.accessToken(ctx, account) + if err != nil { + return mapVertexClientError(err) + } + inputRef := batchImageProviderInputRef(job) + outputRef := batchImageProviderOutputRef(job) + if job != nil { + if inputRef == "" && job.GCSInputURI != nil { + inputRef = strings.TrimSpace(*job.GCSInputURI) + } + if outputRef == "" && job.GCSOutputURI != nil { + outputRef = strings.TrimSpace(*job.GCSOutputURI) + } + } + + switch target { + case CleanupTargetInput: + return p.deleteManagedInput(ctx, accessToken, job, inputRef) + case CleanupTargetOutput: + return p.deleteManagedOutput(ctx, accessToken, job, outputRef) + case CleanupTargetAll: + if err := p.deleteManagedInput(ctx, accessToken, job, inputRef); err != nil { + return err + } + return p.deleteManagedOutput(ctx, accessToken, job, outputRef) + default: + return ErrUnsupportedCleanupTarget + } +} + +func (p *VertexBatchImageProvider) validateAccount(account *Account) error { + if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeServiceAccount { + return ErrBatchImageProviderUnsupportedAccount + } + if _, err := parseVertexServiceAccountKey(account); err != nil { + return ErrBatchImageProviderMissingServiceAccount + } + return nil +} + +func (p *VertexBatchImageProvider) accessToken(ctx context.Context, account *Account) (string, error) { + return getVertexServiceAccountAccessToken(ctx, p.tokenCache, account) +} + +func (p *VertexBatchImageProvider) deleteManagedInput(ctx context.Context, accessToken string, job *BatchImageJob, uri string) error { + if strings.TrimSpace(uri) == "" { + return nil + } + if !p.isSafeManagedInput(job, uri) { + return ErrBatchImageProviderUnsafeCleanupPath + } + return mapVertexClientError(p.objectStore.DeleteObject(ctx, accessToken, uri)) +} + +func (p *VertexBatchImageProvider) deleteManagedOutput(ctx context.Context, accessToken string, job *BatchImageJob, uri string) error { + if strings.TrimSpace(uri) == "" { + return nil + } + if !p.isSafeManagedOutput(job, uri) { + return ErrBatchImageProviderUnsafeCleanupPath + } + return mapVertexClientError(p.objectStore.DeletePrefix(ctx, accessToken, uri)) +} + +func (p *VertexBatchImageProvider) isSafeManagedInput(job *BatchImageJob, uri string) bool { + if job == nil || strings.TrimSpace(job.BatchID) == "" { + return false + } + refs, err := p.managedRefs(job.BatchID) + return err == nil && strings.TrimSpace(uri) == refs.InputURI +} + +func (p *VertexBatchImageProvider) isSafeManagedOutput(job *BatchImageJob, uri string) bool { + if job == nil || strings.TrimSpace(job.BatchID) == "" { + return false + } + refs, err := p.managedRefs(job.BatchID) + return err == nil && strings.HasPrefix(strings.TrimSpace(uri), refs.OutputPrefixURI) +} + +type vertexManagedRefs struct { + Prefix string + InputURI string + OutputPrefixURI string +} + +func (p *VertexBatchImageProvider) managedRefs(batchID string) (vertexManagedRefs, error) { + batchID = strings.TrimSpace(batchID) + if !IsValidBatchImageID(batchID) { + return vertexManagedRefs{}, batchImageProviderInputError("valid batch_id is required") + } + bucket := strings.Trim(strings.TrimSpace(p.opts.ManagedGCSBucket), "/") + if bucket == "" || strings.Contains(bucket, "://") { + return vertexManagedRefs{}, vertexProviderError("VERTEX_MANAGED_GCS_BUCKET_MISSING", "Vertex managed GCS bucket is not configured", nil) + } + prefix := buildVertexManagedGCSPrefix(p.opts.ManagedGCSPrefix, p.opts.Environment, batchID) + if !strings.Contains(prefix, batchID) { + return vertexManagedRefs{}, batchImageProviderInputError("managed GCS prefix must contain batch_id") + } + base := "gs://" + bucket + "/" + strings.Trim(prefix, "/") + return vertexManagedRefs{ + Prefix: strings.Trim(prefix, "/"), + InputURI: base + "/input/requests.jsonl", + OutputPrefixURI: base + "/output/", + }, nil +} + +func buildVertexManagedGCSPrefix(template, env, batchID string) string { + template = strings.Trim(strings.TrimSpace(template), "/") + if template == "" { + template = defaultVertexManagedGCSPrefix + } + env = sanitizeVertexGCSPathSegment(env) + batchID = sanitizeVertexGCSPathSegment(batchID) + prefix := strings.ReplaceAll(template, "{env}", env) + prefix = strings.ReplaceAll(prefix, "{batch_id}", batchID) + return strings.Trim(prefix, "/") +} + +func sanitizeVertexGCSPathSegment(v string) string { + v = strings.TrimSpace(v) + if v == "" { + return "default" + } + var b strings.Builder + for _, r := range v { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9': + _, _ = b.WriteRune(r) + case r == '-', r == '_', r == '.': + _, _ = b.WriteRune(r) + default: + _ = b.WriteByte('-') + } + } + return b.String() +} + +func vertexBatchDisplayName(input BatchImageInput) string { + if v := strings.TrimSpace(input.DisplayName); v != "" { + return v + } + if v := strings.TrimSpace(input.BatchID); v != "" { + return "sub2api-" + v + } + return "sub2api-image-batch" +} + +func BuildVertexBatchJSONL(input BatchImageInput) ([]byte, error) { + if strings.TrimSpace(input.Model) == "" { + return nil, batchImageProviderInputError("model is required") + } + if len(input.Items) == 0 { + return nil, batchImageProviderInputError("at least one item is required") + } + seen := make(map[string]struct{}, len(input.Items)) + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + for _, item := range input.Items { + customID := strings.TrimSpace(item.CustomID) + if customID == "" { + return nil, batchImageProviderInputError("custom_id is required") + } + if _, ok := seen[customID]; ok { + return nil, batchImageProviderInputError("duplicate custom_id %q", customID) + } + seen[customID] = struct{}{} + prompt := strings.TrimSpace(item.Prompt) + if prompt == "" { + return nil, batchImageProviderInputError("prompt is required for custom_id %q", customID) + } + parts, err := vertexBatchImageParts(prompt, item.ReferenceImages) + if err != nil { + return nil, err + } + line := map[string]any{ + "key": customID, + "request": map[string]any{ + "contents": []any{map[string]any{ + "role": "user", + "parts": parts, + }}, + "generationConfig": map[string]any{ + "responseModalities": []string{"TEXT", "IMAGE"}, + }, + }, + } + if err := enc.Encode(line); err != nil { + return nil, err + } + } + return buf.Bytes(), nil +} + +func vertexBatchImageParts(prompt string, refs []BatchImageReference) ([]any, error) { + parts := []any{map[string]any{"text": prompt}} + for _, ref := range refs { + mimeType := normalizeBatchImageReferenceMimeType(ref.MimeType) + if mimeType == "" { + return nil, batchImageProviderInputError("reference image mime_type is required") + } + fileURI := strings.TrimSpace(ref.FileURI) + switch { + case len(ref.Data) > 0 && fileURI == "": + parts = append(parts, map[string]any{ + "inlineData": map[string]any{ + "mimeType": mimeType, + "data": base64.StdEncoding.EncodeToString(ref.Data), + }, + }) + case len(ref.Data) == 0 && fileURI != "": + parts = append(parts, map[string]any{ + "fileData": map[string]any{ + "mimeType": mimeType, + "fileUri": fileURI, + }, + }) + default: + return nil, batchImageProviderInputError("reference image must contain exactly one of data or file_uri") + } + } + return parts, nil +} + +func NormalizeVertexBatchModelPath(model string) string { + model = strings.Trim(strings.TrimSpace(model), "/") + if strings.HasPrefix(model, "publishers/") || strings.HasPrefix(model, "projects/") { + return model + } + return "publishers/google/models/" + model +} + +func BuildVertexBatchPredictionJobsEndpoint(baseURL, projectID, location string) (string, error) { + projectID = strings.TrimSpace(projectID) + location = strings.TrimSpace(location) + if projectID == "" { + return "", errors.New("vertex project_id is required") + } + if location == "" { + location = defaultVertexBatchLocation + } + if !vertexLocationPattern.MatchString(location) { + return "", fmt.Errorf("invalid vertex location: %s", location) + } + if strings.TrimSpace(baseURL) != "" { + return strings.TrimRight(strings.TrimSpace(baseURL), "/") + "/v1/projects/" + url.PathEscape(projectID) + "/locations/" + url.PathEscape(location) + "/batchPredictionJobs", nil + } + host := fmt.Sprintf("%s-aiplatform.googleapis.com", location) + if location == "global" { + host = "aiplatform.googleapis.com" + } + return fmt.Sprintf("https://%s/v1/projects/%s/locations/%s/batchPredictionJobs", host, url.PathEscape(projectID), url.PathEscape(location)), nil +} + +func mapVertexBatchState(job *VertexBatchPredictionJob) *BatchProviderStatus { + state := strings.TrimSpace(job.State) + status := &BatchProviderStatus{ + RawState: state, + InternalState: BatchProviderStateRunning, + SuggestedRequeueAfter: defaultVertexBatchRequeueAfter, + } + switch strings.ToUpper(state) { + case "JOB_STATE_PENDING", "JOB_STATE_QUEUED": + status.InternalState = BatchProviderStateQueued + case "JOB_STATE_RUNNING", "JOB_STATE_PAUSED": + status.InternalState = BatchProviderStateRunning + case "JOB_STATE_SUCCEEDED": + status.InternalState = BatchProviderStateSucceeded + status.Done = true + status.SuggestedRequeueAfter = 0 + case "JOB_STATE_FAILED": + status.InternalState = BatchProviderStateFailed + status.Done = true + status.ErrorCode = "VERTEX_BATCH_FAILED" + status.SuggestedRequeueAfter = 0 + case "JOB_STATE_CANCELLED": + status.InternalState = BatchProviderStateCancelled + status.Done = true + status.ErrorCode = "VERTEX_BATCH_CANCELLED" + status.SuggestedRequeueAfter = 0 + case "JOB_STATE_EXPIRED": + status.InternalState = BatchProviderStateExpired + status.Done = true + status.ErrorCode = "VERTEX_BATCH_EXPIRED" + status.SuggestedRequeueAfter = 0 + default: + if job.Error != nil && strings.TrimSpace(job.Error.Message) != "" { + status.InternalState = BatchProviderStateFailed + status.Done = true + status.ErrorCode = "VERTEX_BATCH_FAILED" + status.SuggestedRequeueAfter = 0 + } + } + if job.Error != nil { + if code := strings.TrimSpace(job.Error.Status); code != "" { + status.ErrorCode = code + } + status.ErrorMessage = strings.TrimSpace(job.Error.Message) + } + return status +} + +func vertexProviderError(reason, message string, cause error) error { + err := infraerrors.New(http.StatusBadGateway, reason, message) + if cause != nil { + return err.WithCause(cause) + } + return err +} + +func mapVertexClientError(err error) error { + if err == nil { + return nil + } + if errors.Is(err, ErrBatchImageProviderMissingServiceAccount) || + errors.Is(err, ErrBatchImageProviderMissingJobName) || + errors.Is(err, ErrBatchImageProviderMissingResultRef) || + errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) || + errors.Is(err, ErrUnsupportedCleanupTarget) { + return err + } + var apiErr *VertexAPIError + if errors.As(err, &apiErr) { + switch apiErr.StatusCode { + case http.StatusUnauthorized: + return vertexProviderError("VERTEX_AUTH_FAILED", "Vertex authentication failed", nil) + case http.StatusForbidden: + return vertexProviderError("VERTEX_PERMISSION_DENIED", "Vertex permission denied", nil) + case http.StatusTooManyRequests: + return vertexProviderError("VERTEX_RATE_LIMITED", "Vertex rate limit exceeded", nil) + case http.StatusNotFound: + return vertexProviderError("VERTEX_BATCH_NOT_FOUND", "Vertex batch resource was not found", nil) + default: + return vertexProviderError("VERTEX_INVALID_RESPONSE", "Vertex API request failed", nil) + } + } + return vertexProviderError("VERTEX_INVALID_RESPONSE", "Vertex API request failed", err) +} + +type vertexCombinedJSONLReadCloser struct { + ctx context.Context + accessToken string + objects []string + store VertexBatchObjectStore + index int + current io.ReadCloser + needBoundary bool + closed bool +} + +func (r *vertexCombinedJSONLReadCloser) Read(p []byte) (int, error) { + if r.closed { + return 0, io.ErrClosedPipe + } + if r.needBoundary { + if len(p) == 0 { + return 0, nil + } + p[0] = '\n' + r.needBoundary = false + return 1, nil + } + for { + if r.current == nil { + if r.index >= len(r.objects) { + return 0, io.EOF + } + obj := r.objects[r.index] + r.index++ + rc, _, err := r.store.OpenObject(r.ctx, r.accessToken, obj) + if err != nil { + return 0, err + } + r.current = rc + } + n, err := r.current.Read(p) + if err == io.EOF { + _ = r.current.Close() + r.current = nil + if r.index < len(r.objects) { + if n > 0 { + r.needBoundary = true + return n, nil + } + if len(p) == 0 { + return 0, nil + } + p[0] = '\n' + return 1, nil + } + if n > 0 { + return n, nil + } + continue + } + return n, err + } +} + +func (r *vertexCombinedJSONLReadCloser) Close() error { + r.closed = true + if r.current != nil { + return r.current.Close() + } + return nil +} + +type VertexBatchHTTPClient struct { + baseURL string + client *http.Client +} + +func NewVertexBatchHTTPClient(baseURL string, client *http.Client) *VertexBatchHTTPClient { + if client == nil { + client = batchImageDefaultHTTPClient() + } + return &VertexBatchHTTPClient{baseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"), client: client} +} + +func (c *VertexBatchHTTPClient) CreateBatchPredictionJob(ctx context.Context, accessToken string, req VertexCreateBatchPredictionJobRequest) (*VertexBatchPredictionJob, error) { + endpoint, err := BuildVertexBatchPredictionJobsEndpoint(c.baseURL, req.ProjectID, req.Location) + if err != nil { + return nil, err + } + payload, err := json.Marshal(req) + if err != nil { + return nil, err + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Authorization", "Bearer "+accessToken) + return doVertexJSON[VertexBatchPredictionJob](c.client, httpReq) +} + +func (c *VertexBatchHTTPClient) GetBatchPredictionJob(ctx context.Context, accessToken string, name string) (*VertexBatchPredictionJob, error) { + endpoint := c.vertexResourceURL(name) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + return doVertexJSON[VertexBatchPredictionJob](c.client, req) +} + +func (c *VertexBatchHTTPClient) CancelBatchPredictionJob(ctx context.Context, accessToken string, name string) error { + endpoint := c.vertexResourceURL(name) + ":cancel" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, nil) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + return doVertexNoBody(c.client, req) +} + +func (c *VertexBatchHTTPClient) vertexResourceURL(name string) string { + name = strings.TrimLeft(strings.TrimSpace(name), "/") + if c.baseURL != "" { + return c.baseURL + "/v1/" + name + } + return "https://aiplatform.googleapis.com/v1/" + name +} + +type VertexGCSObjectStore struct { + baseURL string + client *http.Client +} + +func NewVertexGCSObjectStore(baseURL string, client *http.Client) *VertexGCSObjectStore { + if client == nil { + client = batchImageDefaultHTTPClient() + } + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" { + baseURL = "https://storage.googleapis.com" + } + return &VertexGCSObjectStore{baseURL: baseURL, client: client} +} + +func (s *VertexGCSObjectStore) UploadJSONL(ctx context.Context, accessToken string, uri string, r io.Reader) error { + bucket, object, err := parseGCSURI(uri) + if err != nil { + return err + } + endpoint := fmt.Sprintf("%s/upload/storage/v1/b/%s/o?uploadType=media&name=%s", s.baseURL, url.PathEscape(bucket), url.QueryEscape(object)) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, r) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Content-Type", "application/jsonl") + return doVertexNoBody(s.client, req) +} + +func (s *VertexGCSObjectStore) ListJSONLObjects(ctx context.Context, accessToken string, prefixURI string) ([]string, error) { + return s.listObjects(ctx, accessToken, prefixURI, true) +} + +func (s *VertexGCSObjectStore) listObjects(ctx context.Context, accessToken string, prefixURI string, jsonlOnly bool) ([]string, error) { + bucket, prefix, err := parseGCSURI(prefixURI) + if err != nil { + return nil, err + } + var objects []string + pageToken := "" + for { + endpoint := fmt.Sprintf("%s/storage/v1/b/%s/o?prefix=%s", s.baseURL, url.PathEscape(bucket), url.QueryEscape(prefix)) + if pageToken != "" { + endpoint += "&pageToken=" + url.QueryEscape(pageToken) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + var page struct { + Items []struct { + Name string `json:"name"` + } `json:"items"` + NextPageToken string `json:"nextPageToken"` + } + if err := doVertexDecodeJSON(s.client, req, &page); err != nil { + return nil, err + } + for _, item := range page.Items { + if !jsonlOnly || strings.HasSuffix(item.Name, ".jsonl") { + objects = append(objects, "gs://"+bucket+"/"+item.Name) + } + } + if page.NextPageToken == "" { + return objects, nil + } + pageToken = page.NextPageToken + } +} + +func (s *VertexGCSObjectStore) OpenObject(ctx context.Context, accessToken string, uri string) (io.ReadCloser, string, error) { + bucket, object, err := parseGCSURI(uri) + if err != nil { + return nil, "", err + } + endpoint := fmt.Sprintf("%s/storage/v1/b/%s/o/%s?alt=media", s.baseURL, url.PathEscape(bucket), url.PathEscape(object)) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return nil, "", err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + resp, err := s.client.Do(req) + if err != nil { + return nil, "", err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + defer func() { _ = resp.Body.Close() }() + return nil, "", readVertexAPIError(resp) + } + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + contentType = "application/jsonl" + } + return resp.Body, contentType, nil +} + +func (s *VertexGCSObjectStore) DeleteObject(ctx context.Context, accessToken string, uri string) error { + bucket, object, err := parseGCSURI(uri) + if err != nil { + return err + } + endpoint := fmt.Sprintf("%s/storage/v1/b/%s/o/%s", s.baseURL, url.PathEscape(bucket), url.PathEscape(object)) + req, err := http.NewRequestWithContext(ctx, http.MethodDelete, endpoint, nil) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + return doVertexNoBody(s.client, req) +} + +func (s *VertexGCSObjectStore) DeletePrefix(ctx context.Context, accessToken string, prefixURI string) error { + objects, err := s.listObjects(ctx, accessToken, prefixURI, false) + if err != nil { + return err + } + for _, object := range objects { + if err := s.DeleteObject(ctx, accessToken, object); err != nil { + return err + } + } + return nil +} + +func parseGCSURI(uri string) (bucket, object string, err error) { + uri = strings.TrimSpace(uri) + if !strings.HasPrefix(uri, "gs://") { + return "", "", fmt.Errorf("invalid gcs uri") + } + rest := strings.TrimPrefix(uri, "gs://") + parts := strings.SplitN(rest, "/", 2) + if len(parts) != 2 || strings.TrimSpace(parts[0]) == "" || strings.TrimSpace(parts[1]) == "" { + return "", "", fmt.Errorf("invalid gcs uri") + } + return parts[0], parts[1], nil +} + +type VertexAPIError struct { + StatusCode int + Code string + Message string +} + +func (e *VertexAPIError) Error() string { + if e == nil { + return "" + } + if e.Code != "" { + return fmt.Sprintf("vertex api error: status=%d code=%s message=%s", e.StatusCode, e.Code, e.Message) + } + return fmt.Sprintf("vertex api error: status=%d message=%s", e.StatusCode, e.Message) +} + +func doVertexJSON[T any](client *http.Client, req *http.Request) (*T, error) { + var out T + if err := doVertexDecodeJSON(client, req, &out); err != nil { + return nil, err + } + return &out, nil +} + +func doVertexDecodeJSON(client *http.Client, req *http.Request, out any) error { + resp, err := client.Do(req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return readVertexAPIError(resp) + } + return json.NewDecoder(resp.Body).Decode(out) +} + +func doVertexNoBody(client *http.Client, req *http.Request) error { + resp, err := client.Do(req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return readVertexAPIError(resp) + } + return nil +} + +func readVertexAPIError(resp *http.Response) error { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 8192)) + message := string(body) + code := "" + var parsed struct { + Error struct { + Code any `json:"code"` + Message string `json:"message"` + Status string `json:"status"` + } `json:"error"` + } + if err := json.Unmarshal(body, &parsed); err == nil && parsed.Error.Message != "" { + message = parsed.Error.Message + code = parsed.Error.Status + } + return &VertexAPIError{StatusCode: resp.StatusCode, Code: code, Message: message} +} + +var _ BatchImageProvider = (*VertexBatchImageProvider)(nil) +var _ VertexBatchClient = (*VertexBatchHTTPClient)(nil) +var _ VertexBatchObjectStore = (*VertexGCSObjectStore)(nil) diff --git a/backend/internal/service/batch_image_provider_vertex_test.go b/backend/internal/service/batch_image_provider_vertex_test.go new file mode 100644 index 0000000000..4ad68f64a3 --- /dev/null +++ b/backend/internal/service/batch_image_provider_vertex_test.go @@ -0,0 +1,438 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/stretchr/testify/require" +) + +func TestBatchImageProviderRegistry_ReturnsVertex(t *testing.T) { + registry := NewDefaultBatchImageProviderRegistry() + provider, ok := registry.Get(BatchImageProviderVertex) + require.True(t, ok) + require.Equal(t, BatchImageProviderVertex, provider.Name()) +} + +func TestVertexProvider_SupportsOnlyGeminiServiceAccount(t *testing.T) { + provider := newTestVertexProvider(&fakeVertexBatchClient{}, &fakeVertexObjectStore{}) + + require.True(t, provider.SupportsAccount(vertexServiceAccount())) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk"}})) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"access_token": "tok"}})) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformAnthropic, Type: AccountTypeServiceAccount, Credentials: vertexServiceAccount().Credentials})) + require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{}})) +} + +func TestVertexProvider_MissingServiceAccountRejected(t *testing.T) { + provider := newTestVertexProvider(&fakeVertexBatchClient{}, &fakeVertexObjectStore{}) + _, err := provider.Submit(context.Background(), nil, &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Credentials: map[string]any{}}, validVertexBatchInput()) + require.ErrorIs(t, err, ErrBatchImageProviderMissingServiceAccount) +} + +func TestVertexProvider_MissingManagedGCSBucketRejected(t *testing.T) { + provider := NewVertexBatchImageProvider(VertexBatchImageProviderOptions{ProjectID: "proj", Environment: "test"}, &fakeVertexBatchClient{}, &fakeVertexObjectStore{}, &fakeGeminiTokenCache{token: "token"}) + _, err := provider.Submit(context.Background(), nil, vertexServiceAccount(), validVertexBatchInput()) + require.Error(t, err) + require.Equal(t, "VERTEX_MANAGED_GCS_BUCKET_MISSING", infraerrors.Reason(err)) +} + +func TestBuildVertexBatchJSONL_WritesValidLinesAndPreservesCustomID(t *testing.T) { + input := validVertexBatchInput() + input.Items = append(input.Items, BatchImageInputItem{CustomID: "cover_002", Prompt: "Second prompt"}) + + jsonl, err := BuildVertexBatchJSONL(input) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(jsonl)), "\n") + require.Len(t, lines, 2) + requireVertexJSONLLine(t, lines[0], "cover_001", "A clean product hero image") + requireVertexJSONLLine(t, lines[1], "cover_002", "Second prompt") +} + +func TestBuildVertexBatchJSONL_RejectsDuplicateCustomIDs(t *testing.T) { + input := validVertexBatchInput() + input.Items = append(input.Items, BatchImageInputItem{CustomID: "cover_001", Prompt: "Duplicate"}) + _, err := BuildVertexBatchJSONL(input) + require.ErrorIs(t, err, ErrBatchImageProviderInvalidInput) +} + +func TestBuildVertexBatchJSONL_RejectsEmptyPrompt(t *testing.T) { + input := validVertexBatchInput() + input.Items[0].Prompt = " " + _, err := BuildVertexBatchJSONL(input) + require.ErrorIs(t, err, ErrBatchImageProviderInvalidInput) +} + +func TestBuildVertexBatchJSONL_WritesReferenceImages(t *testing.T) { + input := validVertexBatchInput() + input.Items[0].ReferenceImages = []BatchImageReference{ + {MimeType: "image/png", Data: []byte("png-bytes")}, + {MimeType: "image/jpeg", FileURI: "gs://bucket/refs/style.jpg"}, + } + + jsonl, err := BuildVertexBatchJSONL(input) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(jsonl)), "\n") + require.Len(t, lines, 1) + + var got map[string]any + require.NoError(t, json.Unmarshal([]byte(lines[0]), &got)) + request := got["request"].(map[string]any) + contents := request["contents"].([]any) + parts := contents[0].(map[string]any)["parts"].([]any) + require.Len(t, parts, 3) + require.Equal(t, "A clean product hero image", parts[0].(map[string]any)["text"]) + inlineData := parts[1].(map[string]any)["inlineData"].(map[string]any) + require.Equal(t, "image/png", inlineData["mimeType"]) + require.Equal(t, "cG5nLWJ5dGVz", inlineData["data"]) + fileData := parts[2].(map[string]any)["fileData"].(map[string]any) + require.Equal(t, "image/jpeg", fileData["mimeType"]) + require.Equal(t, "gs://bucket/refs/style.jpg", fileData["fileUri"]) +} + +func TestNormalizeVertexBatchModelPath(t *testing.T) { + require.Equal(t, "publishers/google/models/gemini-3.1-flash-image", NormalizeVertexBatchModelPath("gemini-3.1-flash-image")) + require.Equal(t, "publishers/google/models/gemini-2.5-flash-image", NormalizeVertexBatchModelPath("publishers/google/models/gemini-2.5-flash-image")) + require.Equal(t, "projects/p/locations/global/models/m", NormalizeVertexBatchModelPath("projects/p/locations/global/models/m")) +} + +func TestBuildVertexBatchPredictionJobsEndpoint(t *testing.T) { + global, err := BuildVertexBatchPredictionJobsEndpoint("", "my-project", "global") + require.NoError(t, err) + require.Equal(t, "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/batchPredictionJobs", global) + + regional, err := BuildVertexBatchPredictionJobsEndpoint("", "my-project", "asia-northeast1") + require.NoError(t, err) + require.Equal(t, "https://asia-northeast1-aiplatform.googleapis.com/v1/projects/my-project/locations/asia-northeast1/batchPredictionJobs", regional) +} + +func TestVertexProvider_SubmitUploadsJSONLAndCreatesBatchPredictionJob(t *testing.T) { + vertexClient := &fakeVertexBatchClient{created: &VertexBatchPredictionJob{Name: "projects/proj/locations/global/batchPredictionJobs/job-1", State: "JOB_STATE_PENDING"}} + store := &fakeVertexObjectStore{} + provider := newTestVertexProvider(vertexClient, store) + + got, err := provider.Submit(context.Background(), &BatchImageJob{BatchID: "imgbatch_abc123", Model: "gemini-3.1-flash-image"}, vertexServiceAccount(), validVertexBatchInput()) + require.NoError(t, err) + + require.Equal(t, "gs://managed-bucket/batch-image/test/imgbatch_abc123/input/requests.jsonl", store.uploadURI) + require.Equal(t, "projects/proj/locations/global/batchPredictionJobs/job-1", got.ProviderJobName) + require.Equal(t, store.uploadURI, got.ProviderInputRef) + require.Equal(t, "gs://managed-bucket/batch-image/test/imgbatch_abc123/output/", got.ProviderOutputRef) + require.Equal(t, "jsonl", vertexClient.createdReq.InputConfig.InstancesFormat) + require.Equal(t, "jsonl", vertexClient.createdReq.OutputConfig.PredictionsFormat) + require.Equal(t, got.ProviderOutputRef, vertexClient.createdReq.OutputConfig.GCSDestination.OutputURIPrefix) + require.Equal(t, "key", vertexClient.createdReq.InstanceConfig.KeyField) + require.NotContains(t, string(vertexClient.createdPayloadForAssert(t)), "serviceAccount") + require.NotContains(t, string(vertexClient.createdPayloadForAssert(t)), "encryptionSpec") + require.NotContains(t, got.ProviderInputRef+got.ProviderOutputRef+got.ProviderJobName, "A clean product hero image") + require.NotContains(t, string(store.uploadedJSONL), "private_key") +} + +func TestVertexProvider_GetMapsStates(t *testing.T) { + tests := []struct { + name string + state string + err *VertexBatchJobError + wantState BatchProviderInternalState + wantDone bool + wantCode string + }{ + {name: "pending", state: "JOB_STATE_PENDING", wantState: BatchProviderStateQueued}, + {name: "queued", state: "JOB_STATE_QUEUED", wantState: BatchProviderStateQueued}, + {name: "running", state: "JOB_STATE_RUNNING", wantState: BatchProviderStateRunning}, + {name: "succeeded", state: "JOB_STATE_SUCCEEDED", wantState: BatchProviderStateSucceeded, wantDone: true}, + {name: "failed", state: "JOB_STATE_FAILED", err: &VertexBatchJobError{Status: "INVALID_ARGUMENT", Message: "bad request"}, wantState: BatchProviderStateFailed, wantDone: true, wantCode: "INVALID_ARGUMENT"}, + {name: "cancelled", state: "JOB_STATE_CANCELLED", wantState: BatchProviderStateCancelled, wantDone: true, wantCode: "VERTEX_BATCH_CANCELLED"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + output := "gs://managed-bucket/batch-image/test/imgbatch_abc123/output/" + provider := newTestVertexProvider(&fakeVertexBatchClient{got: &VertexBatchPredictionJob{ + Name: "projects/proj/locations/global/batchPredictionJobs/job-1", + State: tt.state, + Error: tt.err, + OutputConfig: VertexBatchOutputConfig{GCSDestination: VertexBatchGCSDestination{OutputURIPrefix: output}}, + }}, &fakeVertexObjectStore{}) + got, err := provider.Get(context.Background(), vertexJobWithName("projects/proj/locations/global/batchPredictionJobs/job-1"), vertexServiceAccount()) + require.NoError(t, err) + require.Equal(t, tt.wantState, got.InternalState) + require.Equal(t, tt.wantDone, got.Done) + require.Equal(t, output, got.ProviderOutputRef) + require.Equal(t, tt.wantCode, got.ErrorCode) + }) + } +} + +func TestVertexProvider_OpenResultReturnsCombinedJSONLStream(t *testing.T) { + output := "gs://managed-bucket/batch-image/test/imgbatch_abc123/output/" + store := &fakeVertexObjectStore{ + listed: []string{ + output + "predictions_2.jsonl", + output + "predictions_1.jsonl", + }, + objects: map[string]string{ + output + "predictions_1.jsonl": `{"key":"1"}` + "\n", + output + "predictions_2.jsonl": `{"key":"2"}` + "\n", + }, + } + provider := newTestVertexProvider(&fakeVertexBatchClient{}, store) + r, contentType, err := provider.OpenResult(context.Background(), &BatchImageJob{ProviderOutputRef: &output}, vertexServiceAccount()) + require.NoError(t, err) + defer r.Close() + + body, err := io.ReadAll(r) + require.NoError(t, err) + require.Equal(t, "application/jsonl", contentType) + require.Equal(t, "{\"key\":\"1\"}\n\n{\"key\":\"2\"}\n", string(body)) +} + +func TestVertexProvider_OpenResultMissingObjectsReturnsTypedError(t *testing.T) { + output := "gs://managed-bucket/batch-image/test/imgbatch_abc123/output/" + provider := newTestVertexProvider(&fakeVertexBatchClient{}, &fakeVertexObjectStore{}) + _, _, err := provider.OpenResult(context.Background(), &BatchImageJob{ProviderOutputRef: &output}, vertexServiceAccount()) + require.Error(t, err) + require.Equal(t, "VERTEX_RESULT_OBJECTS_MISSING", infraerrors.Reason(err)) +} + +func TestVertexProvider_CancelCallsClient(t *testing.T) { + vertexClient := &fakeVertexBatchClient{} + provider := newTestVertexProvider(vertexClient, &fakeVertexObjectStore{}) + + err := provider.Cancel(context.Background(), vertexJobWithName("projects/proj/locations/global/batchPredictionJobs/job-1"), vertexServiceAccount()) + require.NoError(t, err) + require.Equal(t, "projects/proj/locations/global/batchPredictionJobs/job-1", vertexClient.cancelledName) +} + +func TestVertexProvider_CleanupDeletesOnlyManagedPaths(t *testing.T) { + input := "gs://managed-bucket/batch-image/test/imgbatch_abc123/input/requests.jsonl" + output := "gs://managed-bucket/batch-image/test/imgbatch_abc123/output/" + store := &fakeVertexObjectStore{} + provider := newTestVertexProvider(&fakeVertexBatchClient{}, store) + + err := provider.Cleanup(context.Background(), &BatchImageJob{BatchID: "imgbatch_abc123", ProviderInputRef: &input, ProviderOutputRef: &output}, vertexServiceAccount(), CleanupTargetAll) + require.NoError(t, err) + require.Equal(t, []string{input}, store.deletedObjects) + require.Equal(t, []string{output}, store.deletedPrefixes) +} + +func TestVertexProvider_CleanupRejectsUnsafePath(t *testing.T) { + input := "gs://other-bucket/batch-image/test/imgbatch_abc123/input/requests.jsonl" + provider := newTestVertexProvider(&fakeVertexBatchClient{}, &fakeVertexObjectStore{}) + + err := provider.Cleanup(context.Background(), &BatchImageJob{BatchID: "imgbatch_abc123", ProviderInputRef: &input}, vertexServiceAccount(), CleanupTargetInput) + require.ErrorIs(t, err, ErrBatchImageProviderUnsafeCleanupPath) +} + +func TestVertexProvider_ErrorsDoNotExposeServiceAccountSecrets(t *testing.T) { + privateKey := "-----BEGIN PRIVATE KEY-----secret-----END PRIVATE KEY-----" + account := vertexServiceAccount() + account.Credentials["service_account_json"] = map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": privateKey, + "client_email": "svc@proj.iam.gserviceaccount.com", + } + provider := newTestVertexProvider(&fakeVertexBatchClient{createErr: &VertexAPIError{StatusCode: 403, Message: "do not expose " + privateKey}}, &fakeVertexObjectStore{}) + + _, err := provider.Submit(context.Background(), nil, account, validVertexBatchInput()) + require.Error(t, err) + require.Equal(t, "VERTEX_PERMISSION_DENIED", infraerrors.Reason(err)) + require.NotContains(t, err.Error(), privateKey) + require.NotContains(t, err.Error(), "svc@proj") +} + +func TestVertexProvider_MetadataDoesNotStoreImageBytesOrBase64(t *testing.T) { + vertexClient := &fakeVertexBatchClient{created: &VertexBatchPredictionJob{Name: "projects/proj/locations/global/batchPredictionJobs/job-1", State: "JOB_STATE_PENDING"}} + provider := newTestVertexProvider(vertexClient, &fakeVertexObjectStore{}) + + got, err := provider.Submit(context.Background(), nil, vertexServiceAccount(), validVertexBatchInput()) + require.NoError(t, err) + metadata := got.ProviderJobName + got.ProviderInputRef + got.ProviderOutputRef + require.NotContains(t, metadata, "iVBOR") + require.NotContains(t, metadata, "base64") + require.NotContains(t, metadata, "A clean product hero image") +} + +func validVertexBatchInput() BatchImageInput { + return BatchImageInput{ + BatchID: "imgbatch_abc123", + Model: "gemini-3.1-flash-image", + DisplayName: "test vertex batch", + Items: []BatchImageInputItem{{ + CustomID: "cover_001", + Prompt: "A clean product hero image", + }}, + } +} + +func requireVertexJSONLLine(t *testing.T, line, wantKey, wantPrompt string) { + t.Helper() + var got map[string]any + require.NoError(t, json.Unmarshal([]byte(line), &got)) + require.Equal(t, wantKey, got["key"]) + request := got["request"].(map[string]any) + contents := request["contents"].([]any) + require.Equal(t, "user", contents[0].(map[string]any)["role"]) + parts := contents[0].(map[string]any)["parts"].([]any) + require.Equal(t, wantPrompt, parts[0].(map[string]any)["text"]) + config := request["generationConfig"].(map[string]any) + require.Equal(t, []any{"TEXT", "IMAGE"}, config["responseModalities"]) +} + +func newTestVertexProvider(client *fakeVertexBatchClient, store *fakeVertexObjectStore) *VertexBatchImageProvider { + return NewVertexBatchImageProvider(VertexBatchImageProviderOptions{ + ProjectID: "proj", + Location: "global", + ManagedGCSBucket: "managed-bucket", + ManagedGCSPrefix: "batch-image/{env}/{batch_id}", + Environment: "test", + }, client, store, &fakeGeminiTokenCache{token: "ya29.test-token"}) +} + +func vertexServiceAccount() *Account { + return &Account{ + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{ + "service_account_json": map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "client_email": "svc@proj.iam.gserviceaccount.com", + }, + }, + } +} + +func vertexJobWithName(name string) *BatchImageJob { + return &BatchImageJob{ProviderJobName: &name} +} + +type fakeVertexBatchClient struct { + created *VertexBatchPredictionJob + got *VertexBatchPredictionJob + createErr error + getErr error + cancelErr error + createdReq VertexCreateBatchPredictionJobRequest + cancelledName string +} + +func (f *fakeVertexBatchClient) CreateBatchPredictionJob(_ context.Context, accessToken string, req VertexCreateBatchPredictionJobRequest) (*VertexBatchPredictionJob, error) { + if strings.TrimSpace(accessToken) == "" { + return nil, errors.New("missing token") + } + f.createdReq = req + if f.createErr != nil { + return nil, f.createErr + } + if f.created != nil { + return f.created, nil + } + return &VertexBatchPredictionJob{Name: "projects/proj/locations/global/batchPredictionJobs/job-1", State: "JOB_STATE_PENDING"}, nil +} + +func (f *fakeVertexBatchClient) GetBatchPredictionJob(_ context.Context, _ string, _ string) (*VertexBatchPredictionJob, error) { + if f.getErr != nil { + return nil, f.getErr + } + return f.got, nil +} + +func (f *fakeVertexBatchClient) CancelBatchPredictionJob(_ context.Context, _ string, name string) error { + f.cancelledName = name + return f.cancelErr +} + +func (f *fakeVertexBatchClient) createdPayloadForAssert(t *testing.T) []byte { + t.Helper() + b, err := json.Marshal(f.createdReq) + require.NoError(t, err) + return b +} + +type fakeVertexObjectStore struct { + uploadURI string + uploadedJSONL []byte + uploadErr error + listed []string + objects map[string]string + listErr error + openErr error + deleteErr error + deletedObjects []string + deletedPrefixes []string +} + +func (f *fakeVertexObjectStore) UploadJSONL(_ context.Context, _ string, uri string, r io.Reader) error { + f.uploadURI = uri + f.uploadedJSONL, _ = io.ReadAll(r) + return f.uploadErr +} + +func (f *fakeVertexObjectStore) ListJSONLObjects(_ context.Context, _ string, _ string) ([]string, error) { + if f.listErr != nil { + return nil, f.listErr + } + out := make([]string, 0, len(f.listed)) + for _, item := range f.listed { + if strings.HasSuffix(item, ".jsonl") { + out = append(out, item) + } + } + return out, nil +} + +func (f *fakeVertexObjectStore) OpenObject(_ context.Context, _ string, uri string) (io.ReadCloser, string, error) { + if f.openErr != nil { + return nil, "", f.openErr + } + return io.NopCloser(bytes.NewBufferString(f.objects[uri])), "application/jsonl", nil +} + +func (f *fakeVertexObjectStore) DeleteObject(_ context.Context, _ string, uri string) error { + f.deletedObjects = append(f.deletedObjects, uri) + return f.deleteErr +} + +func (f *fakeVertexObjectStore) DeletePrefix(_ context.Context, _ string, uri string) error { + f.deletedPrefixes = append(f.deletedPrefixes, uri) + return f.deleteErr +} + +type fakeGeminiTokenCache struct { + token string +} + +func (f *fakeGeminiTokenCache) GetAccessToken(context.Context, string) (string, error) { + if strings.TrimSpace(f.token) == "" { + return "", errors.New("missing token") + } + return f.token, nil +} + +func (f *fakeGeminiTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (f *fakeGeminiTokenCache) DeleteAccessToken(context.Context, string) error { + return nil +} + +func (f *fakeGeminiTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (f *fakeGeminiTokenCache) ReleaseRefreshLock(context.Context, string) error { + return nil +} diff --git a/backend/internal/service/batch_image_public.go b/backend/internal/service/batch_image_public.go new file mode 100644 index 0000000000..2960d4cd9f --- /dev/null +++ b/backend/internal/service/batch_image_public.go @@ -0,0 +1,1438 @@ +package service + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "sort" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + defaultBatchImageMaxItems = 200 + defaultBatchImageMaxOutputImages = 200 + defaultBatchImageMaxOutputCount = 4 + defaultBatchImageMaxPromptChars = 8000 + defaultBatchImageResponseMime = "image/png" + defaultBatchImageImageSize = "1K" + defaultBatchImageDiscountMultiplier = 0.5 + defaultBatchImageHoldMultiplier = 0.6 + maxBatchImagePublicErrorChars = 500 + maxBatchImageReferenceImageBytes = 10 * 1024 * 1024 + defaultBatchImageMaxReferenceImages = 1000 + defaultBatchImageMaxReferenceBytes = 128 * 1024 * 1024 +) + +type BatchImageAccountSelectionRepository interface { + GetByID(ctx context.Context, id int64) (*Account, error) + ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) + ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) +} + +type BatchImageGroupPricingRepository interface { + GetByIDLite(ctx context.Context, id int64) (*Group, error) +} + +type BatchImageUserGroupRateRepository interface { + GetByUserAndGroup(ctx context.Context, userID, groupID int64) (*float64, error) +} + +type BatchImageSubmitRequest struct { + Model string `json:"model"` + TaskName string `json:"task_name"` + ParentBatchID string `json:"parent_batch_id"` + Provider string `json:"provider"` + Items []BatchImageSubmitItem `json:"items"` + ResponseMimeType string `json:"response_mime_type"` + AspectRatio string `json:"aspect_ratio"` + ImageSize string `json:"image_size"` + Metadata map[string]string `json:"metadata"` +} + +type BatchImageSubmitItem struct { + CustomID string `json:"custom_id"` + Prompt string `json:"prompt"` + OutputCount int `json:"output_count,omitempty"` + ReferenceImages []BatchImageReferenceInput `json:"reference_images,omitempty"` +} + +type BatchImageReferenceInput struct { + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + MimeType string `json:"mime_type"` + Data []byte `json:"data,omitempty"` + FileURI string `json:"file_uri,omitempty"` +} + +type BatchImageOwner struct { + UserID int64 + APIKeyID int64 + GroupID *int64 +} + +type BatchImagePublicService struct { + Repo BatchImageRepository + AccountRepo BatchImageAccountSelectionRepository + GroupRepo BatchImageGroupPricingRepository + UserGroupRateRepo BatchImageUserGroupRateRepository + Queue BatchImageQueue + ProviderRegistry *BatchImageProviderRegistry + Pricing BatchImagePricingResolver + BillingRepo UsageBillingRepository + AuthCache APIKeyAuthCacheInvalidator + Config *config.Config +} + +type BatchImagePricingSnapshot struct { + BaseUnitPrice float64 + GroupRateMultiplier float64 + AccountRateMultiplier float64 + BatchDiscountMultiplier float64 + HoldMultiplier float64 + BillableUnitPrice float64 + HoldUnitPrice float64 + EstimatedCost float64 + HoldAmount float64 +} + +type BatchImagePublicBatch struct { + ID string `json:"id"` + Object string `json:"object"` + TaskName string `json:"task_name"` + ParentBatchID *string `json:"parent_batch_id,omitempty"` + Status string `json:"status"` + Model string `json:"model"` + Provider string `json:"provider"` + ItemCount int `json:"item_count"` + SuccessCount int `json:"success_count"` + FailCount int `json:"fail_count"` + EstimatedCost float64 `json:"estimated_cost"` + HoldAmount float64 `json:"hold_amount"` + ActualCost *float64 `json:"actual_cost"` + CreatedAt int64 `json:"created_at"` + SubmittedAt *int64 `json:"submitted_at"` + SettledAt *int64 `json:"settled_at"` + DownloadedAt *int64 `json:"downloaded_at,omitempty"` + OutputDeletedAt *int64 `json:"output_deleted_at,omitempty"` +} + +type BatchImagePublicItem struct { + CustomID string `json:"custom_id"` + Status string `json:"status"` + PromptPreview *string `json:"prompt_preview,omitempty"` + MimeType *string `json:"mime_type"` + FileExtension *string `json:"file_extension"` + ImageCount int `json:"image_count"` + Error *BatchImagePublicError `json:"error"` +} + +type BatchImagePublicError struct { + Code string `json:"code"` + Message string `json:"message"` + Source string `json:"source,omitempty"` +} + +type BatchImagePublicItemsResponse struct { + Object string `json:"object"` + Data []BatchImagePublicItem `json:"data"` + HasMore bool `json:"has_more"` +} + +type BatchImagePublicListResponse struct { + Object string `json:"object"` + Data []*BatchImagePublicBatch `json:"data"` + HasMore bool `json:"has_more"` +} + +type BatchImagePublicModel struct { + ID string `json:"id"` + Object string `json:"object"` + Provider string `json:"provider"` +} + +type BatchImagePublicModelsResponse struct { + Object string `json:"object"` + Data []BatchImagePublicModel `json:"data"` +} + +type BatchImageJobsQuery struct { + Status string + TaskName string + Downloaded string + From string + To string + Limit int + Cursor string +} + +type BatchImageItemsQuery struct { + Status string + Limit int + Cursor string +} + +func NewBatchImagePublicService(repo BatchImageRepository, accountRepo AccountRepository, groupRepo GroupRepository, userGroupRateRepo UserGroupRateRepository, queue BatchImageQueue, pricing *BatchImageModelPricingResolver, billingRepo UsageBillingRepository, authCache APIKeyAuthCacheInvalidator, cfg *config.Config) *BatchImagePublicService { + return &BatchImagePublicService{ + Repo: repo, + AccountRepo: accountRepo, + GroupRepo: groupRepo, + UserGroupRateRepo: userGroupRateRepo, + Queue: queue, + ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg), + Pricing: pricing, + BillingRepo: billingRepo, + AuthCache: authCache, + Config: cfg, + } +} + +func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOwner, req BatchImageSubmitRequest, idempotencyKey string) (*BatchImagePublicBatch, error) { + if !s.enabled() { + return nil, ErrBatchImageDisabled + } + normalized, err := s.validateSubmitRequest(req) + if err != nil { + return nil, err + } + // 与 ListModels 使用同一鉴权谓词(AllowBatchImageGeneration + Platform==Gemini), + // 避免两个入口校验口径不一致留下防御纵深缺口。 + if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil { + return nil, err + } + requestHash := HashBatchImageSubmitRequest(normalized) + idempotencyKey = strings.TrimSpace(idempotencyKey) + if idempotencyKey != "" { + existing, err := s.Repo.GetBatchImageJobByIdempotencyKey(ctx, owner.UserID, owner.APIKeyID, idempotencyKey) + if err == nil { + if batchImageDerefString(existing.RequestHash) != requestHash { + return nil, ErrBatchImageIdempotencyConflict + } + if existing.Status == BatchImageJobStatusSubmitted && s.Queue != nil { + if enqueueErr := s.Queue.Enqueue(ctx, existing.BatchID); enqueueErr != nil && !errors.Is(enqueueErr, ErrBatchImageAlreadyQueued) { + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, existing.BatchID, "QUEUE_FAILED", sanitizeBatchImagePublicMessage(enqueueErr.Error()), false) + return nil, ErrBatchImageQueueFailed + } + } + return BatchImageJobToPublic(existing), nil + } + if !errors.Is(err, ErrBatchImageJobNotFound) { + return nil, err + } + } + + provider, account, err := s.selectProviderAndAccount(ctx, owner, normalized.Provider, normalized.Model) + if err != nil { + return nil, err + } + pricingSnapshot, err := s.resolvePricingSnapshot(ctx, owner, normalized, provider.Name(), account) + if err != nil { + return nil, err + } + parentBatchID := batchImageOptionalStringPtr(normalized.ParentBatchID) + if parentBatchID != nil { + parent, parentErr := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, *parentBatchID) + if parentErr != nil { + return nil, parentErr + } + if parent.ParentBatchID != nil && strings.TrimSpace(*parent.ParentBatchID) != "" { + parentBatchID = batchImageOptionalStringPtr(*parent.ParentBatchID) + } + } + batchID, err := NewBatchImageID() + if err != nil { + return nil, err + } + apiKeyID := owner.APIKeyID + accountID := account.ID + holdID := BatchImageHoldRequestID(batchID) + holdAmount := pricingSnapshot.HoldAmount + job, err := s.Repo.CreateBatchImageJob(ctx, CreateBatchImageJobParams{ + BatchID: batchID, + UserID: owner.UserID, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: provider.Name(), + Model: normalized.Model, + TaskName: normalized.TaskName, + ParentBatchID: parentBatchID, + Status: BatchImageJobStatusCreated, + ItemCount: len(normalized.Items), + EstimatedCost: pricingSnapshot.EstimatedCost, + HoldAmount: &holdAmount, + BaseUnitPrice: pricingSnapshot.BaseUnitPrice, + GroupRateMultiplier: pricingSnapshot.GroupRateMultiplier, + AccountRateMultiplier: pricingSnapshot.AccountRateMultiplier, + BatchDiscountMultiplier: pricingSnapshot.BatchDiscountMultiplier, + HoldMultiplier: pricingSnapshot.HoldMultiplier, + BillableUnitPrice: pricingSnapshot.BillableUnitPrice, + HoldUnitPrice: pricingSnapshot.HoldUnitPrice, + PricingSnapshotVersion: 1, + Currency: "USD", + HoldID: &holdID, + IdempotencyKey: batchImageOptionalStringPtr(idempotencyKey), + RequestHash: batchImageStringPtr(requestHash), + }) + if err != nil { + return nil, err + } + if err := reserveBatchImageBalanceHold(ctx, s.BillingRepo, job, requestHash); err != nil { + code := "BILLING_HOLD_FAILED" + if errors.Is(err, ErrBatchImageInsufficientBalance) { + code = "INSUFFICIENT_BALANCE" + } + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, code, sanitizeBatchImagePublicMessage(err.Error()), true) + s.hidePreUpstreamSubmitFailure(ctx, owner, job) + return nil, err + } + s.invalidateAuthCache(ctx, owner.UserID) + if err := s.createPendingItems(ctx, job.BatchID, requestHash, normalized.Items); err != nil { + if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil { + return nil, releaseErr + } + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "ITEM_CREATE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true) + s.hidePreUpstreamSubmitFailure(ctx, owner, job) + return nil, ErrBatchImageQueueFailed + } + + input := BatchImageInput{ + BatchID: job.BatchID, + Model: normalized.Model, + DisplayName: job.BatchID, + ResponseMimeType: normalized.ResponseMimeType, + AspectRatio: normalized.AspectRatio, + ImageSize: normalized.ImageSize, + Metadata: normalized.Metadata, + Items: make([]BatchImageInputItem, 0, len(normalized.Items)), + } + for _, item := range normalized.Items { + refs := make([]BatchImageReference, 0, len(item.ReferenceImages)) + for _, ref := range item.ReferenceImages { + refs = append(refs, BatchImageReference(ref)) + } + input.Items = append(input.Items, BatchImageInputItem{ + CustomID: item.CustomID, + Prompt: item.Prompt, + ReferenceImages: refs, + }) + } + + // 上游提交(上传参考图 + 创建批任务)可能长达数分钟且不刷新 updated_at, + // 会被 stale 恢复扫描误判为滞留并退款。提交前转入 uploading 刷新时间戳, + // 提交期间用心跳持续续期。 + if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusUploading, BatchImageTransitionOptions{ + EventType: "upload_started", + EventPayload: map[string]any{"batch_id": job.BatchID}, + }); err != nil { + if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil { + return nil, releaseErr + } + // 并发 Cancel 等导致的非法转换:job 已处于终态,不再覆盖其状态。 + if !errors.Is(err, ErrBatchImageInvalidTransition) { + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "UPLOAD_TRANSITION_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true) + s.hidePreUpstreamSubmitFailure(ctx, owner, job) + } + return nil, err + } + job.Status = BatchImageJobStatusUploading + + hbCtx, hbCancel := context.WithCancel(ctx) + hbDone := make(chan struct{}) + go s.runSubmitHeartbeat(hbCtx, job.BatchID, hbDone) + providerJob, err := provider.Submit(ctx, job, account, input) + hbCancel() + <-hbDone + if err != nil { + if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil { + return nil, releaseErr + } + publicErr := batchImageProviderSubmitPublicError(err) + reason := batchImageProviderSubmitRecordCode(publicErr) + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, reason, sanitizeBatchImagePublicMessage(err.Error()), true) + s.hidePreUpstreamSubmitFailure(ctx, owner, job) + return nil, publicErr + } + if providerJob == nil || strings.TrimSpace(providerJob.ProviderJobName) == "" { + if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil { + return nil, releaseErr + } + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "PROVIDER_SUBMIT_FAILED", "provider job name missing", true) + s.hidePreUpstreamSubmitFailure(ctx, owner, job) + return nil, ErrBatchImageProviderSubmitFailed + } + + if err := s.Repo.UpdateBatchImageJobProviderSubmit(ctx, UpdateBatchImageJobProviderSubmitParams{ + BatchID: job.BatchID, + ProviderJobName: providerJob.ProviderJobName, + ProviderInputRef: providerJob.ProviderInputRef, + ProviderOutputRef: providerJob.ProviderOutputRef, + GCSInputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef), + GCSOutputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderOutputRef), + EventPayload: map[string]any{"provider": provider.Name()}, + }); err != nil { + // job 可能已被恢复扫描转 failed 并退款:上游批任务已创建成功, + // 必须尽力取消并清理输入,否则上游照常产生成本(孤儿任务)。 + s.abortOrphanProviderJob(ctx, provider, job, account, providerJob) + return nil, err + } + + if s.Queue != nil { + if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) { + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "QUEUE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), false) + return nil, ErrBatchImageQueueFailed + } + } + + created, err := s.Repo.GetBatchImageJobByBatchID(ctx, job.BatchID) + if err != nil { + return nil, err + } + return BatchImageJobToPublic(created), nil +} + +func (s *BatchImagePublicService) releaseFailedSubmitHold(ctx context.Context, job *BatchImageJob, requestHash string) error { + if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, requestHash); err != nil { + _ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "BILLING_RELEASE_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true) + s.enqueueBillingRetry(ctx, job.BatchID) + return ErrBatchImageBillingHoldFailed + } + s.invalidateAuthCache(ctx, job.UserID) + return nil +} + +// runSubmitHeartbeat 在 provider.Submit 期间周期性刷新 job 的 updated_at, +// 使 stale 恢复扫描能区分"仍在慢提交"与"进程死亡后的滞留"。 +func (s *BatchImagePublicService) runSubmitHeartbeat(ctx context.Context, batchID string, done chan<- struct{}) { + defer close(done) + interval := s.submitHeartbeatInterval() + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if err := s.Repo.TouchBatchImageJobSubmitting(ctx, batchID); err != nil && ctx.Err() == nil { + logger.L().Warn("batch_image.submit_heartbeat_failed", + zap.String("batch_id", batchID), + zap.Error(err), + ) + } + } + } +} + +func (s *BatchImagePublicService) submitHeartbeatInterval() time.Duration { + staleAfter := 10 * time.Minute + if s != nil && s.Config != nil && s.Config.BatchImage.StaleActiveAfterSeconds > 0 { + staleAfter = time.Duration(s.Config.BatchImage.StaleActiveAfterSeconds) * time.Second + } + interval := staleAfter / 3 + if interval < 15*time.Second { + interval = 15 * time.Second + } + return interval +} + +// abortOrphanProviderJob 在上游任务创建成功但本地状态推进失败时, +// 尽力取消上游批任务并清理已上传的输入文件,避免孤儿任务持续产生成本。 +func (s *BatchImagePublicService) abortOrphanProviderJob(ctx context.Context, provider BatchImageProvider, job *BatchImageJob, account *Account, providerJob *BatchProviderJob) { + if s == nil || provider == nil || job == nil || providerJob == nil { + return + } + orphan := *job + orphan.ProviderJobName = batchImageOptionalStringPtr(providerJob.ProviderJobName) + orphan.ProviderInputRef = batchImageOptionalStringPtr(providerJob.ProviderInputRef) + orphan.GCSInputURI = batchImageOptionalStringPtr(batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef)) + if err := provider.Cancel(ctx, &orphan, account); err != nil { + logger.L().Warn("batch_image.orphan_provider_job_cancel_failed", + zap.String("batch_id", job.BatchID), + zap.String("provider", provider.Name()), + zap.Error(err), + ) + } + if err := provider.Cleanup(ctx, &orphan, account, CleanupTargetInput); err != nil { + logger.L().Warn("batch_image.orphan_provider_job_cleanup_failed", + zap.String("batch_id", job.BatchID), + zap.String("provider", provider.Name()), + zap.Error(err), + ) + } + if err := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "provider_job_aborted_after_submit", map[string]any{ + "batch_id": job.BatchID, + "provider": provider.Name(), + }); err != nil { + logger.L().Warn("batch_image.orphan_provider_job_event_failed", + zap.String("batch_id", job.BatchID), + zap.Error(err), + ) + } +} + +func (s *BatchImagePublicService) createPendingItems(ctx context.Context, batchID, requestHash string, items []BatchImageSubmitItem) error { + if s == nil || s.Repo == nil || len(items) == 0 { + return nil + } + params := make([]CreateBatchImageItemParams, 0, len(items)) + for _, item := range items { + preview := truncateBatchImageMessage(item.Prompt, s.maxPromptChars()) + params = append(params, CreateBatchImageItemParams{ + JobID: batchID, + CustomID: item.CustomID, + Status: BatchImageItemStatusPending, + RequestHash: batchImageStringPtr(requestHash), + PromptPreview: batchImageStringPtr(preview), + ImageCount: 0, + }) + } + return s.Repo.BulkCreateBatchImageItems(ctx, params) +} + +func (s *BatchImagePublicService) enqueueBillingRetry(ctx context.Context, batchID string) { + if s == nil || s.Queue == nil { + return + } + if err := s.Queue.Enqueue(ctx, batchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) { + logger.L().Warn("batch_image.billing_retry_enqueue_failed", + zap.String("batch_id", batchID), + zap.Error(err), + ) + if eventErr := s.Repo.AppendBatchImageEvent(ctx, batchID, "billing_retry_enqueue_failed", map[string]any{ + "batch_id": batchID, + "error": sanitizeBatchImagePublicMessage(err.Error()), + }); eventErr != nil { + logger.L().Warn("batch_image.billing_retry_event_failed", + zap.String("batch_id", batchID), + zap.Error(eventErr), + ) + } + } +} + +func (s *BatchImagePublicService) hidePreUpstreamSubmitFailure(ctx context.Context, owner BatchImageOwner, job *BatchImageJob) { + if s == nil || s.Repo == nil || job == nil || job.ProviderJobName != nil { + return + } + if err := s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now()); err != nil { + logger.L().Warn("batch_image.hide_pre_upstream_failure_failed", + zap.String("batch_id", job.BatchID), + zap.Error(err), + ) + } +} + +func (s *BatchImagePublicService) Get(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) { + job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + return BatchImageJobToPublic(job), nil +} + +func (s *BatchImagePublicService) List(ctx context.Context, owner BatchImageOwner, query BatchImageJobsQuery) (*BatchImagePublicListResponse, error) { + filter := BatchImageJobFilter{Limit: query.Limit, Offset: parseBatchImageCursor(query.Cursor), ExcludeDeleted: true} + filter.TaskNameLike = strings.TrimSpace(query.TaskName) + switch strings.TrimSpace(query.Status) { + case "", "all": + case "queued": + filter.Status = BatchImageJobStatusSubmitted + case "processing_results": + filter.Status = BatchImageJobStatusIndexing + case "completed": + filter.Status = BatchImageJobStatusCompleted + case "failed": + filter.Status = BatchImageJobStatusFailed + case "cancelled": + filter.Status = BatchImageJobStatusCancelled + case "output_deleted": + filter.Status = BatchImageJobStatusOutputDeleted + default: + filter.Status = strings.TrimSpace(query.Status) + } + switch strings.TrimSpace(strings.ToLower(query.Downloaded)) { + case "", "all": + case "true", "1", "yes", "downloaded": + downloaded := true + filter.Downloaded = &downloaded + case "false", "0", "no", "not_downloaded": + downloaded := false + filter.Downloaded = &downloaded + default: + return nil, ErrBatchImageInvalidItems + } + if from := parseBatchImageListTime(query.From); from != nil { + filter.CreatedAfter = from + } + if to := parseBatchImageListTime(query.To); to != nil { + filter.CreatedBefore = to + } + if filter.Limit <= 0 || filter.Limit > 100 { + filter.Limit = 20 + } + jobs, err := s.Repo.ListBatchImageJobsForOwner(ctx, owner.UserID, owner.APIKeyID, filter) + if err != nil { + return nil, err + } + data := make([]*BatchImagePublicBatch, 0, len(jobs)) + for _, job := range jobs { + data = append(data, BatchImageJobToPublic(job)) + } + return &BatchImagePublicListResponse{ + Object: "list", + Data: data, + HasMore: len(data) == filter.Limit, + }, nil +} + +func (s *BatchImagePublicService) MarkDownloaded(ctx context.Context, owner BatchImageOwner, batchID string) error { + job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return err + } + return s.Repo.MarkBatchImageDownloaded(ctx, job.BatchID, time.Now()) +} + +func (s *BatchImagePublicService) DeleteRecord(ctx context.Context, owner BatchImageOwner, batchID string) error { + job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return err + } + if !isBatchImageProcessorDoneStatus(job.Status) { + return ErrBatchImageRecordDeleteNotReady + } + return s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now()) +} + +func (s *BatchImagePublicService) ListModels(ctx context.Context, owner BatchImageOwner) (*BatchImagePublicModelsResponse, error) { + if !s.enabled() { + return nil, ErrBatchImageDisabled + } + if s.Pricing == nil { + return nil, ErrBatchImageSettlementPricingMissing + } + if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil { + return nil, err + } + + modelsByProvider := make(map[string]map[string]struct{}) + for _, providerName := range batchImageProviderSelectionOrder("") { + provider, ok := s.ProviderRegistry.Get(providerName) + if !ok || provider == nil { + continue + } + accounts, err := s.listCandidateAccounts(ctx, owner.GroupID, batchImageProviderPlatform(providerName)) + if err != nil { + return nil, err + } + for i := range accounts { + account := accounts[i] + if !account.IsSchedulable() || !provider.SupportsAccount(&account) { + continue + } + for _, model := range batchImageModelsFromAccountMapping(&account) { + if _, err := s.Pricing.BatchImageUnitPrice(ctx, &BatchImageJob{Provider: providerName, Model: model}); err != nil { + continue + } + if !account.IsModelSupported(model) { + continue + } + if modelsByProvider[providerName] == nil { + modelsByProvider[providerName] = make(map[string]struct{}) + } + modelsByProvider[providerName][model] = struct{}{} + } + } + } + + out := make([]BatchImagePublicModel, 0) + for _, providerName := range batchImageProviderSelectionOrder("") { + models := make([]string, 0, len(modelsByProvider[providerName])) + for model := range modelsByProvider[providerName] { + models = append(models, model) + } + sort.Strings(models) + for _, model := range models { + out = append(out, BatchImagePublicModel{ + ID: model, + Object: "image.batch.model", + Provider: providerName, + }) + } + } + return &BatchImagePublicModelsResponse{Object: "list", Data: out}, nil +} + +func (s *BatchImagePublicService) ListItems(ctx context.Context, owner BatchImageOwner, batchID string, query BatchImageItemsQuery) (*BatchImagePublicItemsResponse, error) { + filter := BatchImageItemFilter{Limit: query.Limit, Offset: parseBatchImageCursor(query.Cursor)} + switch strings.TrimSpace(query.Status) { + case "", "all": + case "succeeded", "success": + filter.Status = BatchImageItemStatusSuccess + case "pending": + filter.Status = BatchImageItemStatusPending + case "failed": + filter.Status = BatchImageItemStatusFailed + default: + return nil, ErrBatchImageInvalidItems + } + if filter.Limit <= 0 || filter.Limit > 500 { + filter.Limit = 100 + } + items, err := s.Repo.ListBatchImageItemsForOwner(ctx, owner.UserID, owner.APIKeyID, batchID, filter) + if err != nil { + return nil, err + } + data := make([]BatchImagePublicItem, 0, len(items)) + for _, item := range items { + data = append(data, BatchImageItemToPublic(item)) + } + return &BatchImagePublicItemsResponse{ + Object: "list", + Data: data, + HasMore: len(data) == filter.Limit, + }, nil +} + +func (s *BatchImagePublicService) Cancel(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) { + job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + if isBatchImageProcessorDoneStatus(job.Status) { + if job.Status == BatchImageJobStatusFailed || job.Status == BatchImageJobStatusCancelled { + if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil { + s.enqueueBillingRetry(ctx, job.BatchID) + return nil, ErrBatchImageCancelFailed + } + s.invalidateAuthCache(ctx, owner.UserID) + } + return BatchImageJobToPublic(job), nil + } + if job.ProviderJobName != nil && strings.TrimSpace(*job.ProviderJobName) != "" { + provider, ok := s.ProviderRegistry.Get(job.Provider) + if !ok || provider == nil { + return nil, ErrBatchImageUnsupportedProvider + } + if job.AccountID == nil { + return nil, ErrBatchImageCancelFailed + } + account, err := s.AccountRepo.GetByID(ctx, *job.AccountID) + if err != nil { + return nil, ErrBatchImageCancelFailed + } + if err := provider.Cancel(ctx, job, account); err != nil { + return nil, ErrBatchImageCancelFailed + } + if eventErr := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "job_cancel_requested", map[string]any{"batch_id": job.BatchID}); eventErr != nil { + logger.L().Warn("batch_image.cancel_event_failed", + zap.String("batch_id", job.BatchID), + zap.Error(eventErr), + ) + } + if s.Queue != nil { + if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) { + return nil, ErrBatchImageCancelFailed + } + } + updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + return BatchImageJobToPublic(updated), nil + } + if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusCancelled, BatchImageTransitionOptions{ + EventType: "job_cancelled", + EventPayload: map[string]any{"batch_id": job.BatchID}, + }); err != nil { + return nil, err + } + if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil { + s.enqueueBillingRetry(ctx, job.BatchID) + return nil, ErrBatchImageCancelFailed + } + s.invalidateAuthCache(ctx, owner.UserID) + updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID) + if err != nil { + return nil, err + } + return BatchImageJobToPublic(updated), nil +} + +func (s *BatchImagePublicService) validateSubmitRequest(req BatchImageSubmitRequest) (BatchImageSubmitRequest, error) { + req.Model = strings.TrimSpace(req.Model) + req.TaskName = strings.TrimSpace(req.TaskName) + req.ParentBatchID = strings.TrimSpace(req.ParentBatchID) + req.Provider = strings.TrimSpace(req.Provider) + req.ResponseMimeType = strings.TrimSpace(req.ResponseMimeType) + req.AspectRatio = strings.TrimSpace(req.AspectRatio) + req.ImageSize = strings.TrimSpace(req.ImageSize) + if req.Model == "" { + return req, ErrBatchImageInvalidModel + } + if req.TaskName == "" { + req.TaskName = defaultBatchImageTaskName(time.Now()) + } + if len(req.TaskName) > 255 { + req.TaskName = truncateBatchImageMessage(req.TaskName, 255) + } + if req.Provider != "" && !IsSupportedBatchImageProvider(req.Provider) { + return req, ErrBatchImageUnsupportedProvider + } + if len(req.Items) == 0 { + return req, ErrBatchImageInvalidItems + } + maxItems := s.maxItems() + if len(req.Items) > maxItems { + return req, ErrBatchImageInvalidItems + } + if req.ResponseMimeType == "" { + req.ResponseMimeType = s.defaultResponseMimeType() + } + if req.ImageSize == "" { + req.ImageSize = s.defaultImageSize() + } + if !strings.EqualFold(req.ImageSize, defaultBatchImageImageSize) { + return req, ErrBatchImageInvalidItems + } + req.ImageSize = defaultBatchImageImageSize + req.Metadata = sanitizeBatchImageMetadata(req.Metadata) + + seen := make(map[string]struct{}, len(req.Items)) + totalReferenceImages := 0 + totalInlineReferenceBytes := 0 + totalOutputImages := 0 + expandedItems := make([]BatchImageSubmitItem, 0, len(req.Items)) + for i := range req.Items { + req.Items[i].CustomID = strings.TrimSpace(req.Items[i].CustomID) + if req.Items[i].CustomID == "" { + req.Items[i].CustomID = fmt.Sprintf("item_%06d", i+1) + } + outputCount := req.Items[i].OutputCount + if outputCount == 0 { + outputCount = 1 + } + if outputCount < 1 || outputCount > s.maxOutputImagesPerItem() { + return req, ErrBatchImageInvalidItems + } + totalOutputImages += outputCount + if totalOutputImages > s.maxOutputImagesPerJob() { + return req, ErrBatchImageTooManyOutputImages + } + req.Items[i].Prompt = strings.TrimSpace(req.Items[i].Prompt) + if req.Items[i].Prompt == "" { + return req, ErrBatchImageInvalidItems + } + if len(req.Items[i].Prompt) > s.maxPromptChars() { + return req, ErrBatchImagePromptTooLong + } + referenceCount, inlineReferenceBytes, err := normalizeBatchImageReferenceInputs(req.Model, &req.Items[i]) + if err != nil { + return req, err + } + totalReferenceImages += referenceCount * outputCount + if totalReferenceImages > s.maxReferenceImagesPerJob() { + return req, ErrBatchImageTooManyReferenceImages + } + totalInlineReferenceBytes += inlineReferenceBytes * outputCount + if totalInlineReferenceBytes > s.maxReferenceInlineBytesPerJob() { + return req, ErrBatchImageReferenceImagesTooLarge + } + for repeatIndex := 1; repeatIndex <= outputCount; repeatIndex++ { + expanded := req.Items[i] + expanded.OutputCount = 0 + if outputCount > 1 { + expanded.CustomID = fmt.Sprintf("%s_%0*d", req.Items[i].CustomID, batchImageRepeatSuffixWidth(outputCount), repeatIndex) + } + if _, ok := seen[expanded.CustomID]; ok { + return req, ErrBatchImageDuplicateCustomIDInRequest + } + seen[expanded.CustomID] = struct{}{} + expandedItems = append(expandedItems, expanded) + } + } + req.Items = expandedItems + return req, nil +} + +func normalizeBatchImageReferenceInputs(model string, item *BatchImageSubmitItem) (int, int, error) { + if item == nil || len(item.ReferenceImages) == 0 { + return 0, 0, nil + } + maxRefs := maxBatchImageReferenceImagesForModel(model) + if maxRefs <= 0 || len(item.ReferenceImages) > maxRefs { + return 0, 0, ErrBatchImageTooManyReferenceImages + } + out := make([]BatchImageReferenceInput, 0, len(item.ReferenceImages)) + inlineBytes := 0 + for _, ref := range item.ReferenceImages { + ref.ID = truncateBatchImageMessage(strings.TrimSpace(ref.ID), 80) + ref.Type = truncateBatchImageMessage(strings.TrimSpace(ref.Type), 40) + ref.MimeType = normalizeBatchImageReferenceMimeType(ref.MimeType) + ref.FileURI = strings.TrimSpace(ref.FileURI) + if ref.MimeType == "" { + return 0, 0, ErrBatchImageInvalidReferenceImage + } + if len(ref.Data) == 0 && ref.FileURI == "" { + return 0, 0, ErrBatchImageInvalidReferenceImage + } + if len(ref.Data) > 0 && ref.FileURI != "" { + return 0, 0, ErrBatchImageInvalidReferenceImage + } + if len(ref.Data) > maxBatchImageReferenceImageBytes { + return 0, 0, ErrBatchImageInvalidReferenceImage + } + if ref.FileURI != "" && !strings.HasPrefix(ref.FileURI, "gs://") { + return 0, 0, ErrBatchImageInvalidReferenceImage + } + inlineBytes += len(ref.Data) + out = append(out, ref) + } + item.ReferenceImages = out + return len(out), inlineBytes, nil +} + +func normalizeBatchImageReferenceMimeType(v string) string { + switch strings.ToLower(strings.TrimSpace(v)) { + case "image/jpeg", "image/jpg": + return "image/jpeg" + case "image/png": + return "image/png" + case "image/webp": + return "image/webp" + default: + return "" + } +} + +func batchImageRepeatSuffixWidth(count int) int { + if count < 10 { + return 2 + } + return len(strconv.Itoa(count)) +} + +func maxBatchImageReferenceImagesForModel(model string) int { + model = strings.ToLower(strings.TrimSpace(model)) + if strings.Contains(model, "pro-image") { + return 14 + } + if strings.Contains(model, "flash-image") { + return 3 + } + return 0 +} + +func (s *BatchImagePublicService) selectProviderAndAccount(ctx context.Context, owner BatchImageOwner, requestedProvider, model string) (BatchImageProvider, *Account, error) { + providers := batchImageProviderSelectionOrder(requestedProvider) + for _, providerName := range providers { + provider, ok := s.ProviderRegistry.Get(providerName) + if !ok || provider == nil { + continue + } + accounts, err := s.listCandidateAccounts(ctx, owner.GroupID, batchImageProviderPlatform(providerName)) + if err != nil { + return nil, nil, err + } + sort.SliceStable(accounts, func(i, j int) bool { + if accounts[i].Priority != accounts[j].Priority { + return accounts[i].Priority > accounts[j].Priority + } + return accounts[i].ID < accounts[j].ID + }) + for i := range accounts { + account := accounts[i] + if !account.IsSchedulable() || !account.IsModelSupported(model) { + continue + } + if provider.SupportsAccount(&account) { + return provider, &account, nil + } + } + } + if requestedProvider != "" { + return nil, nil, ErrBatchImageNoAccountAvailable + } + return nil, nil, ErrBatchImageNoAccountAvailable +} + +func (s *BatchImagePublicService) listCandidateAccounts(ctx context.Context, groupID *int64, platform string) ([]Account, error) { + if s.AccountRepo == nil { + return nil, ErrBatchImageNoAccountAvailable + } + if groupID != nil && *groupID > 0 { + return s.AccountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, platform) + } + return s.AccountRepo.ListSchedulableByPlatform(ctx, platform) +} + +func (s *BatchImagePublicService) ensureGroupAllowsBatchImage(ctx context.Context, groupID *int64) error { + if groupID == nil || *groupID <= 0 { + return nil + } + if s.GroupRepo == nil { + return ErrBatchImageSettlementPricingMissing + } + group, err := s.GroupRepo.GetByIDLite(ctx, *groupID) + if err != nil || group == nil { + return ErrBatchImageSettlementPricingMissing + } + if !group.AllowBatchImageGeneration { + return ErrBatchImageGroupDisabled + } + if group.Platform != PlatformGemini { + return ErrBatchImageGroupDisabled + } + return nil +} + +func (s *BatchImagePublicService) resolvePricingSnapshot(ctx context.Context, owner BatchImageOwner, req BatchImageSubmitRequest, provider string, account *Account) (*BatchImagePricingSnapshot, error) { + unit := -1.0 + groupMultiplier := 1.0 + discountMultiplier := defaultBatchImageDiscountMultiplier + holdMultiplier := defaultBatchImageHoldMultiplier + if owner.GroupID != nil && *owner.GroupID > 0 { + if s.GroupRepo == nil { + return nil, ErrBatchImageSettlementPricingMissing + } + group, err := s.GroupRepo.GetByIDLite(ctx, *owner.GroupID) + if err != nil || group == nil { + return nil, ErrBatchImageSettlementPricingMissing + } + if !group.AllowBatchImageGeneration { + return nil, ErrBatchImageGroupDisabled + } + groupDefaultMultiplier := group.RateMultiplier + if groupDefaultMultiplier < 0 { + groupDefaultMultiplier = 0 + } + effectiveGroupMultiplier := groupDefaultMultiplier + if s.UserGroupRateRepo != nil { + userRate, rateErr := s.UserGroupRateRepo.GetByUserAndGroup(ctx, owner.UserID, group.ID) + if rateErr != nil { + return nil, ErrBatchImageSettlementPricingMissing + } + if userRate != nil { + effectiveGroupMultiplier = *userRate + } + } + groupMultiplier = effectiveGroupMultiplier + if group.ImageRateIndependent { + groupMultiplier = group.ImageRateMultiplier + } + if groupMultiplier < 0 { + groupMultiplier = 0 + } + discountMultiplier = group.BatchImageDiscountMultiplier + if discountMultiplier < 0 { + discountMultiplier = 0 + } + if group.BatchImageHoldMultiplier >= 0 { + holdMultiplier = group.BatchImageHoldMultiplier + } + if configuredUnit := group.GetImagePrice(req.ImageSize); configuredUnit != nil && *configuredUnit >= 0 { + unit = *configuredUnit + } + } + if unit < 0 { + if s.Pricing == nil { + return nil, ErrBatchImageSettlementPricingMissing + } + resolvedUnit, err := s.Pricing.BatchImageUnitPrice(ctx, &BatchImageJob{Provider: provider, Model: req.Model}) + if err != nil || resolvedUnit < 0 { + return nil, ErrBatchImageSettlementPricingMissing + } + unit = resolvedUnit + } + // 定价不变式:hold 比例不得低于 discount 比例,否则成功率足够高时 + // actualCost > holdAmount,结算永远失败、冻结余额无法解冻。 + // 管理端已校验新配置,此处兜底钳制存量脏数据。 + if holdMultiplier < discountMultiplier { + logger.L().Warn("batch_image.hold_multiplier_below_discount_clamped", + zap.Float64("hold_multiplier", holdMultiplier), + zap.Float64("discount_multiplier", discountMultiplier), + ) + holdMultiplier = discountMultiplier + } + accountMultiplier := 1.0 + if account != nil { + accountMultiplier = account.BillingRateMultiplier() + } + if accountMultiplier < 0 { + accountMultiplier = 0 + } + standardUnitPrice := unit * groupMultiplier * accountMultiplier + billableUnitPrice := standardUnitPrice * discountMultiplier + holdUnitPrice := standardUnitPrice * holdMultiplier + return &BatchImagePricingSnapshot{ + BaseUnitPrice: unit, + GroupRateMultiplier: groupMultiplier, + AccountRateMultiplier: accountMultiplier, + BatchDiscountMultiplier: discountMultiplier, + HoldMultiplier: holdMultiplier, + BillableUnitPrice: billableUnitPrice, + HoldUnitPrice: holdUnitPrice, + EstimatedCost: billableUnitPrice * float64(len(req.Items)), + HoldAmount: holdUnitPrice * float64(len(req.Items)), + }, nil +} + +func (s *BatchImagePublicService) enabled() bool { + return s != nil && s.Repo != nil && s.AccountRepo != nil && s.Config != nil && s.Config.BatchImage.Enabled +} + +func (s *BatchImagePublicService) invalidateAuthCache(ctx context.Context, userID int64) { + if s != nil && s.AuthCache != nil && userID > 0 { + s.AuthCache.InvalidateAuthCacheByUserID(ctx, userID) + } +} + +func (s *BatchImagePublicService) maxItems() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxItemsPerJobDefault > 0 { + return s.Config.BatchImage.MaxItemsPerJobDefault + } + return defaultBatchImageMaxItems +} + +func (s *BatchImagePublicService) maxOutputImagesPerJob() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxOutputImagesPerJob > 0 { + return s.Config.BatchImage.MaxOutputImagesPerJob + } + return defaultBatchImageMaxOutputImages +} + +func (s *BatchImagePublicService) maxOutputImagesPerItem() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxOutputImagesPerItem > 0 { + return s.Config.BatchImage.MaxOutputImagesPerItem + } + return defaultBatchImageMaxOutputCount +} + +func (s *BatchImagePublicService) maxPromptChars() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxPromptCharsPerItem > 0 { + return s.Config.BatchImage.MaxPromptCharsPerItem + } + return defaultBatchImageMaxPromptChars +} + +func (s *BatchImagePublicService) maxReferenceImagesPerJob() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxReferenceImagesPerJob > 0 { + return s.Config.BatchImage.MaxReferenceImagesPerJob + } + return defaultBatchImageMaxReferenceImages +} + +func (s *BatchImagePublicService) maxReferenceInlineBytesPerJob() int { + if s != nil && s.Config != nil && s.Config.BatchImage.MaxReferenceInlineBytesPerJob > 0 { + return s.Config.BatchImage.MaxReferenceInlineBytesPerJob + } + return defaultBatchImageMaxReferenceBytes +} + +func (s *BatchImagePublicService) defaultResponseMimeType() string { + if s != nil && s.Config != nil && strings.TrimSpace(s.Config.BatchImage.DefaultResponseMimeType) != "" { + return strings.TrimSpace(s.Config.BatchImage.DefaultResponseMimeType) + } + return defaultBatchImageResponseMime +} + +func (s *BatchImagePublicService) defaultImageSize() string { + if s != nil && s.Config != nil && strings.TrimSpace(s.Config.BatchImage.DefaultImageSize) != "" { + return strings.TrimSpace(s.Config.BatchImage.DefaultImageSize) + } + return defaultBatchImageImageSize +} + +func BatchImageJobToPublic(job *BatchImageJob) *BatchImagePublicBatch { + if job == nil { + return nil + } + holdAmount := job.EstimatedCost + if job.HoldAmount != nil { + holdAmount = *job.HoldAmount + } + return &BatchImagePublicBatch{ + ID: job.BatchID, + Object: "image.batch", + TaskName: batchImagePublicTaskName(job), + ParentBatchID: job.ParentBatchID, + Status: PublicBatchImageStatus(job.Status), + Model: job.Model, + Provider: job.Provider, + ItemCount: job.ItemCount, + SuccessCount: job.SuccessCount, + FailCount: job.FailCount, + EstimatedCost: job.EstimatedCost, + HoldAmount: holdAmount, + ActualCost: job.ActualCost, + CreatedAt: job.CreatedAt.Unix(), + SubmittedAt: batchImageUnixPtr(job.SubmittedAt), + SettledAt: batchImageUnixPtr(job.SettledAt), + DownloadedAt: batchImageUnixPtr(job.DownloadedAt), + OutputDeletedAt: batchImageUnixPtr(job.OutputDeletedAt), + } +} + +func BatchImageItemToPublic(item *BatchImageItem) BatchImagePublicItem { + out := BatchImagePublicItem{ + CustomID: item.CustomID, + Status: "failed", + PromptPreview: item.PromptPreview, + MimeType: item.MimeType, + FileExtension: item.FileExtension, + ImageCount: item.ImageCount, + } + if item.Status == BatchImageItemStatusPending { + out.Status = "pending" + return out + } + if item.Status == BatchImageItemStatusSuccess { + out.Status = "succeeded" + return out + } + out.Error = &BatchImagePublicError{ + Code: batchImageDerefString(item.ErrorCode), + Message: sanitizeBatchImagePublicMessage(batchImageDerefString(item.ErrorMessage)), + Source: batchImageItemErrorSource(item), + } + return out +} + +func batchImageItemErrorSource(item *BatchImageItem) string { + if item == nil || item.ErrorCode == nil { + return "" + } + code := strings.TrimSpace(*item.ErrorCode) + if batchImageDerefString(item.ProviderSourceObject) != "" { + return "provider" + } + switch code { + case "EMPTY_IMAGE_OUTPUT", "PROVIDER_ITEM_FAILED": + return "provider" + case "INDEX_OUTPUT_MISSING", "INDEX_PARSE_FAILED", "DUPLICATE_CUSTOM_ID_IN_OUTPUT": + return "system" + default: + return "" + } +} + +func PublicBatchImageStatus(status string) string { + switch status { + case BatchImageJobStatusCreated, BatchImageJobStatusUploading, BatchImageJobStatusSubmitted: + return "queued" + case BatchImageJobStatusRunning: + return "running" + case BatchImageJobStatusIndexing: + return "processing_results" + case BatchImageJobStatusSettling: + return "settling" + case BatchImageJobStatusCompleted: + return "completed" + case BatchImageJobStatusFailed: + return "failed" + case BatchImageJobStatusCancelled: + return "cancelled" + case BatchImageJobStatusOutputDeleted: + return "output_deleted" + default: + return status + } +} + +func HashBatchImageSubmitRequest(req BatchImageSubmitRequest) string { + req.Metadata = sanitizeBatchImageMetadata(req.Metadata) + b, _ := json.Marshal(req) + sum := sha256.Sum256(b) + return hex.EncodeToString(sum[:]) +} + +func batchImageProviderPlatform(provider string) string { + switch provider { + case BatchImageProviderGeminiAPI, BatchImageProviderVertex: + return PlatformGemini + default: + return PlatformGemini + } +} + +func batchImageProviderSelectionOrder(requestedProvider string) []string { + if strings.TrimSpace(requestedProvider) != "" { + return []string{strings.TrimSpace(requestedProvider)} + } + return []string{BatchImageProviderGeminiAPI, BatchImageProviderVertex} +} + +func batchImageModelsFromAccountMapping(account *Account) []string { + if account == nil { + return nil + } + mapping := account.GetModelMapping() + if len(mapping) == 0 { + return nil + } + models := make(map[string]struct{}) + for model := range mapping { + model = strings.TrimSpace(model) + if model == "" { + continue + } + if strings.ContainsAny(model, "*?") { + for _, candidate := range defaultBatchImageModelCandidates() { + if matchWildcard(model, candidate) { + models[candidate] = struct{}{} + } + } + continue + } + models[model] = struct{}{} + } + out := make([]string, 0, len(models)) + for model := range models { + out = append(out, model) + } + sort.Strings(out) + return out +} + +func defaultBatchImageModelCandidates() []string { + return []string{ + "gemini-2.0-flash-exp-image-generation", + "gemini-2.5-flash-image", + "gemini-3-pro-image", + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3.1-flash-lite-image", + } +} + +func batchImageGCSRef(provider, ref string) string { + if provider == BatchImageProviderVertex && strings.HasPrefix(strings.TrimSpace(ref), "gs://") { + return strings.TrimSpace(ref) + } + return "" +} + +func batchImageProviderSubmitPublicError(err error) error { + reason := strings.TrimSpace(infraerrors.Reason(err)) + switch reason { + case "VERTEX_MANAGED_GCS_BUCKET_MISSING": + return ErrBatchImageVertexGCSBucketMissing + case "BATCH_IMAGE_PROVIDER_MISSING_API_KEY": + return ErrBatchImageProviderMissingAPIKey + case "BATCH_IMAGE_PROVIDER_MISSING_SERVICE_ACCOUNT": + return ErrBatchImageProviderMissingServiceAccount + case "BATCH_IMAGE_PROVIDER_UNSUPPORTED_ACCOUNT": + return ErrBatchImageProviderUnsupportedAccount + default: + return ErrBatchImageProviderSubmitFailed + } +} + +func batchImagePublicTaskName(job *BatchImageJob) string { + if job == nil { + return "" + } + if strings.TrimSpace(job.TaskName) != "" { + return strings.TrimSpace(job.TaskName) + } + return defaultBatchImageTaskName(job.CreatedAt) +} + +func defaultBatchImageTaskName(now time.Time) string { + if now.IsZero() { + now = time.Now() + } + return now.Format("2006-01-02 15:04:05") +} + +func batchImageProviderSubmitRecordCode(err error) string { + reason := strings.TrimSpace(infraerrors.Reason(err)) + if reason == "" || reason == "BATCH_IMAGE_PROVIDER_SUBMIT_FAILED" { + return "PROVIDER_SUBMIT_FAILED" + } + return reason +} + +func parseBatchImageListTime(raw string) *time.Time { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil + } + if unix, err := strconv.ParseInt(raw, 10, 64); err == nil && unix > 0 { + t := time.Unix(unix, 0) + return &t + } + if t, err := time.Parse(time.RFC3339, raw); err == nil { + return &t + } + if t, err := time.Parse("2006-01-02", raw); err == nil { + return &t + } + return nil +} + +func sanitizeBatchImageMetadata(in map[string]string) map[string]string { + if len(in) == 0 { + return nil + } + keys := make([]string, 0, len(in)) + for k := range in { + keys = append(keys, k) + } + sort.Strings(keys) + out := make(map[string]string, len(keys)) + for _, k := range keys { + key := strings.TrimSpace(k) + if key == "" || len(key) > 64 { + continue + } + value := strings.TrimSpace(in[k]) + if len(value) > 256 { + value = value[:256] + } + out[key] = value + if len(out) >= 20 { + break + } + } + return out +} + +func sanitizeBatchImagePublicMessage(message string) string { + message = strings.TrimSpace(message) + for _, marker := range []string{"gs://", "files/", "projects/"} { + if strings.Contains(message, marker) { + message = "upstream provider operation failed" + break + } + } + if len(message) > maxBatchImagePublicErrorChars { + message = message[:maxBatchImagePublicErrorChars] + } + return message +} + +func batchImageUnixPtr(t *time.Time) *int64 { + if t == nil { + return nil + } + v := t.Unix() + return &v +} + +func parseBatchImageCursor(cursor string) int { + offset, err := strconv.Atoi(strings.TrimSpace(cursor)) + if err != nil || offset < 0 { + return 0 + } + return offset +} diff --git a/backend/internal/service/batch_image_public_test.go b/backend/internal/service/batch_image_public_test.go new file mode 100644 index 0000000000..2e5b930a7b --- /dev/null +++ b/backend/internal/service/batch_image_public_test.go @@ -0,0 +1,963 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/stretchr/testify/require" +) + +func TestBatchImagePublicService_Submit(t *testing.T) { + ctx := context.Background() + + t.Run("rejects when disabled", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(false) + _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageDisabled) + }) + + t.Run("accepts valid request stores refs and enqueues once", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + + got, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.NoError(t, err) + require.Equal(t, "image.batch", got.Object) + require.Equal(t, "queued", got.Status) + require.Equal(t, BatchImageProviderGeminiAPI, got.Provider) + require.Equal(t, 2, got.ItemCount) + require.Equal(t, 0.25, got.EstimatedCost) + require.Len(t, repo.jobs, 1) + require.Len(t, gemini.submits, 1) + require.Equal(t, []string{got.ID}, queue.enqueued) + billing := svc.BillingRepo.(*fakeBatchImageBillingRepo) + require.Len(t, billing.reserves, 1) + require.Equal(t, BatchImageHoldRequestID(got.ID), billing.reserves[0].RequestID) + require.InDelta(t, 0.3, billing.reserves[0].HoldAmount, 1e-12) + require.Empty(t, billing.releases) + authCache := svc.AuthCache.(*fakeBatchImageAuthCacheInvalidator) + require.Equal(t, []int64{11}, authCache.userIDs) + + job := repo.jobs[got.ID] + require.Equal(t, BatchImageJobStatusSubmitted, job.Status) + require.Equal(t, "providers/gemini_api/job", batchImageDerefString(job.ProviderJobName)) + require.Equal(t, "files/gemini_api/input", batchImageDerefString(job.ProviderInputRef)) + require.Equal(t, "files/gemini_api/output", batchImageDerefString(job.ProviderOutputRef)) + require.NotNil(t, job.AccountID) + require.Equal(t, int64(202), *job.AccountID) + require.Equal(t, 1, job.PricingSnapshotVersion) + require.InDelta(t, 0.25, job.BaseUnitPrice, 1e-12) + require.InDelta(t, 1.0, job.GroupRateMultiplier, 1e-12) + require.InDelta(t, 1.0, job.AccountRateMultiplier, 1e-12) + require.InDelta(t, 0.5, job.BatchDiscountMultiplier, 1e-12) + require.InDelta(t, 0.6, job.HoldMultiplier, 1e-12) + require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12) + require.InDelta(t, 0.15, job.HoldUnitPrice, 1e-12) + }) + + t.Run("combines user group image rate account rate discount and hold margin", func(t *testing.T) { + svc, repo, _, _, _ := newTestBatchImagePublicService(true) + groupID := int64(7) + accountMultiplier := 1.25 + accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo) + accountRepo.accounts[1].RateMultiplier = &accountMultiplier + svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{ + groupID: { + ID: groupID, + Platform: PlatformGemini, + RateMultiplier: 2.0, + AllowImageGeneration: true, + AllowBatchImageGeneration: true, + ImageRateIndependent: false, + BatchImageDiscountMultiplier: 0.8, + BatchImageHoldMultiplier: 0.6, + }, + }} + userRate := 0.5 + svc.UserGroupRateRepo = &publicBatchImageUserGroupRateRepo{rates: map[int64]*float64{groupID: &userRate}} + + got, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "") + require.NoError(t, err) + require.InDelta(t, 0.25, got.EstimatedCost, 1e-12) + + job := repo.jobs[got.ID] + require.InDelta(t, 0.25, job.BaseUnitPrice, 1e-12) + require.InDelta(t, 0.5, job.GroupRateMultiplier, 1e-12) + require.InDelta(t, 1.25, job.AccountRateMultiplier, 1e-12) + require.InDelta(t, 0.8, job.BatchDiscountMultiplier, 1e-12) + // 配置的 hold(0.6) < discount(0.8) 属于会导致结算死锁的脏数据, + // 快照时被钳制为 discount,保证 holdAmount >= 实际成本上限。 + require.InDelta(t, 0.8, job.HoldMultiplier, 1e-12) + require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12) + require.InDelta(t, 0.125, job.HoldUnitPrice, 1e-12) + require.InDelta(t, 0.25, *job.HoldAmount, 1e-12) + }) + + t.Run("uses configured group 1k image price for batch image base price", func(t *testing.T) { + svc, repo, _, _, _ := newTestBatchImagePublicService(true) + groupID := int64(7) + imagePrice := 0.134 + svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{ + groupID: { + ID: groupID, + Platform: PlatformGemini, + RateMultiplier: 1.0, + AllowImageGeneration: true, + AllowBatchImageGeneration: true, + ImagePrice1K: &imagePrice, + BatchImageDiscountMultiplier: 0.5, + BatchImageHoldMultiplier: 0.6, + }, + }} + + got, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "") + require.NoError(t, err) + require.InDelta(t, 0.134, got.EstimatedCost, 1e-12) + + job := repo.jobs[got.ID] + require.InDelta(t, 0.134, job.BaseUnitPrice, 1e-12) + require.InDelta(t, 0.067, job.BillableUnitPrice, 1e-12) + require.InDelta(t, 0.0804, job.HoldUnitPrice, 1e-12) + require.InDelta(t, 0.1608, *job.HoldAmount, 1e-12) + }) + + t.Run("pricing missing rejects before provider submit", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + svc.Pricing = &fakeBatchImagePricingResolver{err: ErrBatchImageSettlementPricingMissing} + + _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageSettlementPricingMissing) + require.Empty(t, repo.jobs) + require.Empty(t, queue.enqueued) + require.Empty(t, gemini.submits) + }) + + t.Run("group batch image disabled rejects before provider submit", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + groupID := int64(7) + svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{ + groupID: { + ID: groupID, + Platform: PlatformGemini, + RateMultiplier: 1, + AllowBatchImageGeneration: false, + BatchImageDiscountMultiplier: 0.5, + BatchImageHoldMultiplier: 0.6, + }, + }} + + _, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageGroupDisabled) + require.Empty(t, repo.jobs) + require.Empty(t, queue.enqueued) + require.Empty(t, gemini.submits) + }) + + t.Run("group pricing load failure rejects before provider submit", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + groupID := int64(404) + + _, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageSettlementPricingMissing) + require.Empty(t, repo.jobs) + require.Empty(t, queue.enqueued) + require.Empty(t, gemini.submits) + }) + + t.Run("generates custom ids deterministically", func(t *testing.T) { + svc, _, _, gemini, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + req.Items[0].CustomID = "" + req.Items[1].CustomID = "" + + _, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.NoError(t, err) + require.Len(t, gemini.submits, 1) + require.Equal(t, "item_000001", gemini.submits[0].Items[0].CustomID) + require.Equal(t, "item_000002", gemini.submits[0].Items[1].CustomID) + }) + + t.Run("expands output count into separate billable items", func(t *testing.T) { + svc, repo, _, gemini, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + req.Items = []BatchImageSubmitItem{ + {CustomID: "cover", Prompt: "hero", OutputCount: 3, ReferenceImages: []BatchImageReferenceInput{{MimeType: "image/png", Data: []byte("ref")}}}, + } + + got, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.NoError(t, err) + require.Equal(t, 3, got.ItemCount) + require.InDelta(t, 0.375, got.EstimatedCost, 1e-12) + require.Len(t, gemini.submits, 1) + require.Len(t, gemini.submits[0].Items, 3) + require.Equal(t, []string{"cover_01", "cover_02", "cover_03"}, []string{ + gemini.submits[0].Items[0].CustomID, + gemini.submits[0].Items[1].CustomID, + gemini.submits[0].Items[2].CustomID, + }) + require.Len(t, gemini.submits[0].Items[0].ReferenceImages, 1) + require.Len(t, repo.items[got.ID], 3) + }) + + t.Run("validates request fields", func(t *testing.T) { + tests := []struct { + name string + mutate func(*BatchImageSubmitRequest) + want error + }{ + {name: "missing_model", mutate: func(r *BatchImageSubmitRequest) { r.Model = "" }, want: ErrBatchImageInvalidModel}, + {name: "empty_items", mutate: func(r *BatchImageSubmitRequest) { r.Items = nil }, want: ErrBatchImageInvalidItems}, + {name: "duplicate_custom_ids", mutate: func(r *BatchImageSubmitRequest) { r.Items[1].CustomID = r.Items[0].CustomID }, want: ErrBatchImageDuplicateCustomIDInRequest}, + {name: "empty_prompt", mutate: func(r *BatchImageSubmitRequest) { r.Items[0].Prompt = " " }, want: ErrBatchImageInvalidItems}, + {name: "prompt_too_long", mutate: func(r *BatchImageSubmitRequest) { r.Items[0].Prompt = strings.Repeat("x", 9) }, want: ErrBatchImagePromptTooLong}, + {name: "unsupported_provider", mutate: func(r *BatchImageSubmitRequest) { r.Provider = "other" }, want: ErrBatchImageUnsupportedProvider}, + {name: "vertex_rejects_2k", mutate: func(r *BatchImageSubmitRequest) { r.Provider = BatchImageProviderVertex; r.ImageSize = "2K" }, want: ErrBatchImageInvalidItems}, + {name: "too_many_outputs_per_item", mutate: func(r *BatchImageSubmitRequest) { + r.Items[0].OutputCount = 5 + }, want: ErrBatchImageInvalidItems}, + {name: "too_many_reference_images_for_flash", mutate: func(r *BatchImageSubmitRequest) { + r.Model = "gemini-2.5-flash-image" + r.Items[0].ReferenceImages = []BatchImageReferenceInput{ + {MimeType: "image/png", Data: []byte("1")}, + {MimeType: "image/png", Data: []byte("2")}, + {MimeType: "image/png", Data: []byte("3")}, + {MimeType: "image/png", Data: []byte("4")}, + } + }, want: ErrBatchImageTooManyReferenceImages}, + {name: "bad_reference_mime", mutate: func(r *BatchImageSubmitRequest) { + r.Items[0].ReferenceImages = []BatchImageReferenceInput{{MimeType: "application/octet-stream", Data: []byte("x")}} + }, want: ErrBatchImageInvalidReferenceImage}, + {name: "reference_requires_data_or_file_uri", mutate: func(r *BatchImageSubmitRequest) { + r.Items[0].ReferenceImages = []BatchImageReferenceInput{{MimeType: "image/png"}} + }, want: ErrBatchImageInvalidReferenceImage}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + tt.mutate(&req) + + _, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.ErrorIs(t, err, tt.want) + }) + } + }) + + t.Run("rejects too many items", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + req.Items = append(req.Items, BatchImageSubmitItem{CustomID: "too_many", Prompt: "x"}) + + _, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.ErrorIs(t, err, ErrBatchImageInvalidItems) + }) + + t.Run("rejects too many output images", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + svc.Config.BatchImage.MaxOutputImagesPerJob = 3 + req := validBatchImageSubmitRequest() + req.Items[0].OutputCount = 2 + req.Items[1].OutputCount = 2 + + _, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.ErrorIs(t, err, ErrBatchImageTooManyOutputImages) + }) + + t.Run("rejects too many reference images across request", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + svc.Config.BatchImage.MaxReferenceImagesPerJob = 3 + req := validBatchImageSubmitRequest() + req.Model = "gemini-2.5-flash-image" + req.Items[0].ReferenceImages = []BatchImageReferenceInput{ + {MimeType: "image/png", Data: []byte("1")}, + {MimeType: "image/png", Data: []byte("2")}, + } + req.Items[1].ReferenceImages = []BatchImageReferenceInput{ + {MimeType: "image/png", Data: []byte("3")}, + {MimeType: "image/png", Data: []byte("4")}, + } + + _, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.ErrorIs(t, err, ErrBatchImageTooManyReferenceImages) + }) + + t.Run("rejects too much inline reference image data across request", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + svc.Config.BatchImage.MaxReferenceImagesPerJob = 10 + svc.Config.BatchImage.MaxReferenceInlineBytesPerJob = 4 + req := validBatchImageSubmitRequest() + req.Model = "gemini-2.5-flash-image" + req.Items[0].ReferenceImages = []BatchImageReferenceInput{{MimeType: "image/png", Data: []byte("123")}} + req.Items[1].ReferenceImages = []BatchImageReferenceInput{{MimeType: "image/png", Data: []byte("456")}} + + _, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.ErrorIs(t, err, ErrBatchImageReferenceImagesTooLarge) + }) + + t.Run("selects requested provider", func(t *testing.T) { + svc, _, _, gemini, vertex := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + req.Provider = BatchImageProviderVertex + + got, err := svc.Submit(ctx, testBatchImageOwner(), req, "") + require.NoError(t, err) + require.Equal(t, BatchImageProviderVertex, got.Provider) + require.Empty(t, gemini.submits) + require.Len(t, vertex.submits, 1) + }) + + t.Run("insufficient balance rejects before provider submit", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + billing := &fakeBatchImageBillingRepo{err: ErrBatchImageInsufficientBalance} + svc.BillingRepo = billing + + _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageInsufficientBalance) + require.Empty(t, queue.enqueued) + require.Empty(t, gemini.submits) + require.Len(t, billing.reserves, 1) + require.Empty(t, billing.releases) + require.Len(t, repo.jobs, 1) + for _, job := range repo.jobs { + require.Equal(t, BatchImageJobStatusFailed, job.Status) + require.Equal(t, "INSUFFICIENT_BALANCE", batchImageDerefString(job.LastErrorCode)) + require.NotNil(t, job.UserDeletedAt) + } + }) + + t.Run("provider failure marks failed and does not enqueue", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + gemini.submitErr = errors.New("projects/secret-provider-job failed") + billing := svc.BillingRepo.(*fakeBatchImageBillingRepo) + + _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageProviderSubmitFailed) + require.Empty(t, queue.enqueued) + require.Len(t, billing.reserves, 1) + require.Len(t, billing.releases, 1) + require.Equal(t, BatchImageReleaseRequestID(billing.reserves[0].BatchID), billing.releases[0].RequestID) + require.Len(t, repo.jobs, 1) + for _, job := range repo.jobs { + require.Equal(t, BatchImageJobStatusFailed, job.Status) + require.Equal(t, "PROVIDER_SUBMIT_FAILED", batchImageDerefString(job.LastErrorCode)) + require.Equal(t, "upstream provider operation failed", batchImageDerefString(job.LastErrorMessage)) + require.NotNil(t, job.UserDeletedAt) + } + }) + + t.Run("provider failure with release failure enqueues billing retry", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + gemini.submitErr = errors.New("projects/secret-provider-job failed") + billing := svc.BillingRepo.(*fakeBatchImageBillingRepo) + billing.releaseErr = errors.New("billing database timeout") + + _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageBillingHoldFailed) + require.Len(t, billing.reserves, 1) + require.Len(t, billing.releases, 1) + require.Len(t, repo.jobs, 1) + for _, job := range repo.jobs { + require.Equal(t, BatchImageJobStatusFailed, job.Status) + require.Equal(t, "BILLING_RELEASE_FAILED", batchImageDerefString(job.LastErrorCode)) + require.Equal(t, []string{job.BatchID}, queue.enqueued) + } + }) + + t.Run("queue failure is recorded after provider submit", func(t *testing.T) { + svc, repo, queue, _, _ := newTestBatchImagePublicService(true) + queue.err = errors.New("redis unavailable") + billing := svc.BillingRepo.(*fakeBatchImageBillingRepo) + + _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.ErrorIs(t, err, ErrBatchImageQueueFailed) + require.Len(t, billing.reserves, 1) + require.Empty(t, billing.releases) + require.Len(t, repo.jobs, 1) + for _, job := range repo.jobs { + require.Equal(t, BatchImageJobStatusSubmitted, job.Status) + require.Equal(t, "QUEUE_FAILED", batchImageDerefString(job.LastErrorCode)) + require.Contains(t, repo.events[job.BatchID], "queue_failed") + } + }) + + t.Run("idempotency returns same batch without provider resubmit", func(t *testing.T) { + svc, _, queue, gemini, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + + first, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key") + require.NoError(t, err) + second, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key") + require.NoError(t, err) + + require.Equal(t, first.ID, second.ID) + require.Len(t, gemini.submits, 1) + require.Equal(t, []string{first.ID}, queue.enqueued) + }) + + t.Run("idempotency conflict rejects changed request", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + first, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key") + require.NoError(t, err) + + req.Items[0].Prompt = "diff" + second, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key") + require.Nil(t, second) + require.ErrorIs(t, err, ErrBatchImageIdempotencyConflict) + require.NotEmpty(t, first.ID) + }) + + t.Run("public response does not expose internals", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + got, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + require.NoError(t, err) + + body, err := json.Marshal(got) + require.NoError(t, err) + requireBatchImagePublicJSONHasNoInternals(t, string(body)) + }) +} + +func TestBatchImagePublicService_List(t *testing.T) { + ctx := context.Background() + svc, repo, _, _, _ := newTestBatchImagePublicService(true) + visibleKeyID := int64(22) + otherKeyID := int64(23) + + repo.jobs["visible-1"] = &BatchImageJob{ + BatchID: "visible-1", + UserID: 11, + APIKeyID: &visibleKeyID, + Status: BatchImageJobStatusCompleted, + Provider: BatchImageProviderVertex, + Model: "gemini-3.1-flash-lite-image", + ItemCount: 1, + CreatedAt: time.Now(), + } + repo.jobs["hidden-other-key"] = &BatchImageJob{ + BatchID: "hidden-other-key", + UserID: 11, + APIKeyID: &otherKeyID, + Status: BatchImageJobStatusCompleted, + Provider: BatchImageProviderVertex, + Model: "gemini-3.1-flash-lite-image", + ItemCount: 1, + CreatedAt: time.Now(), + } + + got, err := svc.List(ctx, BatchImageOwner{UserID: 11, APIKeyID: visibleKeyID}, BatchImageJobsQuery{Limit: 20}) + require.NoError(t, err) + require.Equal(t, "list", got.Object) + require.Len(t, got.Data, 1) + require.Equal(t, "visible-1", got.Data[0].ID) + require.False(t, got.HasMore) +} + +func TestBatchImagePublicService_ListModels(t *testing.T) { + ctx := context.Background() + + t.Run("requires explicit account model mapping", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + + got, err := svc.ListModels(ctx, testBatchImageOwner()) + require.NoError(t, err) + require.Equal(t, "list", got.Object) + require.Empty(t, got.Data) + }) + + t.Run("returns priced models from selected account group", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + groupID := int64(7) + svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{ + groupID: { + ID: groupID, + Platform: PlatformGemini, + RateMultiplier: 1, + AllowImageGeneration: true, + AllowBatchImageGeneration: true, + BatchImageDiscountMultiplier: 0.5, + BatchImageHoldMultiplier: 0.6, + }, + }} + accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo) + accountRepo.accounts = []Account{testBatchImageMappedAccount(303, AccountTypeAPIKey, map[string]any{ + "gemini-2.5-flash-image": "gemini-2.5-flash-image", + })} + + got, err := svc.ListModels(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}) + require.NoError(t, err) + require.Equal(t, []BatchImagePublicModel{{ + ID: "gemini-2.5-flash-image", + Object: "image.batch.model", + Provider: BatchImageProviderGeminiAPI, + }, { + ID: "gemini-2.5-flash-image", + Object: "image.batch.model", + Provider: BatchImageProviderVertex, + }}, got.Data) + }) + + t.Run("expands wildcard mappings against batch image candidates", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo) + accountRepo.accounts = []Account{testBatchImageMappedAccount(303, AccountTypeAPIKey, map[string]any{ + "gemini-3.1-*": "gemini-3.1-flash-lite-image", + })} + + got, err := svc.ListModels(ctx, testBatchImageOwner()) + require.NoError(t, err) + require.NotEmpty(t, got.Data) + ids := make([]string, 0, len(got.Data)) + for _, model := range got.Data { + ids = append(ids, model.ID) + } + require.Contains(t, ids, "gemini-3.1-flash-image") + require.Contains(t, ids, "gemini-3.1-flash-lite-image") + require.NotContains(t, ids, "gemini-2.5-flash-image") + }) + + t.Run("filters models without batch image pricing", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + svc.Pricing = &fakeBatchImagePricingResolver{ + unitPrice: 0.25, + missingModels: map[string]bool{"gemini-3.1-flash-lite-image": true}, + } + accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo) + accountRepo.accounts = []Account{testBatchImageMappedAccount(303, AccountTypeAPIKey, map[string]any{ + "gemini-2.5-flash-image": "gemini-2.5-flash-image", + "gemini-3.1-flash-lite-image": "gemini-3.1-flash-lite-image", + })} + + got, err := svc.ListModels(ctx, testBatchImageOwner()) + require.NoError(t, err) + ids := make([]string, 0, len(got.Data)) + for _, model := range got.Data { + ids = append(ids, model.ID) + } + require.Contains(t, ids, "gemini-2.5-flash-image") + require.NotContains(t, ids, "gemini-3.1-flash-lite-image") + }) + + t.Run("rejects when group disables batch image", func(t *testing.T) { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + groupID := int64(7) + svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{ + groupID: {ID: groupID, AllowBatchImageGeneration: false}, + }} + + _, err := svc.ListModels(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}) + require.ErrorIs(t, err, ErrBatchImageGroupDisabled) + }) +} + +func TestBatchImagePublicService_StatusItemsAndCancel(t *testing.T) { + ctx := context.Background() + + t.Run("status is owner scoped and maps public status", func(t *testing.T) { + svc, repo, _, _, _ := newTestBatchImagePublicService(true) + apiKeyID := int64(22) + accountID := int64(101) + repo.jobs["imgbatch_status"] = &BatchImageJob{ + BatchID: "imgbatch_status", + UserID: 11, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: BatchImageJobStatusIndexing, + ProviderJobName: batchImageStringPtr("providers/internal/job"), + CreatedAt: time.Now(), + } + + got, err := svc.Get(ctx, testBatchImageOwner(), "imgbatch_status") + require.NoError(t, err) + require.Equal(t, "processing_results", got.Status) + body, err := json.Marshal(got) + require.NoError(t, err) + requireBatchImagePublicJSONHasNoInternals(t, string(body)) + + _, err = svc.Get(ctx, BatchImageOwner{UserID: 11, APIKeyID: 999}, "imgbatch_status") + require.ErrorIs(t, err, ErrBatchImageJobNotFound) + }) + + t.Run("items are filtered paginated and sanitized", func(t *testing.T) { + svc, repo, _, _, _ := newTestBatchImagePublicService(true) + apiKeyID := int64(22) + repo.jobs["imgbatch_items"] = &BatchImageJob{ + BatchID: "imgbatch_items", + UserID: 11, + APIKeyID: &apiKeyID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: BatchImageJobStatusCompleted, + CreatedAt: time.Now(), + } + sourceObject := "gs://bucket/internal/output.jsonl" + mime := "image/png" + ext := "png" + code := "SAFETY_BLOCKED" + msg := "blocked in gs://bucket/internal/output.jsonl" + repo.items["imgbatch_items"] = []CreateBatchImageItemParams{ + {JobID: "imgbatch_items", CustomID: "ok_1", Status: BatchImageItemStatusSuccess, ProviderSourceObject: &sourceObject, MimeType: &mime, FileExtension: &ext, ImageCount: 1}, + {JobID: "imgbatch_items", CustomID: "bad_1", Status: BatchImageItemStatusFailed, ProviderSourceObject: &sourceObject, ErrorCode: &code, ErrorMessage: &msg}, + {JobID: "imgbatch_items", CustomID: "ok_2", Status: BatchImageItemStatusSuccess, MimeType: &mime, FileExtension: &ext, ImageCount: 1}, + } + + page, err := svc.ListItems(ctx, testBatchImageOwner(), "imgbatch_items", BatchImageItemsQuery{Limit: 1}) + require.NoError(t, err) + require.True(t, page.HasMore) + require.Len(t, page.Data, 1) + require.Equal(t, "ok_1", page.Data[0].CustomID) + + filtered, err := svc.ListItems(ctx, testBatchImageOwner(), "imgbatch_items", BatchImageItemsQuery{Status: "failed", Limit: 100}) + require.NoError(t, err) + require.False(t, filtered.HasMore) + require.Len(t, filtered.Data, 1) + require.Equal(t, "failed", filtered.Data[0].Status) + require.NotNil(t, filtered.Data[0].Error) + require.Equal(t, "upstream provider operation failed", filtered.Data[0].Error.Message) + + body, err := json.Marshal(filtered) + require.NoError(t, err) + requireBatchImagePublicJSONHasNoInternals(t, string(body)) + require.NotContains(t, string(body), "download_url") + + _, err = svc.ListItems(ctx, BatchImageOwner{UserID: 12, APIKeyID: 22}, "imgbatch_items", BatchImageItemsQuery{}) + require.ErrorIs(t, err, ErrBatchImageJobNotFound) + }) + + t.Run("cancel active job calls provider and waits for confirmed terminal state", func(t *testing.T) { + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + apiKeyID := int64(22) + accountID := int64(101) + holdAmount := 0.5 + holdID := BatchImageHoldRequestID("imgbatch_cancel") + repo.jobs["imgbatch_cancel"] = &BatchImageJob{ + BatchID: "imgbatch_cancel", + UserID: 11, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: BatchImageJobStatusSubmitted, + ProviderJobName: batchImageStringPtr("providers/internal/job"), + EstimatedCost: holdAmount, + HoldAmount: &holdAmount, + HoldID: &holdID, + CreatedAt: time.Now(), + } + + got, err := svc.Cancel(ctx, testBatchImageOwner(), "imgbatch_cancel") + require.NoError(t, err) + require.Equal(t, "queued", got.Status) + require.Equal(t, 1, gemini.cancelCount) + billing := svc.BillingRepo.(*fakeBatchImageBillingRepo) + require.Empty(t, billing.releases) + require.Equal(t, []string{"imgbatch_cancel"}, queue.enqueued) + require.Equal(t, BatchImageJobStatusSubmitted, repo.jobs["imgbatch_cancel"].Status) + require.Contains(t, repo.events["imgbatch_cancel"], "job_cancel_requested") + }) + + t.Run("cancel terminal job is idempotent", func(t *testing.T) { + svc, repo, _, gemini, _ := newTestBatchImagePublicService(true) + apiKeyID := int64(22) + repo.jobs["imgbatch_done"] = &BatchImageJob{ + BatchID: "imgbatch_done", + UserID: 11, + APIKeyID: &apiKeyID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: BatchImageJobStatusCompleted, + CreatedAt: time.Now(), + } + + got, err := svc.Cancel(ctx, testBatchImageOwner(), "imgbatch_done") + require.NoError(t, err) + require.Equal(t, "completed", got.Status) + require.Zero(t, gemini.cancelCount) + }) + + t.Run("cancel hides provider raw errors behind public error", func(t *testing.T) { + svc, repo, _, gemini, _ := newTestBatchImagePublicService(true) + gemini.cancelErr = errors.New("projects/secret-provider-job not found") + apiKeyID := int64(22) + accountID := int64(101) + repo.jobs["imgbatch_cancel_error"] = &BatchImageJob{ + BatchID: "imgbatch_cancel_error", + UserID: 11, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-2.5-flash-image", + Status: BatchImageJobStatusSubmitted, + ProviderJobName: batchImageStringPtr("providers/internal/job"), + CreatedAt: time.Now(), + } + + _, err := svc.Cancel(ctx, testBatchImageOwner(), "imgbatch_cancel_error") + require.ErrorIs(t, err, ErrBatchImageCancelFailed) + require.Equal(t, "BATCH_IMAGE_CANCEL_FAILED", infraerrors.Reason(err)) + require.NotContains(t, infraerrors.Message(err), "projects/") + }) +} + +func newTestBatchImagePublicService(enabled bool) (*BatchImagePublicService, *fakeBatchImageRepository, *publicBatchImageQueue, *publicBatchImageProvider, *publicBatchImageProvider) { + repo := newFakeBatchImageRepository() + queue := &publicBatchImageQueue{} + gemini := &publicBatchImageProvider{name: BatchImageProviderGeminiAPI} + vertex := &publicBatchImageProvider{name: BatchImageProviderVertex} + svc := &BatchImagePublicService{ + Repo: repo, + AccountRepo: &publicBatchImageAccountRepo{accounts: []Account{testBatchImageAccount(101, AccountTypeAPIKey), testBatchImageAccount(202, AccountTypeServiceAccount)}}, + Queue: queue, + ProviderRegistry: NewBatchImageProviderRegistry( + gemini, + vertex, + ), + Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}, + BillingRepo: &fakeBatchImageBillingRepo{}, + AuthCache: &fakeBatchImageAuthCacheInvalidator{}, + Config: &config.Config{BatchImage: config.BatchImageConfig{ + Enabled: enabled, + MaxItemsPerJobDefault: 2, + MaxPromptCharsPerItem: 8, + DefaultResponseMimeType: "image/png", + DefaultImageSize: "1K", + }}, + } + return svc, repo, queue, gemini, vertex +} + +func testBatchImageOwner() BatchImageOwner { + return BatchImageOwner{UserID: 11, APIKeyID: 22} +} + +type fakeBatchImageAuthCacheInvalidator struct { + keys []string + userIDs []int64 + groupIDs []int64 +} + +func (f *fakeBatchImageAuthCacheInvalidator) InvalidateAuthCacheByKey(_ context.Context, key string) { + f.keys = append(f.keys, key) +} + +func (f *fakeBatchImageAuthCacheInvalidator) InvalidateAuthCacheByUserID(_ context.Context, userID int64) { + f.userIDs = append(f.userIDs, userID) +} + +func (f *fakeBatchImageAuthCacheInvalidator) InvalidateAuthCacheByGroupID(_ context.Context, groupID int64) { + f.groupIDs = append(f.groupIDs, groupID) +} + +func validBatchImageSubmitRequest() BatchImageSubmitRequest { + return BatchImageSubmitRequest{ + Model: "gemini-2.5-flash-image", + Provider: BatchImageProviderGeminiAPI, + ResponseMimeType: "image/png", + AspectRatio: "1:1", + ImageSize: "1K", + Metadata: map[string]string{"project": "campaign-a", "secret": strings.Repeat("x", 300)}, + Items: []BatchImageSubmitItem{ + {CustomID: "cover_001", Prompt: "hero"}, + {CustomID: "cover_002", Prompt: "clean"}, + }, + } +} + +func testBatchImageAccount(id int64, accountType string) Account { + return Account{ + ID: id, + Platform: PlatformGemini, + Type: accountType, + Status: StatusActive, + Schedulable: true, + Priority: int(id), + Credentials: map[string]any{"api_key": "test-secret"}, + Concurrency: 1, + RateLimitedAt: nil, + } +} + +func testBatchImageMappedAccount(id int64, accountType string, mapping map[string]any) Account { + account := testBatchImageAccount(id, accountType) + account.Credentials["model_mapping"] = mapping + return account +} + +func requireBatchImagePublicJSONHasNoInternals(t *testing.T, body string) { + t.Helper() + for _, forbidden := range []string{ + "provider_job_name", + "provider_input_ref", + "provider_output_ref", + "gcs_input_uri", + "gcs_output_uri", + "account_id", + "service_account", + "api_key", + "download_url", + "providers/", + "files/", + "gs://", + } { + require.NotContains(t, body, forbidden) + } +} + +type publicBatchImageAccountRepo struct { + accounts []Account +} + +func (r *publicBatchImageAccountRepo) GetByID(_ context.Context, id int64) (*Account, error) { + for i := range r.accounts { + if r.accounts[i].ID == id { + return &r.accounts[i], nil + } + } + return nil, errors.New("account not found") +} + +func (r *publicBatchImageAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) { + out := make([]Account, 0, len(r.accounts)) + for _, account := range r.accounts { + if account.Platform == platform { + out = append(out, account) + } + } + return out, nil +} + +func (r *publicBatchImageAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, _ int64, platform string) ([]Account, error) { + return r.ListSchedulableByPlatform(ctx, platform) +} + +type publicBatchImageQueue struct { + enqueued []string + err error +} + +func (q *publicBatchImageQueue) Enqueue(_ context.Context, batchID string) error { + if q.err != nil { + return q.err + } + for _, existing := range q.enqueued { + if existing == batchID { + return ErrBatchImageAlreadyQueued + } + } + q.enqueued = append(q.enqueued, batchID) + return nil +} + +func (q *publicBatchImageQueue) Reserve(context.Context, time.Duration) (ReservedBatchImageJob, error) { + return ReservedBatchImageJob{}, ErrBatchImageQueueEmpty +} + +func (q *publicBatchImageQueue) RequeueAfter(context.Context, string, time.Duration) error { + return nil +} + +func (q *publicBatchImageQueue) Ack(context.Context, string) error { + return nil +} + +func (q *publicBatchImageQueue) Heartbeat(context.Context, string) error { + return nil +} + +func (q *publicBatchImageQueue) MoveDueDelayedToReady(context.Context, int) (int, error) { + return 0, nil +} + +func (q *publicBatchImageQueue) RecoverStaleActive(context.Context, time.Duration, int) (int, error) { + return 0, nil +} + +func (q *publicBatchImageQueue) TryAcquireJobLock(context.Context, string, time.Duration) (BatchImageJobLock, bool, error) { + return nil, false, nil +} + +type publicBatchImageProvider struct { + name string + submits []BatchImageInput + submitErr error + cancelCount int + cancelErr error + result string + cleanupTargets []CleanupTarget + cleanupErr error +} + +func (p *publicBatchImageProvider) Name() string { return p.name } + +func (p *publicBatchImageProvider) SupportsAccount(*Account) bool { return true } + +func (p *publicBatchImageProvider) Submit(_ context.Context, _ *BatchImageJob, _ *Account, input BatchImageInput) (*BatchProviderJob, error) { + p.submits = append(p.submits, input) + if p.submitErr != nil { + return nil, p.submitErr + } + return &BatchProviderJob{ + ProviderJobName: "providers/" + p.name + "/job", + ProviderInputRef: "files/" + p.name + "/input", + ProviderOutputRef: "files/" + p.name + "/output", + }, nil +} + +func (p *publicBatchImageProvider) Get(context.Context, *BatchImageJob, *Account) (*BatchProviderStatus, error) { + return &BatchProviderStatus{InternalState: BatchProviderStateQueued}, nil +} + +func (p *publicBatchImageProvider) Cancel(context.Context, *BatchImageJob, *Account) error { + p.cancelCount++ + return p.cancelErr +} + +func (p *publicBatchImageProvider) OpenResult(context.Context, *BatchImageJob, *Account) (io.ReadCloser, string, error) { + return io.NopCloser(strings.NewReader(p.result)), "application/jsonl", nil +} + +func (p *publicBatchImageProvider) Cleanup(_ context.Context, _ *BatchImageJob, _ *Account, target CleanupTarget) error { + p.cleanupTargets = append(p.cleanupTargets, target) + return p.cleanupErr +} + +var _ BatchImageAccountSelectionRepository = (*publicBatchImageAccountRepo)(nil) +var _ BatchImageQueue = (*publicBatchImageQueue)(nil) +var _ BatchImageProvider = (*publicBatchImageProvider)(nil) + +type publicBatchImageGroupRepo struct { + groups map[int64]*Group +} + +func (r *publicBatchImageGroupRepo) GetByIDLite(_ context.Context, id int64) (*Group, error) { + if r != nil && r.groups != nil { + if group, ok := r.groups[id]; ok { + return group, nil + } + } + return nil, ErrGroupNotFound +} + +type publicBatchImageUserGroupRateRepo struct { + rates map[int64]*float64 +} + +func (r *publicBatchImageUserGroupRateRepo) GetByUserAndGroup(_ context.Context, _ int64, groupID int64) (*float64, error) { + if r != nil && r.rates != nil { + return r.rates[groupID], nil + } + return nil, nil +} + +var _ BatchImageGroupPricingRepository = (*publicBatchImageGroupRepo)(nil) +var _ BatchImageUserGroupRateRepository = (*publicBatchImageUserGroupRateRepo)(nil) diff --git a/backend/internal/service/batch_image_queue.go b/backend/internal/service/batch_image_queue.go new file mode 100644 index 0000000000..f5b25ccc08 --- /dev/null +++ b/backend/internal/service/batch_image_queue.go @@ -0,0 +1,64 @@ +package service + +import ( + "context" + "net/http" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +var ( + ErrBatchImageQueueEmpty = infraerrors.New(http.StatusNotFound, "BATCH_IMAGE_QUEUE_EMPTY", "batch image queue is empty") + ErrBatchImageAlreadyQueued = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_ALREADY_QUEUED", "batch image job is already queued") + ErrBatchImageLockNotAcquired = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_LOCK_NOT_ACQUIRED", "batch image job lock was not acquired") + ErrInvalidBatchImageQueuePayload = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_QUEUE_INVALID_PAYLOAD", "invalid batch image queue payload") +) + +type ReservedBatchImageJob struct { + BatchID string +} + +type BatchImageJobLock interface { + Release(ctx context.Context) error +} + +type BatchImageQueue interface { + Enqueue(ctx context.Context, batchID string) error + Reserve(ctx context.Context, blockTimeout time.Duration) (ReservedBatchImageJob, error) + RequeueAfter(ctx context.Context, batchID string, delay time.Duration) error + Ack(ctx context.Context, batchID string) error + Heartbeat(ctx context.Context, batchID string) error + MoveDueDelayedToReady(ctx context.Context, limit int) (int, error) + RecoverStaleActive(ctx context.Context, staleAfter time.Duration, limit int) (int, error) + TryAcquireJobLock(ctx context.Context, batchID string, ttl time.Duration) (BatchImageJobLock, bool, error) +} + +type BatchImageService struct { + repo BatchImageRepository + queue BatchImageQueue +} + +func NewBatchImageService(repo BatchImageRepository, queue BatchImageQueue) *BatchImageService { + return &BatchImageService{repo: repo, queue: queue} +} + +func (s *BatchImageService) EnqueueBatchImageJob(ctx context.Context, batchID string) error { + if !IsValidBatchImageID(batchID) { + return ErrInvalidBatchImageQueuePayload + } + if s == nil || s.queue == nil { + return infraerrors.New(http.StatusInternalServerError, "BATCH_IMAGE_QUEUE_NOT_CONFIGURED", "batch image queue is not configured") + } + if s.repo != nil { + if _, err := s.repo.GetBatchImageJobByBatchID(ctx, batchID); err != nil { + return err + } + } + return s.queue.Enqueue(ctx, batchID) +} + +func IsValidBatchImageID(batchID string) bool { + return strings.HasPrefix(batchID, "imgbatch_") && len(batchID) > len("imgbatch_") +} diff --git a/backend/internal/service/batch_image_settlement.go b/backend/internal/service/batch_image_settlement.go new file mode 100644 index 0000000000..26d2abb0fd --- /dev/null +++ b/backend/internal/service/batch_image_settlement.go @@ -0,0 +1,378 @@ +package service + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + batchImageSettlementRequestPrefix = "batch_image_settlement:" + batchImageSettlementRetryDelay = time.Minute + batchImageSettlementMaxRetries = 5 + batchImageCostEpsilon = 0.00000001 +) + +type BatchImagePricingResolver interface { + BatchImageUnitPrice(ctx context.Context, job *BatchImageJob) (float64, error) +} + +type BatchImageModelPricingResolver struct { + Resolver *ModelPricingResolver +} + +func (r *BatchImageModelPricingResolver) BatchImageUnitPrice(ctx context.Context, job *BatchImageJob) (float64, error) { + if r == nil || r.Resolver == nil || job == nil || strings.TrimSpace(job.Model) == "" { + return 0, ErrBatchImageSettlementPricingMissing + } + resolved := r.Resolver.Resolve(ctx, PricingInput{Model: job.Model}) + if resolved == nil { + return 0, ErrBatchImageSettlementPricingMissing + } + switch resolved.Mode { + case BillingModeImage, BillingModePerRequest: + if resolved.DefaultPerRequestPrice > 0 { + return resolved.DefaultPerRequestPrice, nil + } + if len(resolved.RequestTiers) == 1 && resolved.RequestTiers[0].PerRequestPrice != nil && *resolved.RequestTiers[0].PerRequestPrice >= 0 { + return *resolved.RequestTiers[0].PerRequestPrice, nil + } + case BillingModeToken: + if resolved.BasePricing != nil && (resolved.BasePricing.ImageOutputPriceExplicit || resolved.BasePricing.ImageOutputPricePerToken > 0) { + return resolved.BasePricing.ImageOutputPricePerToken, nil + } + } + return 0, ErrBatchImageSettlementPricingMissing +} + +type BatchImageSettlementService struct { + Repo BatchImageRepository + BillingRepo UsageBillingRepository + UsageLogRepo UsageLogRepository + Pricing BatchImagePricingResolver + AuthCache APIKeyAuthCacheInvalidator + Config *config.Config +} + +type BatchImageSettlementResult struct { + BatchID string + SuccessCount int + FailCount int + ActualCost float64 + ManifestHash string + RequestID string + AlreadySettled bool +} + +func (s *BatchImageSettlementService) Settle(ctx context.Context, batchID string) (*BatchImageSettlementResult, error) { + if s == nil || s.Repo == nil || s.BillingRepo == nil || s.Pricing == nil { + return nil, ErrBatchImageSettlementBillingFailed.WithCause(errors.New("batch image settlement service is not configured")) + } + job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID) + if err != nil { + return nil, err + } + + manifestHash := BuildBatchImageSettlementManifestHash(job) + result := &BatchImageSettlementResult{ + BatchID: job.BatchID, + SuccessCount: job.SuccessCount, + FailCount: job.FailCount, + ManifestHash: manifestHash, + RequestID: BatchImageCaptureRequestID(job.BatchID), + } + if job.ActualCost != nil { + result.ActualCost = *job.ActualCost + } + if job.Status == BatchImageJobStatusCompleted { + result.AlreadySettled = true + return result, nil + } + if job.Status != BatchImageJobStatusSettling { + return nil, ErrBatchImageSettlementInvalidStatus + } + if job.APIKeyID == nil || *job.APIKeyID <= 0 { + return nil, ErrBatchImageSettlementMissingAPIKeyID + } + if job.AccountID == nil || *job.AccountID <= 0 { + return nil, ErrBatchImageSettlementMissingAccountID + } + // 重试耗尽检查必须先于各类可重复失败的校验(counts/manifest/定价/超冻结), + // 否则这些错误路径会绕过耗尽出口,settling job 无限 requeue、冻结余额永不释放。 + if isBatchImageSettlementRetryExhausted(job) { + return nil, s.failExhaustedSettlement(ctx, job, "settlement retry limit reached: "+batchImageDerefString(job.LastErrorCode)) + } + if job.SuccessCount < 0 || job.FailCount < 0 || job.ItemCount < 0 || job.SuccessCount+job.FailCount > job.ItemCount { + if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_INVALID_COUNTS", + fmt.Sprintf("success=%d fail=%d item_count=%d", job.SuccessCount, job.FailCount, job.ItemCount)); failErr != nil { + return nil, failErr + } + return nil, ErrBatchImageSettlementInvalidCounts + } + if strings.TrimSpace(batchImageDerefString(job.ManifestHash)) != "" && batchImageDerefString(job.ManifestHash) != manifestHash { + if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_MANIFEST_CONFLICT", "manifest hash conflict"); failErr != nil { + return nil, failErr + } + return nil, ErrBatchImageSettlementManifestConflict + } + + unitPrice, err := s.settlementUnitPrice(ctx, job) + if err == nil && unitPrice < 0 { + err = ErrBatchImageSettlementPricingMissing + } + if err != nil { + if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_PRICING_MISSING", err.Error()); failErr != nil { + return nil, failErr + } + return nil, err + } + actualCost := float64(job.SuccessCount) * unitPrice + result.ActualCost = actualCost + holdAmount := job.EstimatedCost + if job.HoldAmount != nil { + holdAmount = *job.HoldAmount + } + if actualCost-holdAmount > batchImageCostEpsilon { + msg := fmt.Sprintf("actual cost %.10f exceeds held amount %.10f", actualCost, holdAmount) + if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_COST_EXCEEDS_HOLD", msg); failErr != nil { + return nil, failErr + } + return nil, ErrBatchImageSettlementCostExceedsHold + } + + if err := captureBatchImageBalanceHold(ctx, s.BillingRepo, job, actualCost, manifestHash); err != nil { + msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength) + if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_BILLING_FAILED", msg); failErr != nil { + return nil, failErr + } + return nil, err + } + s.invalidateAuthCache(ctx, job.UserID) + + now := time.Now() + outputExpiresAt := now.Add(s.outputRetentionAfterTerminal()) + if err := s.Repo.MarkBatchImageJobSettled(ctx, MarkBatchImageJobSettledParams{ + BatchID: job.BatchID, + ActualCost: actualCost, + ManifestHash: manifestHash, + Now: &now, + OutputExpiresAt: &outputExpiresAt, + EventPayload: map[string]any{ + "batch_id": job.BatchID, + "request_id": result.RequestID, + "success_count": job.SuccessCount, + "fail_count": job.FailCount, + "actual_cost": actualCost, + "manifest_hash": manifestHash, + }, + }); err != nil { + return nil, err + } + s.recordUsageLog(ctx, job, actualCost, result.RequestID, now) + + return result, nil +} + +// isBatchImageSettlementRetryExhausted 判断 settling job 是否已达重试上限。 +// 必须覆盖所有 SETTLEMENT_* 失败码(而非仅 SETTLEMENT_BILLING_FAILED), +// 否则 SETTLEMENT_COST_EXCEEDS_HOLD / SETTLEMENT_INVALID_COUNTS 等错误会无限 requeue。 +func isBatchImageSettlementRetryExhausted(job *BatchImageJob) bool { + return job != nil && + job.Status == BatchImageJobStatusSettling && + job.RetryCount >= batchImageSettlementMaxRetries && + strings.HasPrefix(batchImageDerefString(job.LastErrorCode), "SETTLEMENT_") +} + +// recordSettlementFailure 记录一次结算失败并递增 retry_count。 +// 重试达到上限时立即走耗尽出口(释放冻结余额并转 failed); +// 返回非 nil 时调用方应直接返回该错误。 +func (s *BatchImageSettlementService) recordSettlementFailure(ctx context.Context, job *BatchImageJob, code, message string) error { + retryCount, recordErr := s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, code, truncateBatchImageMessage(message, batchImageMaxErrorMessageLength)) + if recordErr != nil { + logger.L().Warn("batch_image.settlement_failure_record_failed", + zap.String("batch_id", job.BatchID), + zap.String("code", code), + zap.Error(recordErr), + ) + return nil + } + job.RetryCount = retryCount + job.LastErrorCode = &code + if retryCount >= batchImageSettlementMaxRetries { + return s.failExhaustedSettlement(ctx, job, message) + } + return nil +} + +func (s *BatchImageSettlementService) failExhaustedSettlement(ctx context.Context, job *BatchImageJob, message string) error { + if s == nil || s.Repo == nil { + return ErrBatchImageSettlementBillingFailed + } + // 释放指纹必须与其余所有释放点(processor/Cancel/recovery)一致地使用 RequestHash: + // 它们共享同一 request id,payloadHash 不同会触发 ErrUsageBillingRequestConflict, + // 导致后续 Cancel/重试永远失败、terminal job 变成毒消息。 + if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil { + msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength) + if _, recordErr := s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_RELEASE_FAILED", msg); recordErr != nil { + logger.L().Warn("batch_image.settlement_release_failure_record_failed", + zap.String("batch_id", job.BatchID), + zap.Error(recordErr), + ) + } + return ErrBatchImageSettlementBillingFailed.WithCause(err) + } + s.invalidateAuthCache(ctx, job.UserID) + msg := strings.TrimSpace(message) + if msg == "" { + msg = "settlement billing retry limit reached" + } + if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{ + ErrorCode: batchImageStringPtr("SETTLEMENT_BILLING_RETRY_EXHAUSTED"), + ErrorMessage: batchImageStringPtr(msg), + EventType: "settlement_retry_exhausted", + EventPayload: map[string]any{ + "batch_id": job.BatchID, + "retry_count": job.RetryCount, + }, + }); err != nil { + return err + } + return ErrBatchImageSettlementBillingFailed +} + +func (s *BatchImageSettlementService) recordUsageLog(ctx context.Context, job *BatchImageJob, actualCost float64, requestID string, createdAt time.Time) { + if s == nil || s.UsageLogRepo == nil || job == nil || job.APIKeyID == nil || job.AccountID == nil { + return + } + billingMode := string(BillingModeImage) + accountRateMultiplier := job.AccountRateMultiplier + inboundEndpoint := "/v1/images/batches" + upstreamEndpoint := "vertex:batchPredictionJobs" + imageSize := "1K" + usageLog := &UsageLog{ + UserID: job.UserID, + APIKeyID: *job.APIKeyID, + AccountID: *job.AccountID, + RequestID: strings.TrimSpace(requestID), + Model: job.Model, + RequestedModel: job.Model, + InboundEndpoint: &inboundEndpoint, + UpstreamEndpoint: &upstreamEndpoint, + ImageCount: job.SuccessCount, + ImageOutputCost: actualCost, + TotalCost: actualCost, + ActualCost: actualCost, + RateMultiplier: job.GroupRateMultiplier * job.BatchDiscountMultiplier, + AccountRateMultiplier: &accountRateMultiplier, + BillingType: BillingTypeBalance, + RequestType: RequestTypeSync, + BillingMode: &billingMode, + ImageSize: &imageSize, + CreatedAt: createdAt, + } + writeUsageLogBestEffort(ctx, s.UsageLogRepo, usageLog, "service.batch_image_settlement") +} + +func (s *BatchImageSettlementService) invalidateAuthCache(ctx context.Context, userID int64) { + if s != nil && s.AuthCache != nil && userID > 0 { + s.AuthCache.InvalidateAuthCacheByUserID(ctx, userID) + } +} + +func (s *BatchImageSettlementService) settlementUnitPrice(ctx context.Context, job *BatchImageJob) (float64, error) { + if job != nil && job.PricingSnapshotVersion >= 1 { + if job.BillableUnitPrice < 0 { + return 0, ErrBatchImageSettlementPricingMissing + } + return job.BillableUnitPrice, nil + } + unitPrice, err := s.Pricing.BatchImageUnitPrice(ctx, job) + if err != nil { + return 0, err + } + return unitPrice, nil +} + +func (s *BatchImageSettlementService) outputRetentionAfterTerminal() time.Duration { + if s != nil && s.Config != nil && s.Config.BatchImage.OutputRetentionAfterTerminalHours > 0 { + return time.Duration(s.Config.BatchImage.OutputRetentionAfterTerminalHours) * time.Hour + } + return 72 * time.Hour +} + +func BatchImageSettlementRequestID(batchID string) string { + return batchImageSettlementRequestPrefix + strings.TrimSpace(batchID) +} + +func BuildBatchImageSettlementManifestHash(job *BatchImageJob) string { + if job == nil { + return "" + } + parts := []string{ + strings.TrimSpace(job.BatchID), + strings.TrimSpace(job.Provider), + strings.TrimSpace(job.Model), + batchImageDerefString(job.ProviderJobName), + batchImageDerefString(job.ProviderOutputRef), + strconv.Itoa(job.SuccessCount), + strconv.Itoa(job.FailCount), + strconv.Itoa(job.ItemCount), + } + sum := sha256.Sum256([]byte(strings.Join(parts, "\x00"))) + return hex.EncodeToString(sum[:]) +} + +type BatchImagePipelineProcessor struct { + ProviderProcessor *BatchImageProviderProcessor + SettlementService *BatchImageSettlementService + RetryDelay time.Duration +} + +func (p *BatchImagePipelineProcessor) Process(ctx context.Context, batchID string) (BatchImageProcessResult, error) { + if p == nil || p.ProviderProcessor == nil { + return BatchImageProcessResult{}, errors.New("batch image pipeline processor is not configured") + } + job, err := p.ProviderProcessor.Repo.GetBatchImageJobByBatchID(ctx, batchID) + if err != nil { + return BatchImageProcessResult{}, err + } + if job.Status == BatchImageJobStatusSettling { + if p.SettlementService == nil { + return BatchImageProcessResult{Terminal: true}, nil + } + _, err := p.SettlementService.Settle(ctx, batchID) + if err != nil { + if errors.Is(err, ErrBatchImageSettlementBillingFailed) { + updated, getErr := p.ProviderProcessor.Repo.GetBatchImageJobByBatchID(ctx, batchID) + if getErr == nil && IsTerminalBatchImageJobStatus(updated.Status) { + return BatchImageProcessResult{Terminal: true}, nil + } + delay := p.RetryDelay + if delay <= 0 { + delay = batchImageSettlementRetryDelay + } + return BatchImageProcessResult{RequeueAfter: delay}, nil + } + return BatchImageProcessResult{}, err + } + return BatchImageProcessResult{Terminal: true}, nil + } + return p.ProviderProcessor.Process(ctx, batchID) +} + +func (r *BatchImageSettlementResult) String() string { + if r == nil { + return "" + } + return fmt.Sprintf("batch_id=%s success=%d fail=%d actual_cost=%0.10f already_settled=%t", + r.BatchID, r.SuccessCount, r.FailCount, r.ActualCost, r.AlreadySettled) +} diff --git a/backend/internal/service/batch_image_settlement_test.go b/backend/internal/service/batch_image_settlement_test.go new file mode 100644 index 0000000000..a1fe3e899d --- /dev/null +++ b/backend/internal/service/batch_image_settlement_test.go @@ -0,0 +1,507 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBatchImageSettlementService_SettlesAndChargesSuccessfulImagesOnly(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_settle") + job.SuccessCount = 3 + job.FailCount = 2 + job.ItemCount = 5 + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + result, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.Equal(t, 0.75, result.ActualCost) + require.Equal(t, BatchImageCaptureRequestID(job.BatchID), result.RequestID) + require.False(t, result.AlreadySettled) + require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status) + require.NotNil(t, repo.jobs[job.BatchID].ActualCost) + require.Equal(t, 0.75, *repo.jobs[job.BatchID].ActualCost) + require.NotEmpty(t, batchImageDerefString(repo.jobs[job.BatchID].ManifestHash)) + require.NotNil(t, repo.jobs[job.BatchID].SettledAt) + require.Len(t, billing.captures, 1) + require.Equal(t, int64(321), billing.captures[0].APIKeyID) + require.Equal(t, job.UserID, billing.captures[0].UserID) + require.Equal(t, job.BatchID, billing.captures[0].BatchID) + require.Equal(t, 0.75, billing.captures[0].ActualAmount) + require.Equal(t, 1.25, billing.captures[0].HoldAmount) + require.NotContains(t, fmt.Sprintf("%+v", billing.captures[0]), batchImageTestData) + require.NotContains(t, fmt.Sprintf("%+v", billing.captures[0]), "gs://") + require.NotContains(t, fmt.Sprintf("%+v", billing.captures[0]), "prompt") +} + +func TestBatchImageSettlementService_ZeroSuccessCanComplete(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_zero") + job.SuccessCount = 0 + job.FailCount = 4 + job.ItemCount = 4 + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + result, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.Equal(t, 0.0, result.ActualCost) + require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status) + require.Len(t, billing.captures, 1) + require.Equal(t, 0.0, billing.captures[0].ActualAmount) +} + +func TestBatchImageSettlementService_CompletedJobReturnsAlreadySettledWithoutBilling(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_done") + job.Status = BatchImageJobStatusCompleted + cost := 0.5 + job.ActualCost = &cost + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + result, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.True(t, result.AlreadySettled) + require.Equal(t, 0.5, result.ActualCost) + require.Empty(t, billing.captures) +} + +func TestBatchImageSettlementService_IdempotentAfterBillingCrash(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_crash") + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{alreadyApplied: map[string]bool{BatchImageCaptureRequestID(job.BatchID): true}} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + result, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.Equal(t, 0.5, result.ActualCost) + require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status) + require.Len(t, billing.captures, 1) +} + +func TestBatchImageSettlementService_ValidationErrors(t *testing.T) { + tests := []struct { + name string + mutate func(*BatchImageJob) + pricing BatchImagePricingResolver + want error + }{ + {name: "invalid_status", mutate: func(j *BatchImageJob) { j.Status = BatchImageJobStatusRunning }, want: ErrBatchImageSettlementInvalidStatus}, + {name: "negative_success_count", mutate: func(j *BatchImageJob) { j.SuccessCount = -1 }, want: ErrBatchImageSettlementInvalidCounts}, + {name: "negative_fail_count", mutate: func(j *BatchImageJob) { j.FailCount = -1 }, want: ErrBatchImageSettlementInvalidCounts}, + {name: "counts_exceed_item_count", mutate: func(j *BatchImageJob) { j.SuccessCount = 2; j.FailCount = 2; j.ItemCount = 3 }, want: ErrBatchImageSettlementInvalidCounts}, + {name: "missing_api_key", mutate: func(j *BatchImageJob) { j.APIKeyID = nil }, want: ErrBatchImageSettlementMissingAPIKeyID}, + {name: "missing_account", mutate: func(j *BatchImageJob) { j.AccountID = nil }, want: ErrBatchImageSettlementMissingAccountID}, + {name: "pricing_missing", pricing: &fakeBatchImagePricingResolver{err: ErrBatchImageSettlementPricingMissing}, want: ErrBatchImageSettlementPricingMissing}, + {name: "manifest_conflict", mutate: func(j *BatchImageJob) { v := "different"; j.ManifestHash = &v }, want: ErrBatchImageSettlementManifestConflict}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_" + tt.name) + if tt.mutate != nil { + tt.mutate(job) + } + repo.jobs[job.BatchID] = job + pricing := tt.pricing + if pricing == nil { + pricing = &fakeBatchImagePricingResolver{unitPrice: 0.25} + } + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: pricing} + + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, tt.want) + require.Empty(t, billing.captures) + require.NotEqual(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status) + }) + } +} + +func TestBatchImageSettlementService_CostExceedingHoldDoesNotCharge(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_cost_over_hold") + job.SuccessCount = 2 + job.FailCount = 0 + job.ItemCount = 2 + holdAmount := 0.5 + job.HoldAmount = &holdAmount + job.EstimatedCost = holdAmount + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}} + + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementCostExceedsHold) + require.Empty(t, billing.captures) + require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status) + require.Equal(t, "SETTLEMENT_COST_EXCEEDS_HOLD", batchImageDerefString(repo.jobs[job.BatchID].LastErrorCode)) +} + +func TestBatchImageSettlementService_UsesSubmittedPricingSnapshot(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_snapshot") + job.SuccessCount = 2 + job.FailCount = 0 + job.ItemCount = 2 + job.PricingSnapshotVersion = 1 + job.BaseUnitPrice = 0.25 + job.GroupRateMultiplier = 1 + job.AccountRateMultiplier = 1 + job.BatchDiscountMultiplier = 1 + job.HoldMultiplier = 1.1 + job.BillableUnitPrice = 0.25 + job.HoldUnitPrice = 0.275 + holdAmount := 0.55 + job.HoldAmount = &holdAmount + job.EstimatedCost = 0.5 + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}} + + result, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.InDelta(t, 0.5, result.ActualCost, 1e-12) + require.Len(t, billing.captures, 1) + require.InDelta(t, 0.5, billing.captures[0].ActualAmount, 1e-12) + require.InDelta(t, 0.55, billing.captures[0].HoldAmount, 1e-12) +} + +func TestBatchImageSettlementService_BillingFailureLeavesSettlingAndRecordsError(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_billing_fail") + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{err: errors.New("temporary billing timeout with gs://hidden-output")} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed) + require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status) + require.Equal(t, "SETTLEMENT_BILLING_FAILED", batchImageDerefString(repo.jobs[job.BatchID].LastErrorCode)) + require.Contains(t, batchImageDerefString(repo.jobs[job.BatchID].LastErrorMessage), "temporary billing timeout") + require.NotNil(t, billing.captures[0]) +} + +func TestBatchImagePipelineProcessor_SettlesQueuedSettlingJob(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_pipeline") + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + settlement := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + processor := &BatchImagePipelineProcessor{ + ProviderProcessor: &BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}}, + SettlementService: settlement, + } + + result, err := processor.Process(context.Background(), job.BatchID) + require.NoError(t, err) + require.True(t, result.Terminal) + require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status) + require.Len(t, billing.captures, 1) +} + +func TestBatchImagePipelineProcessor_RequeuesTransientSettlementFailure(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_pipeline_retry") + repo.jobs[job.BatchID] = job + settlement := &BatchImageSettlementService{Repo: repo, BillingRepo: &fakeBatchImageBillingRepo{err: errors.New("temporary")}, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + processor := &BatchImagePipelineProcessor{ + ProviderProcessor: &BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}}, + SettlementService: settlement, + } + + result, err := processor.Process(context.Background(), job.BatchID) + require.NoError(t, err) + require.False(t, result.Terminal) + require.Equal(t, batchImageSettlementRetryDelay, result.RequeueAfter) + require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status) +} + +func TestBatchImagePipelineProcessor_FailsAndReleasesAfterSettlementRetryLimit(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_pipeline_retry_exhausted") + job.RetryCount = batchImageSettlementMaxRetries - 1 + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{captureErr: errors.New("temporary billing timeout")} + settlement := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + processor := &BatchImagePipelineProcessor{ + ProviderProcessor: &BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}}, + SettlementService: settlement, + } + + result, err := processor.Process(context.Background(), job.BatchID) + require.NoError(t, err) + require.True(t, result.Terminal) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status) + require.Equal(t, "SETTLEMENT_BILLING_RETRY_EXHAUSTED", batchImageDerefString(repo.jobs[job.BatchID].LastErrorCode)) + require.Len(t, billing.captures, 1) + require.Len(t, billing.releases, 1) + require.Equal(t, BatchImageReleaseRequestID(job.BatchID), billing.releases[0].RequestID) +} + +func TestBatchImageSettlementRetryExhaustedReleaseIsIdempotentAfterTransitionFailure(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_retry_exhausted_transition_fail") + job.RetryCount = batchImageSettlementMaxRetries + job.LastErrorCode = batchImageStringPtr("SETTLEMENT_BILLING_FAILED") + repo.jobs[job.BatchID] = job + repo.transitionErr = errors.New("temporary transition failure") + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorContains(t, err, "temporary transition failure") + require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status) + require.Len(t, billing.releases, 1) + require.Len(t, billing.seen, 1) + + repo.transitionErr = nil + _, err = svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status) + require.Len(t, billing.releases, 2) + require.Equal(t, billing.releases[0].RequestID, billing.releases[1].RequestID) + require.Len(t, billing.seen, 1) +} + +func TestBatchImageSettlementService_CostExceedsHoldExhaustsAndReleases(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_over_hold_exhausted") + job.SuccessCount = 2 + job.FailCount = 0 + job.ItemCount = 2 + holdAmount := 0.5 + job.HoldAmount = &holdAmount + job.EstimatedCost = holdAmount + requestHash := "request-hash-over-hold" + job.RequestHash = &requestHash + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}} + + // 前 N-1 次:记录失败并返回错误(等待 worker 重试)。 + for i := 0; i < batchImageSettlementMaxRetries-1; i++ { + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementCostExceedsHold) + require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status) + } + // 达到上限:必须走耗尽出口释放冻结并转 failed,而不是无限 requeue。 + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status) + require.Empty(t, billing.captures) + require.Len(t, billing.releases, 1) + require.Equal(t, BatchImageReleaseRequestID(job.BatchID), billing.releases[0].RequestID) + // 释放指纹必须与 processor/Cancel/recovery 一致地使用 RequestHash, + // 否则共享同一 request id 的后续释放会命中指纹冲突(毒消息)。 + require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash) +} + +func TestBatchImageSettlementService_InvalidCountsExhaustsAndReleases(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("imgbatch_bad_counts_exhausted") + job.SuccessCount = 2 + job.FailCount = 2 + job.ItemCount = 3 + requestHash := "request-hash-bad-counts" + job.RequestHash = &requestHash + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + for i := 0; i < batchImageSettlementMaxRetries-1; i++ { + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementInvalidCounts) + } + _, err := svc.Settle(context.Background(), job.BatchID) + require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed) + require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status) + require.Empty(t, billing.captures) + require.Len(t, billing.releases, 1) + require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash) +} + +func TestReleaseBatchImageBalanceHold_TreatsFingerprintConflictAsReleased(t *testing.T) { + job := testSettlingBatchImageJob("imgbatch_release_conflict") + // 历史版本用 manifestHash 释放过一次:同一 request id 再以 RequestHash + // 释放会命中指纹冲突。资金已归还,必须视为幂等成功而非毒消息。 + billing := &fakeBatchImageBillingRepo{releaseErr: ErrUsageBillingRequestConflict} + err := releaseBatchImageBalanceHold(context.Background(), billing, job, "request-hash") + require.NoError(t, err) + require.Len(t, billing.releases, 1) +} + +func TestBatchImageSettlementManifestHash(t *testing.T) { + job := testSettlingBatchImageJob("imgbatch_hash") + first := BuildBatchImageSettlementManifestHash(job) + job.CreatedAt = job.CreatedAt.AddDate(0, 0, 1) + job.UpdatedAt = job.UpdatedAt.AddDate(0, 0, 1) + require.Equal(t, first, BuildBatchImageSettlementManifestHash(job)) + + job.SuccessCount++ + require.NotEqual(t, first, BuildBatchImageSettlementManifestHash(job)) + + job.SuccessCount-- + promptOrBase64 := first + " prompt " + batchImageTestData + require.NotContains(t, BuildBatchImageSettlementManifestHash(job), promptOrBase64) +} + +func TestBatchImageSettlementBillingRequestIDs(t *testing.T) { + repo := newFakeBatchImageRepository() + first := testSettlingBatchImageJob("imgbatch_unique_1") + second := testSettlingBatchImageJob("imgbatch_unique_2") + repo.jobs[first.BatchID] = first + repo.jobs[second.BatchID] = second + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + + _, err := svc.Settle(context.Background(), first.BatchID) + require.NoError(t, err) + _, err = svc.Settle(context.Background(), first.BatchID) + require.NoError(t, err) + _, err = svc.Settle(context.Background(), second.BatchID) + require.NoError(t, err) + + require.Len(t, billing.captures, 2) + require.Equal(t, BatchImageCaptureRequestID(first.BatchID), billing.captures[0].RequestID) + require.Equal(t, BatchImageCaptureRequestID(second.BatchID), billing.captures[1].RequestID) + require.NotEqual(t, billing.captures[0].RequestID, billing.captures[1].RequestID) + require.Len(t, billing.seen, 2) +} + +func testSettlingBatchImageJob(batchID string) *BatchImageJob { + apiKeyID := int64(321) + accountID := int64(654) + providerJobName := "providers/job" + outputRef := "files/output" + holdAmount := 1.25 + holdID := BatchImageHoldRequestID(batchID) + return &BatchImageJob{ + BatchID: batchID, + UserID: 123, + APIKeyID: &apiKeyID, + AccountID: &accountID, + Provider: BatchImageProviderGeminiAPI, + Model: "gemini-image", + Status: BatchImageJobStatusSettling, + ProviderJobName: &providerJobName, + ProviderOutputRef: &outputRef, + ItemCount: 3, + SuccessCount: 2, + FailCount: 1, + EstimatedCost: holdAmount, + HoldAmount: &holdAmount, + HoldID: &holdID, + } +} + +type fakeBatchImagePricingResolver struct { + unitPrice float64 + missingModels map[string]bool + err error +} + +func (r *fakeBatchImagePricingResolver) BatchImageUnitPrice(_ context.Context, job *BatchImageJob) (float64, error) { + if r.err != nil { + return 0, r.err + } + if job != nil && r.missingModels[job.Model] { + return 0, ErrBatchImageSettlementPricingMissing + } + return r.unitPrice, nil +} + +type fakeBatchImageBillingRepo struct { + commands []*UsageBillingCommand + reserves []*BatchImageBalanceHoldCommand + captures []*BatchImageBalanceHoldCommand + releases []*BatchImageBalanceHoldCommand + seen map[string]struct{} + alreadyApplied map[string]bool + err error + reserveErr error + captureErr error + releaseErr error +} + +func (r *fakeBatchImageBillingRepo) Apply(_ context.Context, cmd *UsageBillingCommand) (*UsageBillingApplyResult, error) { + if r.seen == nil { + r.seen = make(map[string]struct{}) + } + if r.err != nil { + r.commands = append(r.commands, cmd) + return nil, r.err + } + if cmd != nil { + cmd.Normalize() + if _, ok := r.seen[cmd.RequestID]; ok || r.alreadyApplied[cmd.RequestID] { + r.commands = append(r.commands, cmd) + return &UsageBillingApplyResult{Applied: false}, nil + } + r.seen[cmd.RequestID] = struct{}{} + } + r.commands = append(r.commands, cmd) + return &UsageBillingApplyResult{Applied: true}, nil +} + +func (r *fakeBatchImageBillingRepo) ReserveBatchImageBalance(_ context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) { + if r.reserveErr != nil { + r.reserves = append(r.reserves, cmd) + return nil, r.reserveErr + } + return r.applyHold(cmd, &r.reserves) +} + +func (r *fakeBatchImageBillingRepo) CaptureBatchImageBalance(_ context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) { + if r.captureErr != nil { + r.captures = append(r.captures, cmd) + return nil, r.captureErr + } + return r.applyHold(cmd, &r.captures) +} + +func (r *fakeBatchImageBillingRepo) ReleaseBatchImageBalance(_ context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) { + if r.releaseErr != nil { + r.releases = append(r.releases, cmd) + return nil, r.releaseErr + } + return r.applyHold(cmd, &r.releases) +} + +func (r *fakeBatchImageBillingRepo) applyHold(cmd *BatchImageBalanceHoldCommand, calls *[]*BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) { + if r.seen == nil { + r.seen = make(map[string]struct{}) + } + if r.err != nil { + *calls = append(*calls, cmd) + return nil, r.err + } + if cmd != nil { + cmd.Normalize() + if _, ok := r.seen[cmd.RequestID]; ok || r.alreadyApplied[cmd.RequestID] { + *calls = append(*calls, cmd) + return &BatchImageBalanceHoldResult{Applied: false}, nil + } + r.seen[cmd.RequestID] = struct{}{} + } + *calls = append(*calls, cmd) + return &BatchImageBalanceHoldResult{Applied: true}, nil +} + +var _ UsageBillingRepository = (*fakeBatchImageBillingRepo)(nil) +var _ BatchImagePricingResolver = (*fakeBatchImagePricingResolver)(nil) +var _ = strings.TrimSpace diff --git a/backend/internal/service/batch_image_test.go b/backend/internal/service/batch_image_test.go new file mode 100644 index 0000000000..dca17ec85f --- /dev/null +++ b/backend/internal/service/batch_image_test.go @@ -0,0 +1,63 @@ +//go:build unit + +package service + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCanTransitionBatchImageJob(t *testing.T) { + tests := []struct { + name string + from string + to string + want bool + }{ + {name: "created_to_uploading", from: BatchImageJobStatusCreated, to: BatchImageJobStatusUploading, want: true}, + {name: "uploading_to_submitted", from: BatchImageJobStatusUploading, to: BatchImageJobStatusSubmitted, want: true}, + {name: "submitted_to_running", from: BatchImageJobStatusSubmitted, to: BatchImageJobStatusRunning, want: true}, + {name: "running_self_poll", from: BatchImageJobStatusRunning, to: BatchImageJobStatusRunning, want: true}, + {name: "running_to_indexing", from: BatchImageJobStatusRunning, to: BatchImageJobStatusIndexing, want: true}, + {name: "indexing_to_settling", from: BatchImageJobStatusIndexing, to: BatchImageJobStatusSettling, want: true}, + {name: "settling_to_completed", from: BatchImageJobStatusSettling, to: BatchImageJobStatusCompleted, want: true}, + {name: "submitted_to_cancelled", from: BatchImageJobStatusSubmitted, to: BatchImageJobStatusCancelled, want: true}, + {name: "non_terminal_to_failed", from: BatchImageJobStatusCreated, to: BatchImageJobStatusFailed, want: true}, + {name: "completed_to_output_deleted", from: BatchImageJobStatusCompleted, to: BatchImageJobStatusOutputDeleted, want: true}, + {name: "failed_to_output_deleted", from: BatchImageJobStatusFailed, to: BatchImageJobStatusOutputDeleted, want: true}, + {name: "cancelled_to_output_deleted", from: BatchImageJobStatusCancelled, to: BatchImageJobStatusOutputDeleted, want: true}, + {name: "created_to_running_invalid", from: BatchImageJobStatusCreated, to: BatchImageJobStatusRunning, want: false}, + {name: "completed_to_running_invalid", from: BatchImageJobStatusCompleted, to: BatchImageJobStatusRunning, want: false}, + {name: "output_deleted_to_failed_invalid", from: BatchImageJobStatusOutputDeleted, to: BatchImageJobStatusFailed, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, CanTransitionBatchImageJob(tt.from, tt.to)) + }) + } +} + +func TestIsTerminalBatchImageJobStatus(t *testing.T) { + require.True(t, IsTerminalBatchImageJobStatus(BatchImageJobStatusCompleted)) + require.True(t, IsTerminalBatchImageJobStatus(BatchImageJobStatusFailed)) + require.True(t, IsTerminalBatchImageJobStatus(BatchImageJobStatusCancelled)) + require.True(t, IsTerminalBatchImageJobStatus(BatchImageJobStatusOutputDeleted)) + require.False(t, IsTerminalBatchImageJobStatus(BatchImageJobStatusRunning)) +} + +func TestIsSupportedBatchImageProvider(t *testing.T) { + require.True(t, IsSupportedBatchImageProvider(BatchImageProviderGeminiAPI)) + require.True(t, IsSupportedBatchImageProvider(BatchImageProviderVertex)) + require.False(t, IsSupportedBatchImageProvider("gemini_oauth")) + require.False(t, IsSupportedBatchImageProvider("")) +} + +func TestNewBatchImageID(t *testing.T) { + id, err := NewBatchImageID() + require.NoError(t, err) + require.True(t, strings.HasPrefix(id, "imgbatch_")) + require.Len(t, id, len("imgbatch_")+32) +} diff --git a/backend/internal/service/batch_image_worker.go b/backend/internal/service/batch_image_worker.go new file mode 100644 index 0000000000..f090a67a52 --- /dev/null +++ b/backend/internal/service/batch_image_worker.go @@ -0,0 +1,286 @@ +package service + +import ( + "context" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" +) + +const ( + defaultBatchImageWorkerLockTTL = 5 * time.Minute + defaultBatchImageWorkerLockConflictDelay = 5 * time.Second + defaultBatchImageWorkerErrorRetryDelay = time.Minute + defaultBatchImageWorkerRequeueDelay = 30 * time.Second + defaultBatchImageWorkerDelayedPollInterval = 5 * time.Second + defaultBatchImageWorkerRecoveryInterval = 5 * time.Minute + defaultBatchImageWorkerStaleActiveAfter = 10 * time.Minute + defaultBatchImageWorkerDelayedMoveLimit = 100 + defaultBatchImageWorkerRecoverLimit = 100 + defaultBatchImageWorkerErrorBackoff = time.Second + defaultBatchImageWorkerReserveBlockTimeout = 5 * time.Second +) + +type BatchImageProcessor interface { + Process(ctx context.Context, batchID string) (BatchImageProcessResult, error) +} + +type BatchImageProcessResult struct { + RequeueAfter time.Duration + Terminal bool +} + +type BatchImageWorkerOptions struct { + ReserveBlockTimeout time.Duration + JobLockTTL time.Duration + LockConflictDelay time.Duration + DefaultRequeueDelay time.Duration + ErrorRetryDelay time.Duration + ErrorBackoff time.Duration + DelayedPollInterval time.Duration + RecoveryInterval time.Duration + StaleActiveAfter time.Duration + DelayedMoveLimit int + RecoverLimit int +} + +type BatchImageWorker struct { + queue BatchImageQueue + processor BatchImageProcessor + opts BatchImageWorkerOptions +} + +func NewBatchImageWorker(queue BatchImageQueue, processor BatchImageProcessor, opts BatchImageWorkerOptions) *BatchImageWorker { + return &BatchImageWorker{ + queue: queue, + processor: processor, + opts: normalizeBatchImageWorkerOptions(opts), + } +} + +func NewBatchImageWorkerOptionsFromConfig(cfg *config.Config) BatchImageWorkerOptions { + if cfg == nil { + return normalizeBatchImageWorkerOptions(BatchImageWorkerOptions{}) + } + return normalizeBatchImageWorkerOptions(BatchImageWorkerOptions{ + JobLockTTL: time.Duration(cfg.BatchImage.JobLockTTLSeconds) * time.Second, + LockConflictDelay: time.Duration(cfg.BatchImage.LockConflictDelaySeconds) * time.Second, + DefaultRequeueDelay: time.Duration(cfg.BatchImage.DefaultRequeueDelaySeconds) * time.Second, + ErrorRetryDelay: time.Duration(cfg.BatchImage.ErrorRetryDelaySeconds) * time.Second, + DelayedPollInterval: time.Duration(cfg.BatchImage.DelayedMoverIntervalSeconds) * time.Second, + RecoveryInterval: time.Duration(cfg.BatchImage.RecoveryIntervalSeconds) * time.Second, + StaleActiveAfter: time.Duration(cfg.BatchImage.StaleActiveAfterSeconds) * time.Second, + DelayedMoveLimit: cfg.BatchImage.DelayedMoveLimit, + RecoverLimit: cfg.BatchImage.RecoverLimit, + }) +} + +func normalizeBatchImageWorkerOptions(opts BatchImageWorkerOptions) BatchImageWorkerOptions { + if opts.ReserveBlockTimeout <= 0 { + opts.ReserveBlockTimeout = defaultBatchImageWorkerReserveBlockTimeout + } + if opts.JobLockTTL <= 0 { + opts.JobLockTTL = defaultBatchImageWorkerLockTTL + } + if opts.LockConflictDelay <= 0 { + opts.LockConflictDelay = defaultBatchImageWorkerLockConflictDelay + } + if opts.DefaultRequeueDelay <= 0 { + opts.DefaultRequeueDelay = defaultBatchImageWorkerRequeueDelay + } + if opts.ErrorRetryDelay <= 0 { + opts.ErrorRetryDelay = defaultBatchImageWorkerErrorRetryDelay + } + if opts.ErrorBackoff <= 0 { + opts.ErrorBackoff = defaultBatchImageWorkerErrorBackoff + } + if opts.DelayedPollInterval <= 0 { + opts.DelayedPollInterval = defaultBatchImageWorkerDelayedPollInterval + } + if opts.RecoveryInterval <= 0 { + opts.RecoveryInterval = defaultBatchImageWorkerRecoveryInterval + } + if opts.StaleActiveAfter <= 0 { + opts.StaleActiveAfter = defaultBatchImageWorkerStaleActiveAfter + } + if opts.DelayedMoveLimit <= 0 { + opts.DelayedMoveLimit = defaultBatchImageWorkerDelayedMoveLimit + } + if opts.RecoverLimit <= 0 { + opts.RecoverLimit = defaultBatchImageWorkerRecoverLimit + } + return opts +} + +func (w *BatchImageWorker) Run(ctx context.Context) { + if w == nil { + return + } + for { + if err := ctx.Err(); err != nil { + return + } + if err := w.RunOnce(ctx); err != nil && ctx.Err() == nil { + sleepOrDone(ctx, w.opts.ErrorBackoff) + } + } +} + +func (w *BatchImageWorker) RunOnce(ctx context.Context) error { + if w == nil || w.queue == nil || w.processor == nil { + return nil + } + + reserved, err := w.queue.Reserve(ctx, w.opts.ReserveBlockTimeout) + if errors.Is(err, ErrBatchImageQueueEmpty) { + return nil + } + if err != nil { + return err + } + + lock, ok, err := w.queue.TryAcquireJobLock(ctx, reserved.BatchID, w.opts.JobLockTTL) + if err != nil { + if requeueErr := w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.LockConflictDelay); requeueErr != nil { + return requeueErr + } + return err + } + if !ok { + // 锁被其他实例持有:按冲突延迟重新入队。直接丢弃会让 job 滞留在 + // active zset,最早要等 StaleActiveAfter 才被恢复,造成分钟级停摆。 + return w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.LockConflictDelay) + } + defer func() { + _ = lock.Release(ctx) + }() + + // 处理期间持续心跳:刷新 active zset 时间戳防止 stale 恢复把在处理的 + // job 重投给其他 worker,并对支持续期的锁实现延长锁 TTL。 + hbStop := make(chan struct{}) + hbDone := make(chan struct{}) + go w.runJobHeartbeat(ctx, reserved.BatchID, lock, hbStop, hbDone) + + result, err := w.processor.Process(ctx, reserved.BatchID) + close(hbStop) + <-hbDone + if err != nil { + logger.L().Warn("batch_image.worker_process_failed", + zap.String("batch_id", reserved.BatchID), + zap.Error(err), + ) + return w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.ErrorRetryDelay) + } + if result.Terminal { + return w.queue.Ack(ctx, reserved.BatchID) + } + delay := result.RequeueAfter + if delay <= 0 { + delay = w.opts.DefaultRequeueDelay + } + return w.queue.RequeueAfter(ctx, reserved.BatchID, delay) +} + +// BatchImageJobLockRefresher 是可选的锁续期能力;由具体锁实现按需提供。 +type BatchImageJobLockRefresher interface { + Refresh(ctx context.Context, ttl time.Duration) error +} + +func (w *BatchImageWorker) heartbeatInterval() time.Duration { + interval := w.opts.JobLockTTL + if w.opts.StaleActiveAfter < interval { + interval = w.opts.StaleActiveAfter + } + interval /= 3 + if interval < time.Second { + interval = time.Second + } + return interval +} + +func (w *BatchImageWorker) runJobHeartbeat(ctx context.Context, batchID string, lock BatchImageJobLock, stop <-chan struct{}, done chan<- struct{}) { + defer close(done) + ticker := time.NewTicker(w.heartbeatInterval()) + defer ticker.Stop() + for { + select { + case <-stop: + return + case <-ctx.Done(): + return + case <-ticker.C: + if err := w.queue.Heartbeat(ctx, batchID); err != nil && ctx.Err() == nil { + logger.L().Warn("batch_image.worker_heartbeat_failed", + zap.String("batch_id", batchID), + zap.Error(err), + ) + } + if refresher, ok := lock.(BatchImageJobLockRefresher); ok { + if err := refresher.Refresh(ctx, w.opts.JobLockTTL); err != nil && ctx.Err() == nil { + logger.L().Warn("batch_image.worker_lock_refresh_failed", + zap.String("batch_id", batchID), + zap.Error(err), + ) + } + } + } + } +} + +func (w *BatchImageWorker) MoveDueDelayedOnce(ctx context.Context) (int, error) { + if w == nil || w.queue == nil { + return 0, nil + } + return w.queue.MoveDueDelayedToReady(ctx, w.opts.DelayedMoveLimit) +} + +func (w *BatchImageWorker) RunDelayedMover(ctx context.Context) { + if w == nil { + return + } + for { + if err := ctx.Err(); err != nil { + return + } + moved, _ := w.MoveDueDelayedOnce(ctx) + if moved > 0 { + continue + } + sleepOrDone(ctx, w.opts.DelayedPollInterval) + } +} + +func (w *BatchImageWorker) RecoverStaleActiveOnce(ctx context.Context) (int, error) { + if w == nil || w.queue == nil { + return 0, nil + } + return w.queue.RecoverStaleActive(ctx, w.opts.StaleActiveAfter, w.opts.RecoverLimit) +} + +func (w *BatchImageWorker) RunStaleActiveRecovery(ctx context.Context) { + if w == nil { + return + } + for { + if err := ctx.Err(); err != nil { + return + } + _, _ = w.RecoverStaleActiveOnce(ctx) + sleepOrDone(ctx, w.opts.RecoveryInterval) + } +} + +func sleepOrDone(ctx context.Context, d time.Duration) { + if d <= 0 { + return + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + case <-timer.C: + } +} diff --git a/backend/internal/service/batch_image_worker_runtime.go b/backend/internal/service/batch_image_worker_runtime.go new file mode 100644 index 0000000000..f5373cb6d6 --- /dev/null +++ b/backend/internal/service/batch_image_worker_runtime.go @@ -0,0 +1,143 @@ +package service + +import ( + "context" + "sync" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +type BatchImageWorkerRuntime struct { + worker *BatchImageWorker + billingRecovery *BatchImageBillingRecoveryService + cfg *config.Config + + mu sync.Mutex + cancel context.CancelFunc + done chan struct{} +} + +func NewBatchImageWorkerRuntime(worker *BatchImageWorker, cfg *config.Config) *BatchImageWorkerRuntime { + return &BatchImageWorkerRuntime{worker: worker, cfg: cfg} +} + +func ProvideBatchImageWorkerRuntime( + repo BatchImageRepository, + accountRepo AccountRepository, + queue BatchImageQueue, + billingRepo UsageBillingRepository, + usageLogRepo UsageLogRepository, + pricing *BatchImageModelPricingResolver, + authCache APIKeyAuthCacheInvalidator, + cfg *config.Config, +) *BatchImageWorkerRuntime { + processor := &BatchImagePipelineProcessor{ + ProviderProcessor: &BatchImageProviderProcessor{ + Repo: repo, + ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg), + AccountResolver: &BatchImageAccountRepositoryResolver{Repo: accountRepo}, + BillingRepo: billingRepo, + AuthCache: authCache, + }, + SettlementService: &BatchImageSettlementService{ + Repo: repo, + BillingRepo: billingRepo, + UsageLogRepo: usageLogRepo, + Pricing: pricing, + AuthCache: authCache, + Config: cfg, + }, + } + runtime := NewBatchImageWorkerRuntime(NewBatchImageWorker(queue, processor, NewBatchImageWorkerOptionsFromConfig(cfg)), cfg) + runtime.billingRecovery = &BatchImageBillingRecoveryService{ + Repo: repo, + Billing: billingRepo, + AuthCache: authCache, + Queue: queue, + StaleAfter: NewBatchImageWorkerOptionsFromConfig(cfg).StaleActiveAfter, + Limit: NewBatchImageWorkerOptionsFromConfig(cfg).RecoverLimit, + } + runtime.Start() + return runtime +} + +func (r *BatchImageWorkerRuntime) Start() { + if r == nil || r.worker == nil || r.cfg == nil || !r.cfg.BatchImage.QueueEnabled { + return + } + r.mu.Lock() + defer r.mu.Unlock() + if r.cancel != nil { + return + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + r.cancel = cancel + r.done = done + + var wg sync.WaitGroup + wg.Add(4) + go func() { + defer wg.Done() + r.worker.Run(ctx) + }() + go func() { + defer wg.Done() + r.worker.RunDelayedMover(ctx) + }() + go func() { + defer wg.Done() + r.worker.RunStaleActiveRecovery(ctx) + }() + go func() { + defer wg.Done() + r.runBillingRecovery(ctx) + }() + go func() { + wg.Wait() + close(done) + }() +} + +func (r *BatchImageWorkerRuntime) runBillingRecovery(ctx context.Context) { + if r == nil || r.worker == nil || r.billingRecovery == nil { + return + } + interval := r.worker.opts.RecoveryInterval + for { + if err := ctx.Err(); err != nil { + return + } + _, _ = r.billingRecovery.ReleaseStaleUnsubmittedOnce(ctx) + sleepOrDone(ctx, interval) + } +} + +func (r *BatchImageWorkerRuntime) Stop() { + if r == nil { + return + } + r.mu.Lock() + cancel := r.cancel + done := r.done + r.cancel = nil + r.done = nil + r.mu.Unlock() + + if cancel != nil { + cancel() + } + if done != nil { + <-done + } +} + +func (r *BatchImageWorkerRuntime) Running() bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + return r.cancel != nil +} diff --git a/backend/internal/service/batch_image_worker_runtime_redis_test.go b/backend/internal/service/batch_image_worker_runtime_redis_test.go new file mode 100644 index 0000000000..8905411b94 --- /dev/null +++ b/backend/internal/service/batch_image_worker_runtime_redis_test.go @@ -0,0 +1,58 @@ +//go:build unit + +package service_test + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/repository" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestBatchImageWorkerRuntime_StartupDoesNotCreateRedisBatchImageKeys(t *testing.T) { + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { + _ = rdb.Close() + }) + + cfg := &config.Config{BatchImage: config.BatchImageConfig{ + QueueEnabled: true, + QueueReadyKey: "batch_image:queue:ready", + QueueDelayedKey: "batch_image:queue:delayed", + QueueActiveKey: "batch_image:queue:active", + InflightKeyPrefix: "batch_image:queue:inflight:", + LockKeyPrefix: "batch_image:queue:lock:", + InflightTTLSeconds: 60, + JobLockTTLSeconds: 60, + DelayedMoverIntervalSeconds: 60, + RecoveryIntervalSeconds: 60, + StaleActiveAfterSeconds: 60, + DelayedMoveLimit: 10, + RecoverLimit: 10, + }} + queue := repository.NewBatchImageQueue(rdb, cfg) + worker := service.NewBatchImageWorker(queue, noopBatchImageProcessor{}, service.NewBatchImageWorkerOptionsFromConfig(cfg)) + runtime := service.NewBatchImageWorkerRuntime(worker, cfg) + + runtime.Start() + require.Eventually(t, runtime.Running, time.Second, 10*time.Millisecond) + runtime.Stop() + + for _, key := range mr.Keys() { + require.False(t, strings.HasPrefix(key, "batch_image:"), "unexpected Redis key created at startup: %s", key) + } +} + +type noopBatchImageProcessor struct{} + +func (noopBatchImageProcessor) Process(context.Context, string) (service.BatchImageProcessResult, error) { + return service.BatchImageProcessResult{}, nil +} diff --git a/backend/internal/service/batch_image_worker_runtime_test.go b/backend/internal/service/batch_image_worker_runtime_test.go new file mode 100644 index 0000000000..8e397cef77 --- /dev/null +++ b/backend/internal/service/batch_image_worker_runtime_test.go @@ -0,0 +1,87 @@ +//go:build unit + +package service + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestBatchImageWorkerRuntime_QueueDisabledDoesNotStart(t *testing.T) { + queue := &blockingBatchImageRuntimeQueue{} + runtime := NewBatchImageWorkerRuntime( + NewBatchImageWorker(queue, &fakeBatchImageProcessor{}, BatchImageWorkerOptions{}), + &config.Config{BatchImage: config.BatchImageConfig{QueueEnabled: false}}, + ) + + runtime.Start() + + require.False(t, runtime.Running()) + require.Zero(t, queue.reserveCalls.Load()) + require.NotPanics(t, runtime.Stop) +} + +func TestBatchImageWorkerRuntime_QueueEnabledStartsAndStops(t *testing.T) { + queue := &blockingBatchImageRuntimeQueue{} + processor := &fakeBatchImageProcessor{} + runtime := NewBatchImageWorkerRuntime( + NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{ + DelayedPollInterval: time.Hour, + RecoveryInterval: time.Hour, + }), + &config.Config{BatchImage: config.BatchImageConfig{QueueEnabled: true}}, + ) + + runtime.Start() + + require.Eventually(t, func() bool { + return runtime.Running() && queue.reserveCalls.Load() > 0 + }, time.Second, 10*time.Millisecond) + require.Empty(t, processor.processed) + require.NotPanics(t, runtime.Stop) + require.False(t, runtime.Running()) + require.NotPanics(t, runtime.Stop) +} + +type blockingBatchImageRuntimeQueue struct { + reserveCalls atomic.Int64 +} + +func (q *blockingBatchImageRuntimeQueue) Enqueue(context.Context, string) error { + return nil +} + +func (q *blockingBatchImageRuntimeQueue) Reserve(ctx context.Context, _ time.Duration) (ReservedBatchImageJob, error) { + q.reserveCalls.Add(1) + <-ctx.Done() + return ReservedBatchImageJob{}, ctx.Err() +} + +func (q *blockingBatchImageRuntimeQueue) RequeueAfter(context.Context, string, time.Duration) error { + return nil +} + +func (q *blockingBatchImageRuntimeQueue) Ack(context.Context, string) error { + return nil +} + +func (q *blockingBatchImageRuntimeQueue) Heartbeat(context.Context, string) error { + return nil +} + +func (q *blockingBatchImageRuntimeQueue) MoveDueDelayedToReady(context.Context, int) (int, error) { + return 0, nil +} + +func (q *blockingBatchImageRuntimeQueue) RecoverStaleActive(context.Context, time.Duration, int) (int, error) { + return 0, nil +} + +func (q *blockingBatchImageRuntimeQueue) TryAcquireJobLock(context.Context, string, time.Duration) (BatchImageJobLock, bool, error) { + return nil, false, nil +} diff --git a/backend/internal/service/batch_image_worker_test.go b/backend/internal/service/batch_image_worker_test.go new file mode 100644 index 0000000000..016df7e713 --- /dev/null +++ b/backend/internal/service/batch_image_worker_test.go @@ -0,0 +1,157 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestBatchImageWorker_ProcessesJobOnce(t *testing.T) { + queue := newFakeBatchImageQueue("imgbatch_worker_once") + processor := &fakeBatchImageProcessor{} + worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{ReserveBlockTimeout: time.Millisecond}) + + require.NoError(t, worker.RunOnce(context.Background())) + require.Equal(t, []string{"imgbatch_worker_once"}, processor.processed) + require.Len(t, queue.requeued, 1) + require.Equal(t, defaultBatchImageWorkerRequeueDelay, queue.requeued[0].delay) + require.Equal(t, 1, queue.releaseCount) +} + +func TestBatchImageWorker_RequeuesNonTerminalResultWithRequestedDelay(t *testing.T) { + queue := newFakeBatchImageQueue("imgbatch_worker_requeue") + processor := &fakeBatchImageProcessor{result: BatchImageProcessResult{RequeueAfter: 42 * time.Second}} + worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{}) + + require.NoError(t, worker.RunOnce(context.Background())) + require.Len(t, queue.requeued, 1) + require.Equal(t, "imgbatch_worker_requeue", queue.requeued[0].batchID) + require.Equal(t, 42*time.Second, queue.requeued[0].delay) + require.Empty(t, queue.acked) +} + +func TestBatchImageWorker_AcksTerminalResult(t *testing.T) { + queue := newFakeBatchImageQueue("imgbatch_worker_terminal") + processor := &fakeBatchImageProcessor{result: BatchImageProcessResult{Terminal: true}} + worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{}) + + require.NoError(t, worker.RunOnce(context.Background())) + require.Equal(t, []string{"imgbatch_worker_terminal"}, queue.acked) + require.Empty(t, queue.requeued) +} + +func TestBatchImageWorker_RequeuesOnProcessorError(t *testing.T) { + queue := newFakeBatchImageQueue("imgbatch_worker_error") + processor := &fakeBatchImageProcessor{err: errors.New("processor failed")} + worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{ErrorRetryDelay: 7 * time.Second}) + + require.NoError(t, worker.RunOnce(context.Background())) + require.Len(t, queue.requeued, 1) + require.Equal(t, 7*time.Second, queue.requeued[0].delay) + require.Empty(t, queue.acked) +} + +func TestBatchImageWorker_RequeuesWhenJobLockNotAcquired(t *testing.T) { + queue := newFakeBatchImageQueue("imgbatch_worker_locked") + queue.lockAcquired = false + processor := &fakeBatchImageProcessor{} + worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{LockConflictDelay: 3 * time.Second}) + + // 锁冲突必须按冲突延迟重新入队;直接丢弃会让 job 滞留 active zset, + // 要等 StaleActiveAfter(默认 10 分钟)才被恢复。 + require.NoError(t, worker.RunOnce(context.Background())) + require.Empty(t, processor.processed) + require.Len(t, queue.requeued, 1) + require.Equal(t, 3*time.Second, queue.requeued[0].delay) + require.Empty(t, queue.acked) +} + +func TestNewBatchImageWorkerOptionsFromConfig_UsesFiniteReserveTimeout(t *testing.T) { + opts := NewBatchImageWorkerOptionsFromConfig(nil) + require.Equal(t, defaultBatchImageWorkerReserveBlockTimeout, opts.ReserveBlockTimeout) + require.Positive(t, opts.ReserveBlockTimeout) +} + +type fakeBatchImageQueue struct { + reserved ReservedBatchImageJob + lockAcquired bool + acked []string + requeued []fakeBatchImageRequeue + releaseCount int +} + +type fakeBatchImageRequeue struct { + batchID string + delay time.Duration +} + +func newFakeBatchImageQueue(batchID string) *fakeBatchImageQueue { + return &fakeBatchImageQueue{ + reserved: ReservedBatchImageJob{BatchID: batchID}, + lockAcquired: true, + } +} + +func (q *fakeBatchImageQueue) Enqueue(context.Context, string) error { + return nil +} + +func (q *fakeBatchImageQueue) Reserve(context.Context, time.Duration) (ReservedBatchImageJob, error) { + return q.reserved, nil +} + +func (q *fakeBatchImageQueue) RequeueAfter(_ context.Context, batchID string, delay time.Duration) error { + q.requeued = append(q.requeued, fakeBatchImageRequeue{batchID: batchID, delay: delay}) + return nil +} + +func (q *fakeBatchImageQueue) Ack(_ context.Context, batchID string) error { + q.acked = append(q.acked, batchID) + return nil +} + +func (q *fakeBatchImageQueue) Heartbeat(context.Context, string) error { + return nil +} + +func (q *fakeBatchImageQueue) MoveDueDelayedToReady(context.Context, int) (int, error) { + return 0, nil +} + +func (q *fakeBatchImageQueue) RecoverStaleActive(context.Context, time.Duration, int) (int, error) { + return 0, nil +} + +func (q *fakeBatchImageQueue) TryAcquireJobLock(context.Context, string, time.Duration) (BatchImageJobLock, bool, error) { + if !q.lockAcquired { + return nil, false, nil + } + return fakeBatchImageLock{release: func() { q.releaseCount++ }}, true, nil +} + +type fakeBatchImageLock struct { + release func() +} + +func (l fakeBatchImageLock) Release(context.Context) error { + if l.release != nil { + l.release() + } + return nil +} + +type fakeBatchImageProcessor struct { + result BatchImageProcessResult + err error + processed []string +} + +func (p *fakeBatchImageProcessor) Process(_ context.Context, batchID string) (BatchImageProcessResult, error) { + p.processed = append(p.processed, batchID) + return p.result, p.err +} diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index dc54a1b1f3..8dceebc250 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -512,6 +512,14 @@ func (s *BillingService) initFallbackPricing() { SupportsCacheBreakdown: false, } + // xAI Grok 4.5 (official docs: $2 input / $0.50 cached input / $6 output per MTok) + s.fallbackPrices["grok-4.5"] = &ModelPricing{ + InputPricePerToken: 2e-6, + OutputPricePerToken: 6e-6, + CacheReadPricePerToken: 0.5e-6, + SupportsCacheBreakdown: false, + } + // xAI Grok 4.3 (official docs: $1.25 input / $2.50 output per MTok) s.fallbackPrices["grok-4.3"] = &ModelPricing{ InputPricePerToken: 1.25e-6, @@ -696,7 +704,9 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing { } switch modelLower { - case "grok", "grok-latest", "grok-4.3": + case "grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest": + return s.fallbackPrices["grok-4.5"] + case "grok-4.3": return s.fallbackPrices["grok-4.3"] case "grok-build", "grok-build-0.1": return s.fallbackPrices["grok-build-0.1"] @@ -713,6 +723,14 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) { // 1. 优先从动态价格服务获取 if s.pricingService != nil { litellmPricing := s.pricingService.GetModelPricing(model) + // 仅有图片价、无 token 价的条目(如 LiteLLM 的 imagen 类模型)不能用于 + // token 计费:直接返回会把 token 流量按 $0 计费。跳过后走 fallback, + // 无 fallback 则 fail-closed(ErrModelPricingUnavailable)。 + // 图片计费路径(getDefaultImagePrice / getImageUnitPrice)直接读 + // PricingService,不受影响。 + if litellmPricing != nil && litellmPricing.TokenPricingAbsent { + litellmPricing = nil + } if litellmPricing != nil { // 启用 5m/1h 分类计费的条件: // 1. 存在 1h 价格 @@ -763,6 +781,9 @@ func (s *BillingService) GetModelPricingWithChannel(model string, channelPricing if channelPricing == nil { return pricing, nil } + // 防止修改 fallbackPrices 中的共享指针 + cloned := *pricing + pricing = &cloned if channelPricing.InputPrice != nil { pricing.InputPricePerToken = *channelPricing.InputPrice pricing.InputPricePerTokenPriority = *channelPricing.InputPrice @@ -1220,6 +1241,29 @@ type ImagePriceConfig struct { Price4K *float64 // 4K 尺寸价格(nil 表示使用默认值) } +// VideoPriceConfig 视频生成计费配置。所有价格均为**每秒**单价(USD/s),与 xAI 官方计费口径一致。 +type VideoPriceConfig struct { + Price480P *float64 // 480p 每秒价格(nil 表示使用默认值) + Price720P *float64 // 720p 每秒价格(nil 表示使用默认值) + Price1080P *float64 // 1080p 每秒价格(nil 表示使用默认值) +} + +const ( + defaultImageGenerationPrice = 0.134 + + defaultGrokImagineImagePrice1K = 0.02 + defaultGrokImagineImagePrice2K = 0.02 + defaultGrokImagineImageQualityPrice1K = 0.05 + defaultGrokImagineImageQualityPrice2K = 0.07 + + // 视频默认价为 xAI 官方**每秒**输出价格(USD/s),总价 = 每秒价 × 时长(秒)。 + defaultGrokImagineVideoPrice480P = 0.05 + defaultGrokImagineVideoPrice720P = 0.07 + defaultGrokImagineVideo15Price480P = 0.08 + defaultGrokImagineVideo15Price720P = 0.14 + defaultGrokImagineVideo15Price1080P = 0.25 +) + // CalculateImageCost 计算图片生成费用 // model: 请求的模型名称(用于获取 LiteLLM 默认价格) // imageSize: 图片尺寸 "1K", "2K", "4K" @@ -1251,6 +1295,35 @@ func (s *BillingService) CalculateImageCost(model string, imageSize string, imag } } +// CalculateVideoCost 计算视频生成费用(按秒计费,与 xAI 口径一致)。 +// model: 请求的模型名称(用于获取默认价格) +// resolution: 视频分辨率 "480p", "720p", "1080p" +// videoCount: 生成的视频数量 +// durationSeconds: 单个视频时长(秒),<=0 时按上游默认时长计 +// groupConfig: 分组配置的每秒价格(可能为 nil,表示使用默认值) +// rateMultiplier: 费率倍数 +func (s *BillingService) CalculateVideoCost(model string, resolution string, videoCount int, durationSeconds int, groupConfig *VideoPriceConfig, rateMultiplier float64) *CostBreakdown { + if videoCount <= 0 { + return &CostBreakdown{} + } + resolution = NormalizeVideoBillingResolutionOrDefault(resolution) + durationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(durationSeconds) + + perSecondPrice := s.getVideoUnitPrice(model, resolution, groupConfig) + totalCost := perSecondPrice * float64(durationSeconds) * float64(videoCount) + + if rateMultiplier < 0 { + rateMultiplier = 0 + } + actualCost := totalCost * rateMultiplier + + return &CostBreakdown{ + TotalCost: totalCost, + ActualCost: actualCost, + BillingMode: string(BillingModeVideo), + } +} + // getImageUnitPrice 获取图片单价 func (s *BillingService) getImageUnitPrice(model string, imageSize string, groupConfig *ImagePriceConfig) float64 { // 优先使用分组配置的价格 @@ -1275,8 +1348,33 @@ func (s *BillingService) getImageUnitPrice(model string, imageSize string, group return s.getDefaultImagePrice(model, imageSize) } +func (s *BillingService) getVideoUnitPrice(model string, resolution string, groupConfig *VideoPriceConfig) float64 { + if groupConfig != nil { + switch resolution { + case VideoBillingResolution480P: + if groupConfig.Price480P != nil { + return *groupConfig.Price480P + } + case VideoBillingResolution720P: + if groupConfig.Price720P != nil { + return *groupConfig.Price720P + } + case VideoBillingResolution1080P: + if groupConfig.Price1080P != nil { + return *groupConfig.Price1080P + } + } + } + + return s.getDefaultVideoPrice(model, resolution) +} + // getDefaultImagePrice 获取 LiteLLM 默认图片价格 func (s *BillingService) getDefaultImagePrice(model string, imageSize string) float64 { + if price, ok := getDefaultGrokImagineImagePrice(model, imageSize); ok { + return price + } + basePrice := 0.0 // 从 PricingService 获取 output_cost_per_image @@ -1289,7 +1387,7 @@ func (s *BillingService) getDefaultImagePrice(model string, imageSize string) fl // 如果没有找到价格,使用硬编码默认值($0.134,来自 gemini-3-pro-image-preview) if basePrice <= 0 { - basePrice = 0.134 + basePrice = defaultImageGenerationPrice } // 2K 尺寸 1.5 倍,4K 尺寸翻倍 @@ -1302,3 +1400,74 @@ func (s *BillingService) getDefaultImagePrice(model string, imageSize string) fl return basePrice } + +func (s *BillingService) getDefaultVideoPrice(model string, resolution string) float64 { + if price, ok := getDefaultGrokImagineVideoPrice(model, resolution); ok { + return price + } + + // The bundled LiteLLM schema does not expose an output video generation price. + // Keep the historical model default as the fallback (interpreted as a per-second + // rate; today only Grok models reach video billing, so this path is a safety net), + // while letting group-level video prices override it independently from image prices. + return s.getDefaultImagePrice(model, ImageBillingSize2K) +} + +func getDefaultGrokImagineImagePrice(model string, imageSize string) (float64, bool) { + model = strings.ToLower(strings.TrimSpace(model)) + switch model { + case "grok-imagine-image-quality": + return getGrokImagineImageTierPrice( + imageSize, + defaultGrokImagineImageQualityPrice1K, + defaultGrokImagineImageQualityPrice2K, + ), true + case "grok-imagine", "grok-imagine-image", "grok-imagine-edit": + return getGrokImagineImageTierPrice( + imageSize, + defaultGrokImagineImagePrice1K, + defaultGrokImagineImagePrice2K, + ), true + default: + return 0, false + } +} + +func getGrokImagineImageTierPrice(imageSize string, price1K float64, price2K float64) float64 { + switch NormalizeImageBillingTierOrDefault(imageSize) { + case ImageBillingSize1K: + return price1K + case ImageBillingSize2K, ImageBillingSize4K: + return price2K + default: + return price2K + } +} + +func getDefaultGrokImagineVideoPrice(model string, resolution string) (float64, bool) { + model = strings.ToLower(strings.TrimSpace(model)) + switch { + case strings.HasPrefix(model, "grok-imagine-video-1.5"): + switch NormalizeVideoBillingResolutionOrDefault(resolution) { + case VideoBillingResolution480P: + return defaultGrokImagineVideo15Price480P, true + case VideoBillingResolution720P: + return defaultGrokImagineVideo15Price720P, true + case VideoBillingResolution1080P: + return defaultGrokImagineVideo15Price1080P, true + default: + return defaultGrokImagineVideo15Price480P, true + } + case strings.HasPrefix(model, "grok-imagine-video"): + switch NormalizeVideoBillingResolutionOrDefault(resolution) { + case VideoBillingResolution480P: + return defaultGrokImagineVideoPrice480P, true + case VideoBillingResolution720P, VideoBillingResolution1080P: + return defaultGrokImagineVideoPrice720P, true + default: + return defaultGrokImagineVideoPrice480P, true + } + default: + return 0, false + } +} diff --git a/backend/internal/service/billing_service_test.go b/backend/internal/service/billing_service_test.go index 92c143c6ff..c1f3f6e557 100644 --- a/backend/internal/service/billing_service_test.go +++ b/backend/internal/service/billing_service_test.go @@ -872,6 +872,66 @@ func TestCalculateImageCost(t *testing.T) { require.InDelta(t, 0.134*3, cost.ActualCost, 1e-10) } +func TestCalculateVideoCostUsesSeparateConfig(t *testing.T) { + svc := newTestBillingService() + + imagePrice := 0.4 + videoPrice := 0.08 + imageCost := svc.CalculateImageCost("grok-imagine-video", "2K", 1, &ImagePriceConfig{Price2K: &imagePrice}, 1.0) + videoCost := svc.CalculateVideoCost("grok-imagine-video", "480p", 1, 10, &VideoPriceConfig{Price480P: &videoPrice}, 0.5) + + require.InDelta(t, 0.4, imageCost.TotalCost, 1e-10) + require.InDelta(t, 0.8, videoCost.TotalCost, 1e-10) + require.InDelta(t, 0.4, videoCost.ActualCost, 1e-10) + require.Equal(t, string(BillingModeVideo), videoCost.BillingMode) +} + +func TestCalculateVideoCostBillsPerSecond(t *testing.T) { + svc := newTestBillingService() + + oneSecond := svc.CalculateVideoCost("grok-imagine-video", "720p", 1, 1, nil, 1.0) + fifteenSeconds := svc.CalculateVideoCost("grok-imagine-video", "720p", 1, 15, nil, 1.0) + // duration <=0 时按上游默认 8 秒计费,超出上限按 15 秒收敛。 + defaultDuration := svc.CalculateVideoCost("grok-imagine-video", "720p", 1, 0, nil, 1.0) + clampedDuration := svc.CalculateVideoCost("grok-imagine-video", "720p", 1, 999, nil, 1.0) + + require.InDelta(t, 0.07, oneSecond.TotalCost, 1e-10) + require.InDelta(t, 0.07*15, fifteenSeconds.TotalCost, 1e-10) + require.InDelta(t, 0.07*8, defaultDuration.TotalCost, 1e-10) + require.InDelta(t, 0.07*15, clampedDuration.TotalCost, 1e-10) +} + +func TestCalculateGrokImagineImageCostUsesDefaultRateCard(t *testing.T) { + svc := newTestBillingService() + + standard1K := svc.CalculateImageCost("grok-imagine-image", "1K", 1, nil, 1.0) + standard2K := svc.CalculateImageCost("grok-imagine-image", "2K", 1, nil, 1.0) + quality1K := svc.CalculateImageCost("grok-imagine-image-quality", "1K", 1, nil, 1.0) + quality2K := svc.CalculateImageCost("grok-imagine-image-quality", "2K", 1, nil, 1.0) + + require.InDelta(t, 0.02, standard1K.TotalCost, 1e-10) + require.InDelta(t, 0.02, standard2K.TotalCost, 1e-10) + require.InDelta(t, 0.05, quality1K.TotalCost, 1e-10) + require.InDelta(t, 0.07, quality2K.TotalCost, 1e-10) +} + +func TestCalculateGrokImagineVideoCostUsesDefaultRateCard(t *testing.T) { + svc := newTestBillingService() + + // 默认价目为 xAI 官方每秒价格,按 1 秒时长验证每秒单价。 + standard480P := svc.CalculateVideoCost("grok-imagine-video", "480p", 1, 1, nil, 1.0) + standard720P := svc.CalculateVideoCost("grok-imagine-video", "720p", 1, 1, nil, 1.0) + video15_480P := svc.CalculateVideoCost("grok-imagine-video-1.5", "480p", 1, 1, nil, 1.0) + video15_720P := svc.CalculateVideoCost("grok-imagine-video-1.5", "720p", 1, 1, nil, 1.0) + video15_1080P := svc.CalculateVideoCost("grok-imagine-video-1.5", "1080p", 1, 1, nil, 1.0) + + require.InDelta(t, 0.05, standard480P.TotalCost, 1e-10) + require.InDelta(t, 0.07, standard720P.TotalCost, 1e-10) + require.InDelta(t, 0.08, video15_480P.TotalCost, 1e-10) + require.InDelta(t, 0.14, video15_720P.TotalCost, 1e-10) + require.InDelta(t, 0.25, video15_1080P.TotalCost, 1e-10) +} + func TestIsModelSupported(t *testing.T) { svc := newTestBillingService() @@ -963,6 +1023,22 @@ func TestCalculateCostWithLongContext_PropagatesError(t *testing.T) { require.Contains(t, err.Error(), "pricing not found") } +func TestGetModelPricing_Grok45OfficialFallback(t *testing.T) { + svc := newTestBillingService() + + for _, model := range []string{"grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest"} { + model := model + t.Run(model, func(t *testing.T) { + pricing, err := svc.GetModelPricing(model) + require.NoError(t, err) + require.InDelta(t, 2e-6, pricing.InputPricePerToken, 1e-12) + require.InDelta(t, 6e-6, pricing.OutputPricePerToken, 1e-12) + require.InDelta(t, 0.5e-6, pricing.CacheReadPricePerToken, 1e-12) + require.False(t, pricing.SupportsCacheBreakdown) + }) + } +} + func TestCalculateCost_SupportsCacheBreakdown(t *testing.T) { svc := &BillingService{ cfg: &config.Config{}, diff --git a/backend/internal/service/channel.go b/backend/internal/service/channel.go index 88ed2df79c..1fd9e57068 100644 --- a/backend/internal/service/channel.go +++ b/backend/internal/service/channel.go @@ -14,6 +14,7 @@ const ( BillingModeToken BillingMode = "token" // 按 token 区间计费 BillingModePerRequest BillingMode = "per_request" // 按次计费(支持上下文窗口分层) BillingModeImage BillingMode = "image" // 图片计费(当前按次,预留 token 计费) + BillingModeVideo BillingMode = "video" // 视频生成计费(按视频生成次数) ) // IsValid 检查 BillingMode 是否为合法值 @@ -25,6 +26,15 @@ func (m BillingMode) IsValid() bool { return false } +// IsValidUsageFilter 检查 BillingMode 是否可用于使用记录筛选。 +func (m BillingMode) IsValidUsageFilter() bool { + switch m { + case BillingModeToken, BillingModePerRequest, BillingModeImage, BillingModeVideo, "": + return true + } + return false +} + const ( BillingModelSourceRequested = "requested" BillingModelSourceUpstream = "upstream" diff --git a/backend/internal/service/email_html_escape_test.go b/backend/internal/service/email_html_escape_test.go new file mode 100644 index 0000000000..d1e7aca3bc --- /dev/null +++ b/backend/internal/service/email_html_escape_test.go @@ -0,0 +1,67 @@ +//go:build unit + +package service + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBuildVerifyCodeEmailBody_EscapesSiteName(t *testing.T) { + svc := &EmailService{} + + t.Run("escapes_script_injection", func(t *testing.T) { + body := svc.buildVerifyCodeEmailBody("123456", `

`) + + assert.NotContains(t, body, ""}`) + + result := injectSiteTitle(html, settingsJSON) + + assert.NotContains(t, string(result), "<script>") + assert.Contains(t, string(result), "</title><script>alert(1)</script><title>") + }) + + t.Run("escapes_ampersand_in_site_name", func(t *testing.T) { + html := []byte(`<html><head><title>Sub2API`) + settingsJSON := []byte(`{"site_name":"A&B"}`) + + result := injectSiteTitle(html, settingsJSON) + + assert.Contains(t, string(result), "A&B - AI API Gateway") + }) + t.Run("preserves_rest_of_html", func(t *testing.T) { html := []byte(`Sub2API
`) settingsJSON := []byte(`{"site_name":"TestSite"}`) diff --git a/backend/migrations/159_batch_image_foundation.sql b/backend/migrations/159_batch_image_foundation.sql new file mode 100644 index 0000000000..d2464cc683 --- /dev/null +++ b/backend/migrations/159_batch_image_foundation.sql @@ -0,0 +1,86 @@ +CREATE TABLE IF NOT EXISTS batch_image_jobs ( + id BIGSERIAL PRIMARY KEY, + batch_id VARCHAR(64) NOT NULL UNIQUE, + user_id BIGINT NOT NULL, + api_key_id BIGINT, + account_id BIGINT, + provider VARCHAR(32) NOT NULL, + model VARCHAR(128) NOT NULL, + status VARCHAR(32) NOT NULL DEFAULT 'created', + provider_job_name VARCHAR(512), + gcs_input_uri VARCHAR(1024), + gcs_output_uri VARCHAR(1024), + item_count INTEGER NOT NULL, + success_count INTEGER NOT NULL DEFAULT 0, + fail_count INTEGER NOT NULL DEFAULT 0, + cancelled_count INTEGER NOT NULL DEFAULT 0, + estimated_cost DECIMAL(20,10) NOT NULL DEFAULT 0, + hold_amount DECIMAL(20,10), + actual_cost DECIMAL(20,10), + currency VARCHAR(16) NOT NULL DEFAULT 'USD', + hold_id VARCHAR(128), + idempotency_key VARCHAR(255), + request_hash VARCHAR(128), + manifest_hash VARCHAR(128), + retry_count INTEGER NOT NULL DEFAULT 0, + version INTEGER NOT NULL DEFAULT 0, + output_expires_at TIMESTAMPTZ, + input_deleted_at TIMESTAMPTZ, + output_deleted_at TIMESTAMPTZ, + last_error_code VARCHAR(128), + last_error_message TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + submitted_at TIMESTAMPTZ, + started_at TIMESTAMPTZ, + finished_at TIMESTAMPTZ, + settled_at TIMESTAMPTZ +); + +CREATE INDEX IF NOT EXISTS batch_image_jobs_user_created_at_idx ON batch_image_jobs (user_id, created_at); +CREATE INDEX IF NOT EXISTS batch_image_jobs_status_idx ON batch_image_jobs (status); +CREATE INDEX IF NOT EXISTS batch_image_jobs_provider_status_idx ON batch_image_jobs (provider, status); +CREATE INDEX IF NOT EXISTS batch_image_jobs_idempotency_key_idx ON batch_image_jobs (idempotency_key) + WHERE idempotency_key IS NOT NULL AND idempotency_key <> ''; +CREATE UNIQUE INDEX IF NOT EXISTS batch_image_jobs_manifest_hash_uq ON batch_image_jobs (manifest_hash) + WHERE manifest_hash IS NOT NULL AND manifest_hash <> ''; +CREATE INDEX IF NOT EXISTS batch_image_jobs_output_expires_at_idx ON batch_image_jobs (output_expires_at); + +CREATE TABLE IF NOT EXISTS batch_image_items ( + id BIGSERIAL PRIMARY KEY, + job_id VARCHAR(64) NOT NULL REFERENCES batch_image_jobs(batch_id) ON DELETE CASCADE, + custom_id VARCHAR(255) NOT NULL, + status VARCHAR(32) NOT NULL, + request_hash VARCHAR(128), + prompt_preview TEXT, + provider_source_object VARCHAR(1024), + source_line_number INTEGER, + source_byte_offset BIGINT, + source_byte_length BIGINT, + mime_type VARCHAR(128), + file_extension VARCHAR(32), + image_count INTEGER NOT NULL DEFAULT 0, + error_code VARCHAR(128), + error_message TEXT, + billed_amount DECIMAL(20,10), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + indexed_at TIMESTAMPTZ +); + +CREATE UNIQUE INDEX IF NOT EXISTS batch_image_items_job_custom_uq ON batch_image_items (job_id, custom_id); +CREATE INDEX IF NOT EXISTS batch_image_items_job_status_idx ON batch_image_items (job_id, status); +CREATE INDEX IF NOT EXISTS batch_image_items_provider_source_object_idx ON batch_image_items (provider_source_object); + +CREATE TABLE IF NOT EXISTS batch_image_events ( + id BIGSERIAL PRIMARY KEY, + job_id VARCHAR(64) NOT NULL REFERENCES batch_image_jobs(batch_id) ON DELETE CASCADE, + event_type VARCHAR(64) NOT NULL, + payload JSONB, + event_hash VARCHAR(128), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX IF NOT EXISTS batch_image_events_job_created_at_idx ON batch_image_events (job_id, created_at); +CREATE INDEX IF NOT EXISTS batch_image_events_event_type_idx ON batch_image_events (event_type); +CREATE UNIQUE INDEX IF NOT EXISTS batch_image_events_job_event_hash_uq ON batch_image_events (job_id, event_hash) + WHERE event_hash IS NOT NULL AND event_hash <> ''; diff --git a/backend/migrations/160_add_user_frozen_balance.sql b/backend/migrations/160_add_user_frozen_balance.sql new file mode 100644 index 0000000000..d113efc9f7 --- /dev/null +++ b/backend/migrations/160_add_user_frozen_balance.sql @@ -0,0 +1,2 @@ +ALTER TABLE users + ADD COLUMN IF NOT EXISTS frozen_balance DECIMAL(20,8) NOT NULL DEFAULT 0; diff --git a/backend/migrations/160_batch_image_provider_refs.sql b/backend/migrations/160_batch_image_provider_refs.sql new file mode 100644 index 0000000000..1ec9862e9c --- /dev/null +++ b/backend/migrations/160_batch_image_provider_refs.sql @@ -0,0 +1,3 @@ +ALTER TABLE batch_image_jobs + ADD COLUMN IF NOT EXISTS provider_input_ref VARCHAR(1024), + ADD COLUMN IF NOT EXISTS provider_output_ref VARCHAR(1024); diff --git a/backend/migrations/161_batch_image_pricing_snapshot.sql b/backend/migrations/161_batch_image_pricing_snapshot.sql new file mode 100644 index 0000000000..3ae6d1fbb6 --- /dev/null +++ b/backend/migrations/161_batch_image_pricing_snapshot.sql @@ -0,0 +1,25 @@ +ALTER TABLE groups + ADD COLUMN IF NOT EXISTS batch_image_discount_multiplier DECIMAL(10,4) NOT NULL DEFAULT 0.5, + ADD COLUMN IF NOT EXISTS batch_image_hold_multiplier DECIMAL(10,4) NOT NULL DEFAULT 0.6; + +COMMENT ON COLUMN groups.batch_image_discount_multiplier IS '批量图片生成折扣倍率,最终单价会乘以该值;0 表示免费'; +COMMENT ON COLUMN groups.batch_image_hold_multiplier IS '批量图片生成冻结价格比例,按普通生图原价乘以该比例冻结,结算后释放差额'; + +ALTER TABLE batch_image_jobs + ADD COLUMN IF NOT EXISTS base_unit_price DECIMAL(20,10) NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS group_rate_multiplier DECIMAL(10,4) NOT NULL DEFAULT 1.0, + ADD COLUMN IF NOT EXISTS account_rate_multiplier DECIMAL(10,4) NOT NULL DEFAULT 1.0, + ADD COLUMN IF NOT EXISTS batch_discount_multiplier DECIMAL(10,4) NOT NULL DEFAULT 0.5, + ADD COLUMN IF NOT EXISTS hold_multiplier DECIMAL(10,4) NOT NULL DEFAULT 0.6, + ADD COLUMN IF NOT EXISTS billable_unit_price DECIMAL(20,10) NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS hold_unit_price DECIMAL(20,10) NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS pricing_snapshot_version INTEGER NOT NULL DEFAULT 0; + +COMMENT ON COLUMN batch_image_jobs.base_unit_price IS '提交时快照的基础批量图片单价'; +COMMENT ON COLUMN batch_image_jobs.group_rate_multiplier IS '提交时快照的分组/用户专属图片倍率'; +COMMENT ON COLUMN batch_image_jobs.account_rate_multiplier IS '提交时快照的账号倍率'; +COMMENT ON COLUMN batch_image_jobs.batch_discount_multiplier IS '提交时快照的批量折扣倍率'; +COMMENT ON COLUMN batch_image_jobs.hold_multiplier IS '提交时快照的冻结价格比例,按普通生图原价乘以该比例冻结'; +COMMENT ON COLUMN batch_image_jobs.billable_unit_price IS '提交时快照的实际结算单价'; +COMMENT ON COLUMN batch_image_jobs.hold_unit_price IS '提交时快照的冻结单价'; +COMMENT ON COLUMN batch_image_jobs.pricing_snapshot_version IS '批量图片任务价格快照版本;0 表示旧任务无快照'; diff --git a/backend/migrations/162_add_group_batch_image_generation_gate.sql b/backend/migrations/162_add_group_batch_image_generation_gate.sql new file mode 100644 index 0000000000..e96541b931 --- /dev/null +++ b/backend/migrations/162_add_group_batch_image_generation_gate.sql @@ -0,0 +1,4 @@ +ALTER TABLE groups + ADD COLUMN IF NOT EXISTS allow_batch_image_generation BOOLEAN NOT NULL DEFAULT false; + +COMMENT ON COLUMN groups.allow_batch_image_generation IS '是否允许该分组使用批量图片生成能力'; diff --git a/backend/migrations/163_batch_image_default_discount_and_hold_ratio.sql b/backend/migrations/163_batch_image_default_discount_and_hold_ratio.sql new file mode 100644 index 0000000000..65ac699ba9 --- /dev/null +++ b/backend/migrations/163_batch_image_default_discount_and_hold_ratio.sql @@ -0,0 +1,19 @@ +ALTER TABLE groups + ALTER COLUMN batch_image_discount_multiplier SET DEFAULT 0.5, + ALTER COLUMN batch_image_hold_multiplier SET DEFAULT 0.6; + +UPDATE groups +SET batch_image_discount_multiplier = 0.5 +WHERE batch_image_discount_multiplier = 1.0; + +UPDATE groups +SET batch_image_hold_multiplier = 0.6 +WHERE batch_image_hold_multiplier = 1.05; + +COMMENT ON COLUMN groups.batch_image_hold_multiplier IS '批量图片生成冻结价格比例,按普通生图原价乘以该比例冻结,结算后释放差额'; + +ALTER TABLE batch_image_jobs + ALTER COLUMN batch_discount_multiplier SET DEFAULT 0.5, + ALTER COLUMN hold_multiplier SET DEFAULT 0.6; + +COMMENT ON COLUMN batch_image_jobs.hold_multiplier IS '提交时快照的冻结价格比例,按普通生图原价乘以该比例冻结'; diff --git a/backend/migrations/164_batch_image_download_and_user_delete.sql b/backend/migrations/164_batch_image_download_and_user_delete.sql new file mode 100644 index 0000000000..56848b7c08 --- /dev/null +++ b/backend/migrations/164_batch_image_download_and_user_delete.sql @@ -0,0 +1,9 @@ +ALTER TABLE batch_image_jobs + ADD COLUMN IF NOT EXISTS downloaded_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS user_deleted_at TIMESTAMPTZ; + +CREATE INDEX IF NOT EXISTS batch_image_jobs_downloaded_at_idx ON batch_image_jobs (downloaded_at); +CREATE INDEX IF NOT EXISTS batch_image_jobs_user_deleted_at_idx ON batch_image_jobs (user_deleted_at); + +COMMENT ON COLUMN batch_image_jobs.downloaded_at IS '用户首次成功下载批量图片 ZIP 的时间'; +COMMENT ON COLUMN batch_image_jobs.user_deleted_at IS '用户侧删除/隐藏任务记录的时间;账务记录仍保留'; diff --git a/backend/migrations/165_hide_pre_upstream_batch_image_failures.sql b/backend/migrations/165_hide_pre_upstream_batch_image_failures.sql new file mode 100644 index 0000000000..3cd9293e74 --- /dev/null +++ b/backend/migrations/165_hide_pre_upstream_batch_image_failures.sql @@ -0,0 +1,16 @@ +UPDATE batch_image_jobs +SET user_deleted_at = COALESCE(user_deleted_at, updated_at, created_at, NOW()), + updated_at = NOW() +WHERE user_deleted_at IS NULL + AND provider_job_name IS NULL + AND status = 'failed' + AND last_error_code IN ( + 'INSUFFICIENT_BALANCE', + 'PROVIDER_SUBMIT_FAILED', + 'BATCH_IMAGE_PROVIDER_SUBMIT_FAILED', + 'BATCH_IMAGE_VERTEX_GCS_BUCKET_MISSING', + 'VERTEX_MANAGED_GCS_BUCKET_MISSING', + 'BATCH_IMAGE_PROVIDER_MISSING_API_KEY', + 'BATCH_IMAGE_PROVIDER_MISSING_SERVICE_ACCOUNT', + 'BATCH_IMAGE_PROVIDER_UNSUPPORTED_ACCOUNT' + ); diff --git a/backend/migrations/166_batch_image_task_name.sql b/backend/migrations/166_batch_image_task_name.sql new file mode 100644 index 0000000000..ef942d8dad --- /dev/null +++ b/backend/migrations/166_batch_image_task_name.sql @@ -0,0 +1,10 @@ +ALTER TABLE batch_image_jobs + ADD COLUMN IF NOT EXISTS task_name VARCHAR(255) NOT NULL DEFAULT ''; + +UPDATE batch_image_jobs +SET task_name = TO_CHAR(created_at AT TIME ZONE 'Asia/Shanghai', 'YYYY-MM-DD HH24:MI:SS') +WHERE task_name = ''; + +CREATE INDEX IF NOT EXISTS batch_image_jobs_task_name_idx ON batch_image_jobs (task_name); + +COMMENT ON COLUMN batch_image_jobs.task_name IS '用户可读的批量生图任务名称'; diff --git a/backend/migrations/167_clear_auto_batch_image_task_names.sql b/backend/migrations/167_clear_auto_batch_image_task_names.sql new file mode 100644 index 0000000000..d12eefb48c --- /dev/null +++ b/backend/migrations/167_clear_auto_batch_image_task_names.sql @@ -0,0 +1,5 @@ +UPDATE batch_image_jobs +SET task_name = '' +WHERE task_name = TO_CHAR(created_at AT TIME ZONE 'Asia/Shanghai', 'YYYY-MM-DD HH24:MI:SS'); + +COMMENT ON COLUMN batch_image_jobs.task_name IS '用户填写的批量生图任务名称;为空时用户侧显示未填写'; diff --git a/backend/migrations/168_restore_empty_batch_image_task_names.sql b/backend/migrations/168_restore_empty_batch_image_task_names.sql new file mode 100644 index 0000000000..7b2e34bb61 --- /dev/null +++ b/backend/migrations/168_restore_empty_batch_image_task_names.sql @@ -0,0 +1,5 @@ +UPDATE batch_image_jobs +SET task_name = TO_CHAR(created_at AT TIME ZONE 'Asia/Shanghai', 'YYYY-MM-DD HH24:MI:SS') +WHERE task_name = ''; + +COMMENT ON COLUMN batch_image_jobs.task_name IS '用户填写的批量生图任务名称;提交时为空则默认写入当前时间'; diff --git a/backend/migrations/169_batch_image_parent_batch.sql b/backend/migrations/169_batch_image_parent_batch.sql new file mode 100644 index 0000000000..e089c5e49a --- /dev/null +++ b/backend/migrations/169_batch_image_parent_batch.sql @@ -0,0 +1,8 @@ +ALTER TABLE batch_image_jobs + ADD COLUMN IF NOT EXISTS parent_batch_id VARCHAR(64); + +CREATE INDEX IF NOT EXISTS batch_image_jobs_parent_batch_id_idx + ON batch_image_jobs (parent_batch_id) + WHERE parent_batch_id IS NOT NULL AND parent_batch_id <> ''; + +COMMENT ON COLUMN batch_image_jobs.parent_batch_id IS '父批量生图任务 ID;失败项重试等子任务挂在主任务下展示'; diff --git a/backend/migrations/170_add_grok_video_pricing_controls.sql b/backend/migrations/170_add_grok_video_pricing_controls.sql new file mode 100644 index 0000000000..9c798a9945 --- /dev/null +++ b/backend/migrations/170_add_grok_video_pricing_controls.sql @@ -0,0 +1,16 @@ +-- Add independent group pricing controls for Grok video generation. +-- Video prices intentionally do not backfill from image prices: image and video +-- generation must be priced separately. + +ALTER TABLE groups + ADD COLUMN IF NOT EXISTS video_rate_independent BOOLEAN NOT NULL DEFAULT false, + ADD COLUMN IF NOT EXISTS video_rate_multiplier DECIMAL(10,4) NOT NULL DEFAULT 1.0, + ADD COLUMN IF NOT EXISTS video_price_480p DECIMAL(20,8), + ADD COLUMN IF NOT EXISTS video_price_720p DECIMAL(20,8), + ADD COLUMN IF NOT EXISTS video_price_1080p DECIMAL(20,8); + +COMMENT ON COLUMN groups.video_rate_independent IS '视频生成是否使用独立倍率;false 表示共享分组有效倍率'; +COMMENT ON COLUMN groups.video_rate_multiplier IS '视频生成独立倍率,仅 video_rate_independent=true 时生效'; +COMMENT ON COLUMN groups.video_price_480p IS '480p 视频生成单价 (USD),Grok 平台使用'; +COMMENT ON COLUMN groups.video_price_720p IS '720p 视频生成单价 (USD),Grok 平台使用'; +COMMENT ON COLUMN groups.video_price_1080p IS '1080p 视频生成单价 (USD),Grok 平台使用'; diff --git a/backend/migrations/171_allow_video_usage_without_image_size.sql b/backend/migrations/171_allow_video_usage_without_image_size.sql new file mode 100644 index 0000000000..767a15c00e --- /dev/null +++ b/backend/migrations/171_allow_video_usage_without_image_size.sql @@ -0,0 +1,17 @@ +-- Grok video generation stores billing_mode='video' and keeps image_count=1 +-- only as a legacy media-unit counter. It must not be forced to carry an +-- image_size, because video pricing uses video_resolution/request metadata. + +ALTER TABLE usage_logs + DROP CONSTRAINT IF EXISTS usage_logs_image_billing_size_check; + +ALTER TABLE usage_logs + ADD CONSTRAINT usage_logs_image_billing_size_check + CHECK ( + image_count <= 0 + OR billing_mode = 'video' + OR ( + image_size IS NOT NULL + AND image_size IN ('1K', '2K', '4K', 'mixed') + ) + ) NOT VALID; diff --git a/backend/migrations/172_video_per_second_billing_metadata.sql b/backend/migrations/172_video_per_second_billing_metadata.sql new file mode 100644 index 0000000000..e54b83dafe --- /dev/null +++ b/backend/migrations/172_video_per_second_billing_metadata.sql @@ -0,0 +1,38 @@ +-- Grok video billing is per second of generated output (xAI rate card), so usage +-- rows must record the billed resolution and duration for auditability. The +-- image-size check constraint must also exempt any video row by video_count +-- instead of billing_mode='video' alone: a video request billed through a +-- token-mode channel price produces billing_mode='token' with image_count=1 +-- (legacy media counter) and no image_size, which the previous constraint +-- rejected and silently dropped the whole billing transaction. + +ALTER TABLE usage_logs + ADD COLUMN IF NOT EXISTS video_count INTEGER NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS video_resolution VARCHAR(10), + ADD COLUMN IF NOT EXISTS video_duration_seconds INTEGER; + +COMMENT ON COLUMN usage_logs.video_count IS '视频生成数量;>0 表示本行是视频生成用量'; +COMMENT ON COLUMN usage_logs.video_resolution IS '计费用视频分辨率 480p/720p/1080p'; +COMMENT ON COLUMN usage_logs.video_duration_seconds IS '提交时请求的视频时长(秒),按秒计费的乘数'; + +ALTER TABLE usage_logs + DROP CONSTRAINT IF EXISTS usage_logs_image_billing_size_check; + +ALTER TABLE usage_logs + ADD CONSTRAINT usage_logs_image_billing_size_check + CHECK ( + image_count <= 0 + OR billing_mode = 'video' + OR COALESCE(video_count, 0) > 0 + OR ( + image_size IS NOT NULL + AND image_size IN ('1K', '2K', '4K', 'mixed') + ) + ) NOT VALID; + +-- Group video prices are per-second rates (USD/s), matching the xAI rate card; +-- total cost = per-second price x duration seconds. Clarify the column docs +-- introduced by migration 170, which read as per-video prices. +COMMENT ON COLUMN groups.video_price_480p IS '480p 视频生成每秒单价 (USD/s),Grok 平台使用'; +COMMENT ON COLUMN groups.video_price_720p IS '720p 视频生成每秒单价 (USD/s),Grok 平台使用'; +COMMENT ON COLUMN groups.video_price_1080p IS '1080p 视频生成每秒单价 (USD/s),Grok 平台使用'; diff --git a/backend/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index e7d4e7ded3..b9d18f8bf5 100644 --- a/backend/resources/model-pricing/model_prices_and_context_window.json +++ b/backend/resources/model-pricing/model_prices_and_context_window.json @@ -873,7 +873,7 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", - "output_cost_per_image": 0.039, + "output_cost_per_image": 0.034, "output_cost_per_token": 0.0, "source": "https://ai.google.dev/pricing", "supported_modalities": [ @@ -1625,6 +1625,47 @@ "supports_web_search": true, "web_search_billing_unit": "per_query" }, + "gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, "gemini-3-pro-preview": { "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, @@ -1726,6 +1767,39 @@ "supports_web_search": true, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.0003, + "input_cost_per_token": 3e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.034, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, diff --git a/deploy/Dockerfile b/deploy/Dockerfile index c9222da50f..88d364a7f3 100644 --- a/deploy/Dockerfile +++ b/deploy/Dockerfile @@ -7,7 +7,7 @@ # ============================================================================= ARG NODE_IMAGE=node:24-alpine -ARG GOLANG_IMAGE=golang:1.26.4-alpine +ARG GOLANG_IMAGE=golang:1.26.5-alpine ARG ALPINE_IMAGE=alpine:3.20 ARG GOPROXY=https://goproxy.cn,direct ARG GOSUMDB=sum.golang.google.cn @@ -18,9 +18,11 @@ ARG GOSUMDB=sum.golang.google.cn FROM ${NODE_IMAGE} AS frontend-builder WORKDIR /app/frontend +ENV NODE_OPTIONS=--max-old-space-size=1536 -# Install pnpm -RUN corepack enable && corepack prepare pnpm@latest --activate +# Install pnpm. Keep this pinned to the lockfile-compatible major version so +# Docker builds remain reproducible when pnpm changes config validation rules. +RUN corepack enable && corepack prepare pnpm@9.15.9 --activate # Install dependencies first (better caching) COPY frontend/package.json frontend/pnpm-lock.yaml ./ diff --git a/deploy/docker-compose.dev.yml b/deploy/docker-compose.dev.yml index 07e89e0b51..6f5b3f56f3 100644 --- a/deploy/docker-compose.dev.yml +++ b/deploy/docker-compose.dev.yml @@ -13,6 +13,8 @@ services: build: context: .. dockerfile: Dockerfile + args: + NPM_CONFIG_REGISTRY: ${NPM_CONFIG_REGISTRY:-https://registry.npmmirror.com} container_name: sub2api-dev restart: unless-stopped ports: @@ -41,6 +43,12 @@ services: - SETUP_MIGRATION_TIMEOUT_SECONDS=${SETUP_MIGRATION_TIMEOUT_SECONDS:-0} - TOTP_ENCRYPTION_KEY=${TOTP_ENCRYPTION_KEY:-} - TZ=${TZ:-Asia/Shanghai} + # Local mainland-China development proxy. Containers cannot use + # 127.0.0.1 for the host proxy, so default to Docker Desktop's host name. + - HTTP_PROXY=${SUB2API_DEV_HTTP_PROXY:-http://host.docker.internal:7897} + - HTTPS_PROXY=${SUB2API_DEV_HTTPS_PROXY:-http://host.docker.internal:7897} + - ALL_PROXY=${SUB2API_DEV_ALL_PROXY:-socks5://host.docker.internal:7897} + - NO_PROXY=${SUB2API_DEV_NO_PROXY:-127.0.0.1,localhost,::1,postgres,redis,sub2api,192.168.0.0/16,10.0.0.0/8,172.16.0.0/12,.local} # OpenAI HTTP upstream protocol/timeout - GATEWAY_OPENAI_RESPONSE_HEADER_TIMEOUT=${GATEWAY_OPENAI_RESPONSE_HEADER_TIMEOUT:-0} - GATEWAY_OPENAI_HTTP2_ENABLED=${GATEWAY_OPENAI_HTTP2_ENABLED:-true} @@ -55,6 +63,15 @@ services: - GATEWAY_IMAGE_CONCURRENCY_OVERFLOW_MODE=${GATEWAY_IMAGE_CONCURRENCY_OVERFLOW_MODE:-reject} - GATEWAY_IMAGE_CONCURRENCY_WAIT_TIMEOUT_SECONDS=${GATEWAY_IMAGE_CONCURRENCY_WAIT_TIMEOUT_SECONDS:-30} - GATEWAY_IMAGE_CONCURRENCY_MAX_WAITING_REQUESTS=${GATEWAY_IMAGE_CONCURRENCY_MAX_WAITING_REQUESTS:-100} + - BATCH_IMAGE_ENABLED=${BATCH_IMAGE_ENABLED:-true} + - BATCH_IMAGE_QUEUE_ENABLED=${BATCH_IMAGE_QUEUE_ENABLED:-true} + - BATCH_IMAGE_VERTEX_ENABLED=${BATCH_IMAGE_VERTEX_ENABLED:-true} + - BATCH_IMAGE_VERTEX_PROJECT_ID=${BATCH_IMAGE_VERTEX_PROJECT_ID:-project-28424c50-8df2-46e2-a27} + - BATCH_IMAGE_VERTEX_LOCATION=${BATCH_IMAGE_VERTEX_LOCATION:-global} + - BATCH_IMAGE_VERTEX_MANAGED_GCS_BUCKET=${BATCH_IMAGE_VERTEX_MANAGED_GCS_BUCKET:-sub2-batch-image-prod-project-28424c50-8df2-46e2-a27} + - BATCH_IMAGE_VERTEX_MANAGED_GCS_PREFIX=${BATCH_IMAGE_VERTEX_MANAGED_GCS_PREFIX:-batch-image/prod/{batch_id}} + - BATCH_IMAGE_VERTEX_INPUT_RETENTION_HOURS=${BATCH_IMAGE_VERTEX_INPUT_RETENTION_HOURS:-24} + - BATCH_IMAGE_VERTEX_OUTPUT_RETENTION_HOURS=${BATCH_IMAGE_VERTEX_OUTPUT_RETENTION_HOURS:-72} depends_on: postgres: condition: service_healthy diff --git a/docs/BATCH_IMAGE_MVP.md b/docs/BATCH_IMAGE_MVP.md new file mode 100644 index 0000000000..69b0be724a --- /dev/null +++ b/docs/BATCH_IMAGE_MVP.md @@ -0,0 +1,358 @@ +# Batch Image MVP + +Sub2API Batch Image MVP provides asynchronous Gemini image batch generation through a unified API surface backed by Redis workers, PostgreSQL state, and provider-specific batch backends. + +Supported providers: + +- `gemini_api` +- `vertex` + +API users do not see Gemini file names, Vertex job names, GCS paths, signed URLs, API keys, or service account material. Downloads are proxied through Sub2API in this MVP. + +## API Routes + +```text +POST /v1/images/batches +GET /v1/images/batches/{id} +GET /v1/images/batches/{id}/items +GET /v1/images/batches/{id}/items/{custom_id}/content +GET /v1/images/batches/{id}/download +POST /v1/images/batches/{id}/cancel +DELETE /v1/images/batches/{id}/outputs +``` + +Submit request: + +```json +{ + "model": "gemini-2.5-flash-image", + "provider": "gemini_api", + "items": [ + { + "custom_id": "cover_001", + "prompt": "A clean product hero image...", + "output_count": 1, + "reference_images": [ + { + "id": "product-front", + "type": "subject", + "mime_type": "image/png", + "data": "" + }, + { + "id": "style", + "type": "style", + "mime_type": "image/jpeg", + "file_uri": "gs://internal-managed-bucket/batch-image/refs/style.jpg" + } + ] + } + ], + "image_size": "1K", + "response_mime_type": "image/png" +} +``` + +`reference_images` is optional per item. Inline `data` is a base64 string decoded by the backend; `file_uri` is reserved for internal Google Cloud Storage references and must be a `gs://` URI. Each reference image must use one of `image/png`, `image/jpeg`, or `image/webp`. Current model limits are: + +- `gemini-2.5-flash-image` and other Flash Image aliases: up to 3 reference images per item. +- `gemini-3-pro-image` and other Pro Image aliases: up to 14 reference images per item. +- Per batch job: up to 1000 reference image attachments total after `output_count` expansion across all items. This is an internal Sub2API guardrail for request size and cost control, not the generated-image cap and not a Pro Image per-item capability. The generated-output cap is 200 images per job. +- Per batch job: up to 128 MB decoded inline reference image data total. For large batches or repeated reference images, prefer `gs://` `file_uri` references or split the request into multiple jobs. + +`output_count` is optional per item and defaults to `1`. It means "repeat this prompt and reference image set N times" rather than relying on Gemini to return multiple images from one upstream request. The backend expands each repeat into a separate provider JSONL line with suffixed custom ids such as `cover_001_01`, `cover_001_02`. Current limits are: + +- Per prompt item: up to 4 output images. +- Per batch job: up to 200 expected output images after expansion. This is the hard generated-output cap for a single job; clients and Codex skills must split larger workloads before submission. +- The output-image limit intentionally matches the default ZIP item limit so newly submitted jobs are always downloadable as one ZIP by item count. ZIP byte size is still capped separately by `max_download_bytes_per_request`. + +Public batch response: + +```json +{ + "id": "imgbatch_0123456789abcdef0123456789abcdef", + "object": "image.batch", + "status": "queued", + "model": "gemini-2.5-flash-image", + "provider": "gemini_api", + "item_count": 1, + "success_count": 0, + "fail_count": 0, + "estimated_cost": 0.25, + "actual_cost": null, + "created_at": 1783123200, + "submitted_at": 1783123201, + "settled_at": null +} +``` + +Public items response: + +```json +{ + "object": "list", + "data": [ + { + "custom_id": "cover_001", + "status": "succeeded", + "mime_type": "image/png", + "file_extension": "png", + "image_count": 1, + "error": null + } + ], + "has_more": false +} +``` + +## Lifecycle + +Internal lifecycle: + +```text +created -> uploading -> submitted -> running -> indexing -> settling -> completed +``` + +Terminal and cleanup statuses: + +```text +failed +cancelled +completed -> output_deleted +``` + +Public status mapping: + +```text +created/uploading/submitted -> queued +running -> running +indexing -> processing_results +settling -> settling +completed -> completed +failed -> failed +cancelled -> cancelled +output_deleted -> output_deleted +``` + +`completed -> output_deleted` happens after manual output deletion or TTL cleanup. + +## Redis + +Redis is used for wakeups, retries, worker coordination, per-job locks, and download limiting. PostgreSQL remains the source of truth. + +`batch_image.queue_enabled` defaults to `false`. When it is set to `true`, app startup starts `BatchImageWorker` runtime loops for the Redis ready queue, delayed queue mover, and stale active recovery. The worker reserves jobs from the Redis ready queue and blocks there when no job is available. + +Redis structures: + +- Ready queue: `batch_image.queue_ready_key` +- Delayed queue: `batch_image.queue_delayed_key` +- Active set: `batch_image.queue_active_key` +- Inflight keys: `batch_image.inflight_key_prefix` +- Per-job lock keys: `batch_image.lock_key_prefix` +- Queue idempotency keys: `batch_image.idempotency_key_prefix` +- Download limiter keys managed by the download limiter + +Workers should reserve from Redis. They are not expected to run as a database scan loop. + +The worker does not perform DB scan polling. Database reads happen only after a Redis queue reservation yields a specific batch id. + +## Billing + +MVP billing rules: + +- Submit may estimate cost. +- Settlement runs after result indexing. +- Only successful images are charged. +- Failed items are not charged. +- Reference images are sent to Gemini as input and can create small upstream input-token and temporary storage cost. They are counted once per expanded output request when `output_count > 1`, but the public MVP billing model does not add a separate reference-image surcharge. User-facing estimated, held, and settled amounts are still based on the output image count and configured batch image unit price. +- Settlement request id is `batch_image_settlement:{batch_id}`. +- Settlement is idempotent; re-running settlement must not double charge. +- Settlement billing failures are retried with a bounded retry limit. After the retry limit is reached, the job is failed and the remaining hold is released through the idempotent release path. + +Exact production pricing is resolved through model pricing configuration and is not defined here. + +## Cleanup + +Defaults: + +- Input retention after terminal status: 24 hours. +- Output retention after terminal status: 72 hours. +- Maximum output retention: 7 days. +- Cleanup interval: 30 minutes. +- Cleanup batch size: 100. + +Manual output deletion: + +```text +DELETE /v1/images/batches/{id}/outputs +``` + +After output cleanup, downloads return `410 Gone` with `BATCH_IMAGE_OUTPUT_DELETED`. + +Cleanup never accepts user-supplied provider paths. Provider cleanup must use server-generated refs and prefix-safe deletion. + +For the managed Vertex/GCS batch bucket, disable Cloud Storage soft delete or configure lifecycle carefully to avoid hidden retained storage cost. + +## Provider Notes + +`gemini_api`: + +- Uses Gemini Batch API with JSONL file mode. +- Supports Gemini `apikey` upstream accounts with a configured API key. +- Result file refs are internal. +- API keys are never returned. +- The provider can be selected and submitted through Sub2API when an administrator configures a Gemini API-key upstream account. In the 2026-07-07 PR validation, this path was verified as selectable/callable, but successful image generation was not continued because the test API key had no prepayment. + +`vertex`: + +- Uses Vertex `BatchPredictionJob` with managed GCS JSONL. +- Supports Gemini `service_account` upstream accounts with valid service account JSON. +- GCS bucket and prefix are server-managed. +- Vertex job name and GCS paths are internal. +- Batch image output should be treated as `1K`/default only in MVP. +- Do not promise `2K` or `4K`. + +Other Gemini account/login types are not selected by the current batch image providers unless they expose equivalent API-key or service-account credentials through the same provider flow. They were not covered by the 2026-07-07 PR validation. + +## Official Google Enablement + +Operators must enable Gemini/Vertex capability in Google's official console before turning on Sub2API batch image for any group. Sub2API feature flags and group switches do not create Google-side access by themselves. + +Recommended production path: + +- Use a Google Cloud project with billing enabled. +- Enable the relevant Gemini API / Vertex AI APIs for the project. +- Use a service account or Application Default Credentials for the Sub2API runtime. +- Create one fixed Cloud Storage bucket for batch image input and output, then grant the runtime and Vertex service agent the minimum required bucket permissions. +- Configure Sub2API with the project id, location, managed bucket, provider account, model whitelist, and pricing. +- Enable `BATCH_IMAGE_ENABLED` globally, enable image generation on the intended Gemini group, then enable `allow_batch_image_generation` for that group. Non-Gemini groups are not eligible for batch image generation, and the admin UI only shows the batch image group switch after image generation is enabled on a Gemini group. + +API-key path: + +- Google API keys are suitable for Gemini API development and supported Gemini methods. +- The Sub2API `x-goog-api-key` compatibility header still expects a Sub2API key, not a plain Google key. +- Plain Google API keys should not be documented as the default production credential for Vertex service-account batch jobs. +- If an administrator configures a Gemini API-key upstream account, validate it with one low-cost batch image after the Google account has the required billing/prepayment state. If it has no prepayment, record only that the provider is selectable/callable and that failed submit releases hold. + +Official references: + +- Gemini API key guide: https://ai.google.dev/gemini-api/docs/api-key +- Gemini API Batch API: https://ai.google.dev/gemini-api/docs/batch-api +- Gemini API image generation and batch image notes: https://ai.google.dev/gemini-api/docs/image-generation +- Vertex/Gemini batch inference: https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/capabilities/batch-inference +- Vertex batch predictions API: https://docs.cloud.google.com/gemini-enterprise-agent-platform/reference/models/batch-prediction-api + +## Config + +These keys exist in `backend/internal/config/config.go`: + +```yaml +batch_image: + enabled: false + max_items_per_job_default: 200 + max_items_per_job_trial: 50 + max_output_images_per_job: 200 + max_output_images_per_item: 4 + max_prompt_chars_per_item: 8000 + max_reference_images_per_job: 1000 + max_reference_inline_bytes_per_job: 134217728 + default_response_mime_type: "image/png" + default_image_size: "1K" + + max_download_items_zip: 200 + max_download_bytes_per_request: 536870912 + max_download_duration_seconds: 600 + max_download_concurrency_per_user: 1 + + input_retention_after_terminal_hours: 24 + output_retention_after_terminal_hours: 72 + output_retention_max_days: 7 + cleanup_interval_minutes: 30 + cleanup_batch_size: 100 + + queue_enabled: false + queue_ready_key: "batch_image:queue:ready" + queue_delayed_key: "batch_image:queue:delayed" + queue_active_key: "batch_image:queue:active" + inflight_key_prefix: "batch_image:queue:inflight:" + lock_key_prefix: "batch_image:queue:lock:" + idempotency_key_prefix: "batch_image:queue:idem:" + inflight_ttl_seconds: 604800 + job_lock_ttl_seconds: 300 + default_requeue_delay_seconds: 30 + error_retry_delay_seconds: 60 + lock_conflict_delay_seconds: 5 + stale_active_after_seconds: 600 + delayed_mover_interval_seconds: 5 + recovery_interval_seconds: 300 + delayed_move_limit: 100 + recover_limit: 100 + + vertex_enabled: false + vertex_project_id: "" + vertex_location: "global" + vertex_managed_gcs_bucket: "" + vertex_managed_gcs_prefix: "batch-image/{env}/{batch_id}" + vertex_input_retention_hours: 24 + vertex_output_retention_hours: 72 + vertex_batch_prediction_base_url: "" + vertex_gcs_base_url: "" +``` + +Feature flags default to disabled. + +## Operations Checklist + +- Enable `batch_image.enabled`. +- Configure Redis. +- Enable `batch_image.queue_enabled` when workers should consume queue jobs. +- Configure provider accounts. +- Configure the Vertex managed GCS bucket if using Vertex. +- Ensure bucket permissions are correct. +- Disable or manage GCS soft delete. +- Configure cleanup worker settings. +- Configure max items per job. +- Configure download concurrency. +- Confirm billing pricing. +- Run smoke tests before enabling. + +## Future Optimization + +- Optional object-storage download offload: persist completed image outputs to an operator-configured object store such as GCS, S3, or R2, then issue short-lived signed download links to users. This would avoid routing large image/ZIP downloads through the Sub2API server, which is useful for small-bandwidth deployments. Keep it opt-in because it needs extra storage credentials, lifecycle cleanup, signed-URL expiry policy, access auditing, and compatibility with output deletion. + +## Security Checklist + +- No provider refs in public responses. +- No GCS URI exposure. +- No signed URL exposure. +- No service account exposure. +- No API key exposure. +- No image bytes/base64 in PostgreSQL. +- No base64 in logs. +- Owner-scoped status, item, download, cancel, and delete routes. +- Output deletion is owner-scoped. +- Cleanup paths are server-generated only. + +## Test Commands + +Core smoke and compile commands: + +```bash +go test -tags=unit ./internal/service -run 'BatchImage' -count=1 +go test -tags=unit ./internal/config ./internal/service ./internal/repository -count=1 +go test ./internal/config ./internal/service ./internal/repository ./internal/handler ./internal/server/routes -run '^$' +go test ./... -run '^$' +``` + +These commands should not require Docker, testcontainers, Redis, GCP, Gemini, Vertex, or GCS. + +## PR Hygiene Checklist + +- Do not accidentally commit `rfcs/batch-image-issue-draft.md` unless maintainers explicitly want it. +- Keep migrations ordered: `159_batch_image_foundation.sql`, then `160_batch_image_provider_refs.sql`, then later migrations. +- Include generated Ent code if generated code is committed in this repository. +- Keep generated server and wire files updated. +- Keep feature flags disabled by default unless maintainers ask otherwise. +- Do not commit real secrets, API keys, service account JSON, or local machine paths. +- Keep fixtures tiny and fake; no real cloud refs or credentials. +- Do not add new public routes, providers, dashboards, queues, or billing behavior in this stabilization PR. diff --git a/frontend/src/__tests__/integration/data-import.spec.ts b/frontend/src/__tests__/integration/data-import.spec.ts index bc9de148bd..1decee6760 100644 --- a/frontend/src/__tests__/integration/data-import.spec.ts +++ b/frontend/src/__tests__/integration/data-import.spec.ts @@ -1,14 +1,16 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' -import { mount } from '@vue/test-utils' +import { flushPromises, mount } from '@vue/test-utils' import ImportDataModal from '@/components/admin/account/ImportDataModal.vue' const showError = vi.fn() const showSuccess = vi.fn() +const showWarning = vi.fn() vi.mock('@/stores/app', () => ({ useAppStore: () => ({ showError, - showSuccess + showSuccess, + showWarning }) })) @@ -26,49 +28,187 @@ vi.mock('vue-i18n', () => ({ }) })) +const mountModal = () => + mount(ImportDataModal, { + props: { show: true }, + global: { + stubs: { + BaseDialog: { template: '
' } + } + } + }) + +const makeJsonFile = (name: string, content: string, type = 'application/json') => { + const file = new File([content], name, { type }) + Object.defineProperty(file, 'text', { + value: () => Promise.resolve(content) + }) + return file +} + +const setInputFiles = (element: Element, files: File[]) => { + Object.defineProperty(element, 'files', { + value: files, + configurable: true + }) +} + describe('ImportDataModal', () => { - beforeEach(() => { + beforeEach(async () => { showError.mockReset() showSuccess.mockReset() + showWarning.mockReset() + const { adminAPI } = await import('@/api/admin') + vi.mocked(adminAPI.accounts.importData).mockReset() }) it('未选择文件时提示错误', async () => { - const wrapper = mount(ImportDataModal, { - props: { show: true }, - global: { - stubs: { - BaseDialog: { template: '
' } - } - } - }) + const wrapper = mountModal() await wrapper.find('form').trigger('submit') expect(showError).toHaveBeenCalledWith('admin.accounts.dataImportSelectFile') }) - it('无效 JSON 时提示解析失败', async () => { - const wrapper = mount(ImportDataModal, { - props: { show: true }, - global: { - stubs: { - BaseDialog: { template: '
' } - } - } - }) + it('无效 JSON 时按文件名提示解析失败', async () => { + const { adminAPI } = await import('@/api/admin') + const wrapper = mountModal() const input = wrapper.find('input[type="file"]') - const file = new File(['invalid json'], 'data.json', { type: 'application/json' }) - Object.defineProperty(file, 'text', { - value: () => Promise.resolve('invalid json') - }) - Object.defineProperty(input.element, 'files', { - value: [file] - }) + setInputFiles(input.element, [makeJsonFile('data.json', 'invalid json')]) await input.trigger('change') await wrapper.find('form').trigger('submit') - await Promise.resolve() + await flushPromises() - expect(showError).toHaveBeenCalledWith('admin.accounts.dataImportParseFailed') + expect(showError).toHaveBeenCalledWith('admin.accounts.dataImportParseFailedFile') + expect(adminAPI.accounts.importData).not.toHaveBeenCalled() + }) + + it('不是导出数据的 JSON 按文件名拒绝', async () => { + const { adminAPI } = await import('@/api/admin') + const wrapper = mountModal() + + const input = wrapper.find('input[type="file"]') + setInputFiles(input.element, [makeJsonFile('random.json', JSON.stringify({ name: 'test' }))]) + + await input.trigger('change') + await wrapper.find('form').trigger('submit') + await flushPromises() + + expect(showError).toHaveBeenCalledWith('admin.accounts.dataImportInvalidFile') + expect(adminAPI.accounts.importData).not.toHaveBeenCalled() + }) + + it('无有效 JSON 的选择不清空已有选择', async () => { + const { adminAPI } = await import('@/api/admin') + vi.mocked(adminAPI.accounts.importData).mockResolvedValue({ + proxy_created: 0, + proxy_reused: 0, + proxy_failed: 0, + account_created: 1, + account_failed: 0 + }) + + const wrapper = mountModal() + const input = wrapper.find('input[type="file"]') + + const valid = makeJsonFile( + 'valid.json', + JSON.stringify({ exported_at: '2026-07-05T00:00:00Z', proxies: [], accounts: [{ name: 'a' }] }) + ) + setInputFiles(input.element, [valid]) + await input.trigger('change') + + setInputFiles(input.element, [new File(['hello'], 'notes.txt', { type: 'text/plain' })]) + await input.trigger('change') + expect(showError).toHaveBeenCalledWith('admin.accounts.dataImportSelectFile') + + await wrapper.find('form').trigger('submit') + await flushPromises() + + expect(adminAPI.accounts.importData).toHaveBeenCalledWith({ + data: expect.objectContaining({ + accounts: [{ name: 'a' }] + }), + skip_default_group_bind: true + }) + }) + + it('merges multiple selected JSON files before importing', async () => { + const { adminAPI } = await import('@/api/admin') + vi.mocked(adminAPI.accounts.importData).mockResolvedValue({ + proxy_created: 0, + proxy_reused: 0, + proxy_failed: 0, + account_created: 2, + account_failed: 0 + }) + + const wrapper = mountModal() + + const input = wrapper.find('input[type="file"]') + const first = makeJsonFile( + 'first.json', + JSON.stringify({ exported_at: '2026-07-05T00:00:00Z', proxies: [], accounts: [{ name: 'a' }] }) + ) + const second = makeJsonFile( + 'second.json', + JSON.stringify({ + exported_at: '2026-07-05T00:00:01Z', + proxies: [{ proxy_key: 'p' }], + accounts: [{ name: 'b' }] + }) + ) + setInputFiles(input.element, [first, second]) + + await input.trigger('change') + await wrapper.find('form').trigger('submit') + await flushPromises() + + expect(adminAPI.accounts.importData).toHaveBeenCalledWith({ + data: expect.objectContaining({ + proxies: [{ proxy_key: 'p' }], + accounts: [{ name: 'a' }, { name: 'b' }] + }), + skip_default_group_bind: true + }) + expect(showSuccess).toHaveBeenCalledWith('admin.accounts.dataImportSuccess') + }) + + it('部分成功时关闭弹窗仍通知父组件刷新', async () => { + const { adminAPI } = await import('@/api/admin') + vi.mocked(adminAPI.accounts.importData).mockResolvedValue({ + proxy_created: 0, + proxy_reused: 0, + proxy_failed: 0, + account_created: 1, + account_failed: 1 + }) + + const wrapper = mountModal() + const input = wrapper.find('input[type="file"]') + setInputFiles(input.element, [ + makeJsonFile( + 'mixed.json', + JSON.stringify({ + exported_at: '2026-07-05T00:00:00Z', + proxies: [], + accounts: [{ name: 'a' }, { name: 'b' }] + }) + ) + ]) + + await input.trigger('change') + await wrapper.find('form').trigger('submit') + await flushPromises() + + expect(showError).toHaveBeenCalledWith('admin.accounts.dataImportCompletedWithErrors') + expect(wrapper.emitted('imported')).toBeUndefined() + + // 第二个 btn-secondary 是 footer 的取消按钮(第一个是选择文件) + await wrapper.findAll('button.btn-secondary')[1]!.trigger('click') + + expect(wrapper.emitted('imported')).toHaveLength(1) + expect(wrapper.emitted('close')).toHaveLength(1) }) }) diff --git a/frontend/src/api/__tests__/admin.system.rollback.spec.ts b/frontend/src/api/__tests__/admin.system.rollback.spec.ts new file mode 100644 index 0000000000..15d0989ffb --- /dev/null +++ b/frontend/src/api/__tests__/admin.system.rollback.spec.ts @@ -0,0 +1,55 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { get, post } = vi.hoisted(() => ({ + get: vi.fn(), + post: vi.fn(), +})) + +vi.mock('../client', () => ({ + apiClient: { + get, + post, + }, +})) + +import { getRollbackVersions, rollback, type RollbackVersionInfo } from '@/api/admin/system' + +describe('admin system rollback API', () => { + beforeEach(() => { + get.mockReset() + post.mockReset() + }) + + it('getRollbackVersions fetches the rollback version list', async () => { + const versions: RollbackVersionInfo[] = [ + { + version: '0.1.146', + published_at: '2026-07-07T00:00:00Z', + html_url: 'https://github.com/Wei-Shaw/sub2api/releases/tag/v0.1.146' + } + ] + get.mockResolvedValue({ data: { versions } }) + + const result = await getRollbackVersions() + + expect(get).toHaveBeenCalledWith('/admin/system/rollback-versions') + expect(result.versions).toEqual(versions) + }) + + it('rollback posts the target version in the request body', async () => { + post.mockResolvedValue({ data: { message: 'ok', need_restart: true } }) + + const result = await rollback('0.1.146') + + expect(post).toHaveBeenCalledWith('/admin/system/rollback', { version: '0.1.146' }) + expect(result.need_restart).toBe(true) + }) + + it('rollback without a version posts no body (legacy backup rollback)', async () => { + post.mockResolvedValue({ data: { message: 'ok', need_restart: true } }) + + await rollback() + + expect(post).toHaveBeenCalledWith('/admin/system/rollback', undefined) + }) +}) diff --git a/frontend/src/api/admin/accounts.ts b/frontend/src/api/admin/accounts.ts index 7cdb092820..2f9625b430 100644 --- a/frontend/src/api/admin/accounts.ts +++ b/frontend/src/api/admin/accounts.ts @@ -41,6 +41,7 @@ export async function list( search?: string privacy_mode?: string lite?: string + include_scheduler_score?: string sort_by?: string sort_order?: 'asc' | 'desc' }, @@ -76,6 +77,7 @@ export async function listWithEtag( search?: string privacy_mode?: string lite?: string + include_scheduler_score?: string sort_by?: string sort_order?: 'asc' | 'desc' }, @@ -558,7 +560,9 @@ export async function syncFromCrs(params: { action: string error?: string }> - }>('/admin/accounts/sync/crs', params) + }>('/admin/accounts/sync/crs', params, { + timeout: 180000 // 180s timeout: sync refreshes each existing account's OAuth token serially + }) return data } diff --git a/frontend/src/api/admin/dashboard.ts b/frontend/src/api/admin/dashboard.ts index dda7d8927d..97c16aa362 100644 --- a/frontend/src/api/admin/dashboard.ts +++ b/frontend/src/api/admin/dashboard.ts @@ -167,6 +167,8 @@ export interface UserBreakdownParams { endpoint?: string endpoint_type?: 'inbound' | 'upstream' | 'path' limit?: number + // Sort column for the ranking (allowlisted server-side; falls back to actual_cost) + sort_by?: 'total_tokens' | 'input_tokens' | 'output_tokens' | 'cache_tokens' | 'requests' | 'cost' | 'actual_cost' // Additional filter conditions user_id?: number api_key_id?: number diff --git a/frontend/src/api/admin/grok.ts b/frontend/src/api/admin/grok.ts index 2e0f5349b0..c0055d4dcc 100644 --- a/frontend/src/api/admin/grok.ts +++ b/frontend/src/api/admin/grok.ts @@ -63,6 +63,7 @@ export interface GrokQuotaSnapshot { export interface GrokQuotaProbeResult { source: 'active_probe' + model: string snapshot?: GrokQuotaSnapshot | null status_code?: number headers_observed: boolean diff --git a/frontend/src/api/admin/settings.ts b/frontend/src/api/admin/settings.ts index f5da990930..6f69163994 100644 --- a/frontend/src/api/admin/settings.ts +++ b/frontend/src/api/admin/settings.ts @@ -1273,11 +1273,11 @@ export async function updateRectifierSettings( */ export interface OpenAIFastPolicyRule { service_tier: "all" | "priority" | "flex"; - action: "pass" | "filter" | "block"; + action: "pass" | "filter" | "block" | "force_priority"; scope: "all" | "oauth" | "apikey" | "bedrock"; error_message?: string; model_whitelist?: string[]; - fallback_action?: "pass" | "filter" | "block"; + fallback_action?: "pass" | "filter" | "block" | "force_priority"; fallback_error_message?: string; } diff --git a/frontend/src/api/admin/system.ts b/frontend/src/api/admin/system.ts index 9ea312d568..eb3db9b06c 100644 --- a/frontend/src/api/admin/system.ts +++ b/frontend/src/api/admin/system.ts @@ -45,6 +45,22 @@ export interface UpdateResult { need_restart: boolean } +export interface RollbackVersionInfo { + version: string + published_at: string + html_url: string +} + +/** + * Get versions available for rollback (up to 3 versions older than current) + */ +export async function getRollbackVersions(): Promise<{ versions: RollbackVersionInfo[] }> { + const { data } = await apiClient.get<{ versions: RollbackVersionInfo[] }>( + '/admin/system/rollback-versions' + ) + return data +} + /** * Perform system update * Downloads and applies the latest version @@ -55,10 +71,14 @@ export async function performUpdate(): Promise { } /** - * Rollback to previous version + * Rollback to a previous version + * @param version - Target version (e.g. "0.1.146"); omit to restore the local backup binary */ -export async function rollback(): Promise { - const { data } = await apiClient.post('/admin/system/rollback') +export async function rollback(version?: string): Promise { + const { data } = await apiClient.post( + '/admin/system/rollback', + version ? { version } : undefined + ) return data } @@ -74,6 +94,7 @@ export const systemAPI = { getVersion, checkUpdates, performUpdate, + getRollbackVersions, rollback, restartService } diff --git a/frontend/src/api/admin/users.ts b/frontend/src/api/admin/users.ts index 8ff022ff68..78b4f8177b 100644 --- a/frontend/src/api/admin/users.ts +++ b/frontend/src/api/admin/users.ts @@ -121,6 +121,7 @@ export async function create(userData: { password: string username?: string notes?: string + role?: 'admin' | 'user' balance?: number concurrency?: number rpm_limit?: number diff --git a/frontend/src/api/batchImage.ts b/frontend/src/api/batchImage.ts new file mode 100644 index 0000000000..e6a26d79bd --- /dev/null +++ b/frontend/src/api/batchImage.ts @@ -0,0 +1,245 @@ +import { buildGatewayUrl } from './client' + +export type BatchImageStatus = + | 'queued' + | 'running' + | 'indexing' + | 'processing_results' + | 'settling' + | 'completed' + | 'failed' + | 'cancelled' + | 'output_deleted' + | string + +export interface BatchImageSubmitItem { + custom_id: string + prompt: string + output_count?: number + reference_images?: BatchImageReferenceImage[] +} + +export interface BatchImageReferenceImage { + id?: string + type?: string + mime_type: string + data?: string + file_uri?: string +} + +export interface BatchImageSubmitRequest { + model: string + task_name?: string + parent_batch_id?: string + provider?: '' | 'gemini_api' | 'vertex' | string + image_size?: '1K' | '2K' | '4K' | string + response_mime_type?: string + aspect_ratio?: string + items: BatchImageSubmitItem[] + metadata?: Record +} + +export interface BatchImageJob { + id: string + object: string + task_name: string + parent_batch_id?: string | null + status: BatchImageStatus + model: string + provider: string + item_count: number + success_count: number + fail_count: number + estimated_cost: number + hold_amount: number + actual_cost: number | null + created_at: number + submitted_at: number | null + settled_at: number | null + downloaded_at?: number | null + output_deleted_at?: number | null +} + +export interface BatchImageItem { + batch_id?: string + source_task_name?: string + custom_id: string + status: string + prompt_preview?: string | null + mime_type: string | null + file_extension: string | null + image_count: number + error?: { + code: string + message: string + source?: 'provider' | 'system' | string + } | null +} + +export interface BatchImageItemsResponse { + object: string + data: BatchImageItem[] + has_more: boolean +} + +export interface BatchImageJobsResponse { + object: string + data: BatchImageJob[] + has_more: boolean +} + +export interface BatchImageModel { + id: string + object: string + provider: string +} + +export interface BatchImageModelsResponse { + object: string + data: BatchImageModel[] +} + +export interface BatchImageJobsListOptions { + limit?: number + cursor?: string + status?: string + taskName?: string + downloaded?: '' | 'true' | 'false' | string + from?: string + to?: string +} + +async function parseBatchImageError(response: Response): Promise { + try { + const body = await response.json() + const message = body?.error?.message || body?.message || response.statusText + const error = new Error(message) + ;(error as any).code = body?.error?.code || response.status + ;(error as any).status = response.status + ;(error as any).requestId = response.headers.get('X-Request-Id') || '' + return error + } catch { + const error = new Error(response.statusText || `HTTP ${response.status}`) + ;(error as any).code = response.status + ;(error as any).status = response.status + ;(error as any).requestId = response.headers.get('X-Request-Id') || '' + return error + } +} + +function authHeaders(apiKey: string, extra?: HeadersInit): HeadersInit { + return { + Authorization: `Bearer ${apiKey}`, + ...extra, + } +} + +export async function submitBatchImageJob( + apiKey: string, + payload: BatchImageSubmitRequest, + idempotencyKey: string, +): Promise { + const response = await fetch(buildGatewayUrl('/v1/images/batches'), { + method: 'POST', + headers: authHeaders(apiKey, { + 'Content-Type': 'application/json', + 'Idempotency-Key': idempotencyKey, + }), + body: JSON.stringify(payload), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.json() +} + +export async function getBatchImageJob(apiKey: string, batchId: string): Promise { + const response = await fetch(buildGatewayUrl(`/v1/images/batches/${encodeURIComponent(batchId)}`), { + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.json() +} + +export async function listBatchImageJobs(apiKey: string, options: number | BatchImageJobsListOptions = 20): Promise { + const params = new URLSearchParams() + if (typeof options === 'number') { + params.set('limit', String(options)) + } else { + params.set('limit', String(options.limit || 20)) + if (options.cursor) params.set('cursor', options.cursor) + if (options.status) params.set('status', options.status) + if (options.taskName) params.set('task_name', options.taskName) + if (options.downloaded) params.set('downloaded', options.downloaded) + if (options.from) params.set('from', options.from) + if (options.to) params.set('to', options.to) + } + const response = await fetch(buildGatewayUrl(`/v1/images/batches?${params.toString()}`), { + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.json() +} + +export async function listBatchImageModels(apiKey: string): Promise { + const response = await fetch(buildGatewayUrl('/v1/images/batches/models'), { + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.json() +} + +export async function listBatchImageItems( + apiKey: string, + batchId: string, + status = '', +): Promise { + const query = status ? `?status=${encodeURIComponent(status)}` : '' + const response = await fetch(buildGatewayUrl(`/v1/images/batches/${encodeURIComponent(batchId)}/items${query}`), { + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.json() +} + +export async function cancelBatchImageJob(apiKey: string, batchId: string): Promise { + const response = await fetch(buildGatewayUrl(`/v1/images/batches/${encodeURIComponent(batchId)}/cancel`), { + method: 'POST', + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.json() +} + +export async function downloadBatchImageZip(apiKey: string, batchId: string): Promise { + const response = await fetch(buildGatewayUrl(`/v1/images/batches/${encodeURIComponent(batchId)}/download`), { + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.blob() +} + +export async function getBatchImageItemContent(apiKey: string, batchId: string, customId: string, imageIndex = 0): Promise { + const response = await fetch(buildGatewayUrl(`/v1/images/batches/${encodeURIComponent(batchId)}/items/${encodeURIComponent(customId)}/content?image_index=${encodeURIComponent(String(imageIndex))}`), { + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) + return response.blob() +} + +export async function deleteBatchImageJobRecord(apiKey: string, batchId: string): Promise { + const response = await fetch(buildGatewayUrl(`/v1/images/batches/${encodeURIComponent(batchId)}`), { + method: 'DELETE', + headers: authHeaders(apiKey), + }) + if (!response.ok) throw await parseBatchImageError(response) +} + +export function saveBlob(blob: Blob, filename: string) { + const url = URL.createObjectURL(blob) + const link = document.createElement('a') + link.href = url + link.download = filename + document.body.appendChild(link) + link.click() + document.body.removeChild(link) + URL.revokeObjectURL(url) +} diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index c1b9abcd5e..5df969f188 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -205,7 +205,9 @@ apiClient.interceptors.response.use( const refreshResponse = await axios.post( `${getAPIBaseURL()}/auth/refresh`, { refresh_token: refreshToken }, - { headers: { 'Content-Type': 'application/json' } } + // 显式设置超时:裸 axios 默认无限等待,若刷新请求挂起会导致 isRefreshing + // 永远为 true,所有排队的 401 重试请求永久卡死,页面 loading 无法恢复。 + { headers: { 'Content-Type': 'application/json' }, timeout: 30000 } ) const refreshData = refreshResponse.data as ApiResponse<{ diff --git a/frontend/src/api/index.ts b/frontend/src/api/index.ts index 6702468d8e..71fa27e2a7 100644 --- a/frontend/src/api/index.ts +++ b/frontend/src/api/index.ts @@ -17,6 +17,7 @@ export { redeemAPI, type RedeemHistoryItem } from './redeem' export { paymentAPI } from './payment' export { userGroupsAPI } from './groups' export { userChannelsAPI } from './channels' +export * as batchImageAPI from './batchImage' export { totpAPI } from './totp' export { default as announcementsAPI } from './announcements' export { channelMonitorUserAPI } from './channelMonitor' diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index a4cb70a9c3..81d97efb9c 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -372,6 +372,13 @@ A ${{ formatWindowCost(grokLocalUsage) }} + + U ${{ formatWindowUserCost(grokLocalUsage) }} + { const formatWindowRequests = (stats: WindowStats) => formatCompactNumber(stats.requests, { allowBillions: false }) const formatWindowTokens = (stats: WindowStats) => formatCompactNumber(stats.tokens) const formatWindowCost = (stats: WindowStats) => stats.cost.toFixed(2) +const formatWindowUserCost = (stats: WindowStats) => (stats.user_cost ?? 0).toFixed(2) // 账户类型显示标签 const antigravityTierLabel = computed(() => { diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 55efc197a7..2df3cc07e2 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -566,6 +566,58 @@ describe('AccountUsageCell', () => { expect(badges.some(node => node.attributes('title') === 'usage.userBilled')).toBe(true) }) + it('Grok OAuth 会展示本地 user billed 用量并保留超限百分比', async () => { + getUsage.mockResolvedValue({ + grok_local_usage: { + requests: 4, + tokens: 1200, + cost: 0.12, + standard_cost: 0.12, + user_cost: 0.34 + }, + grok_request_quota: { + limit: 10, + remaining: -2, + reset_at: '2026-07-09T16:00:00Z' + }, + grok_quota_snapshot_state: 'observed' + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 3861, + platform: 'grok', + type: 'oauth', + extra: {} + }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'resetsAt', 'color'], + template: '
{{ label }}|{{ utilization }}|{{ resetsAt }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(getUsage).toHaveBeenCalledWith(3861) + expect(wrapper.text()).toContain('4 req') + expect(wrapper.text()).toContain('1.2K') + expect(wrapper.text()).toContain('A $0.12') + expect(wrapper.text()).toContain('U $0.34') + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokRequests|120|2026-07-09T16:00:00Z') + + const badges = wrapper.findAll('span[title]') + expect(badges.some(node => node.attributes('title') === 'usage.accountBilled')).toBe(true) + expect(badges.some(node => node.attributes('title') === 'usage.userBilled')).toBe(true) + }) + it('Key 账号在 today stats loading 时显示骨架屏', async () => { const wrapper = mount(AccountUsageCell, { props: { diff --git a/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts b/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts index d094f5366d..31f6e3bd26 100644 --- a/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts @@ -107,7 +107,7 @@ describe('BulkEditAccountModal', () => { expect(mappingTab).toBeTruthy() await mappingTab!.trigger('click') - expect(wrapper.text()).toContain('3.1-Flash-Image passthrough') + expect(wrapper.text()).toContain('3.1-Flash-Image透传') expect(wrapper.text()).toContain('3-Pro-Image→3.1') expect(wrapper.text()).not.toContain('GPT-5.3 Codex Spark') }) diff --git a/frontend/src/components/admin/account/ImportDataModal.vue b/frontend/src/components/admin/account/ImportDataModal.vue index 6c120be39a..a0bfe294f6 100644 --- a/frontend/src/components/admin/account/ImportDataModal.vue +++ b/frontend/src/components/admin/account/ImportDataModal.vue @@ -19,13 +19,23 @@
-
- {{ fileName || t('admin.accounts.dataImportSelectFile') }} +
+ {{ selectedFilesLabel || t('admin.accounts.dataImportSelectFile') }} +
+
+ JSON (.json) + · {{ fileListTitle }}
-
JSON (.json)
@@ -90,7 +101,7 @@ import { useI18n } from 'vue-i18n' import BaseDialog from '@/components/common/BaseDialog.vue' import { adminAPI } from '@/api/admin' import { useAppStore } from '@/stores/app' -import type { AdminDataImportResult } from '@/types' +import type { AdminDataImportResult, AdminDataPayload } from '@/types' interface Props { show: boolean @@ -108,11 +119,19 @@ const { t } = useI18n() const appStore = useAppStore() const importing = ref(false) -const file = ref(null) +const files = ref([]) +const dragDepth = ref(0) +const dragActive = computed(() => dragDepth.value > 0) +const hasCreatedData = ref(false) const result = ref(null) const fileInput = ref(null) -const fileName = computed(() => file.value?.name || '') +const selectedFilesLabel = computed(() => { + if (files.value.length === 0) return '' + if (files.value.length === 1) return files.value[0]?.name || '' + return t('admin.accounts.selectedCount', { count: files.value.length }) +}) +const fileListTitle = computed(() => files.value.map((item) => item.name).join(', ')) const errorItems = computed(() => result.value?.errors || []) @@ -120,7 +139,9 @@ watch( () => props.show, (open) => { if (open) { - file.value = null + files.value = [] + dragDepth.value = 0 + hasCreatedData.value = false result.value = null if (fileInput.value) { fileInput.value.value = '' @@ -135,14 +156,56 @@ const openFilePicker = () => { const handleFileChange = (event: Event) => { const target = event.target as HTMLInputElement - file.value = target.files?.[0] || null + setSelectedFiles(target.files) + target.value = '' } const handleClose = () => { if (importing.value) return + if (hasCreatedData.value) { + hasCreatedData.value = false + emit('imported') + } emit('close') } +const isJsonFile = (sourceFile: File) => { + const name = sourceFile.name.toLowerCase() + return name.endsWith('.json') || sourceFile.type === 'application/json' +} + +const setSelectedFiles = (sourceFiles: FileList | File[] | null | undefined) => { + if (importing.value) return + const incoming = Array.from(sourceFiles || []) + const picked = incoming.filter(isJsonFile) + if (!picked.length) { + appStore.showError(t('admin.accounts.dataImportSelectFile')) + return + } + if (picked.length < incoming.length) { + appStore.showWarning( + t('admin.accounts.dataImportIgnoredFiles', { count: incoming.length - picked.length }) + ) + } + files.value = picked + result.value = null +} + +const handleDragEnter = () => { + if (importing.value) return + dragDepth.value += 1 +} + +const handleDragLeave = () => { + dragDepth.value = Math.max(0, dragDepth.value - 1) +} + +const handleDrop = (event: DragEvent) => { + dragDepth.value = 0 + if (importing.value) return + setSelectedFiles(event.dataTransfer?.files) +} + const readFileAsText = async (sourceFile: File): Promise => { if (typeof sourceFile.text === 'function') { return sourceFile.text() @@ -161,16 +224,74 @@ const readFileAsText = async (sourceFile: File): Promise => { }) } +const SUPPORTED_DATA_TYPES = ['sub2api-data', 'sub2api-bundle'] +const SUPPORTED_DATA_VERSION = 1 + +// 与后端 validateDataHeader 对齐:合并前逐文件校验,避免坏文件混入合并 payload 后 +// 报错无法定位来源,或绕过后端本会对单文件做的 type/version 检查。 +const isValidDataPayload = (payload: unknown): payload is AdminDataPayload => { + if (!payload || typeof payload !== 'object' || Array.isArray(payload)) return false + const candidate = payload as Record + if ( + candidate.type !== undefined && + candidate.type !== '' && + !SUPPORTED_DATA_TYPES.includes(candidate.type as string) + ) { + return false + } + if ( + candidate.version !== undefined && + candidate.version !== 0 && + candidate.version !== SUPPORTED_DATA_VERSION + ) { + return false + } + return Array.isArray(candidate.proxies) && Array.isArray(candidate.accounts) +} + +const mergeDataPayloads = (payloads: AdminDataPayload[]): AdminDataPayload => { + const [firstPayload] = payloads + if (payloads.length === 1 && firstPayload) return firstPayload + + return { + type: payloads.find((item) => typeof item.type === 'string')?.type, + version: payloads.find((item) => typeof item.version === 'number')?.version, + exported_at: new Date().toISOString(), + proxies: payloads.flatMap((item) => item.proxies), + accounts: payloads.flatMap((item) => item.accounts), + skipped_shadows: payloads.reduce((sum, item) => { + const count = Number(item.skipped_shadows || 0) + return Number.isFinite(count) ? sum + count : sum + }, 0) + } +} + const handleImport = async () => { - if (!file.value) { + if (files.value.length === 0) { appStore.showError(t('admin.accounts.dataImportSelectFile')) return } importing.value = true try { - const text = await readFileAsText(file.value) - const dataPayload = JSON.parse(text) + const dataPayloads: AdminDataPayload[] = [] + for (const sourceFile of files.value) { + let parsed: unknown + try { + parsed = JSON.parse(await readFileAsText(sourceFile)) + } catch { + appStore.showError( + t('admin.accounts.dataImportParseFailedFile', { name: sourceFile.name }) + ) + return + } + if (!isValidDataPayload(parsed)) { + appStore.showError(t('admin.accounts.dataImportInvalidFile', { name: sourceFile.name })) + return + } + dataPayloads.push(parsed) + } + const dataPayload = mergeDataPayloads(dataPayloads) const res = await adminAPI.accounts.importData({ data: dataPayload, @@ -187,17 +308,17 @@ const handleImport = async () => { proxy_failed: res.proxy_failed, } if (res.account_failed > 0 || res.proxy_failed > 0) { + // 部分成功也创建了数据;弹窗关闭时通过 imported 通知父组件刷新列表 + if (res.account_created > 0 || res.proxy_created > 0) { + hasCreatedData.value = true + } appStore.showError(t('admin.accounts.dataImportCompletedWithErrors', msgParams)) } else { appStore.showSuccess(t('admin.accounts.dataImportSuccess', msgParams)) emit('imported') } } catch (error: any) { - if (error instanceof SyntaxError) { - appStore.showError(t('admin.accounts.dataImportParseFailed')) - } else { - appStore.showError(error?.message || t('admin.accounts.dataImportFailed')) - } + appStore.showError(error?.message || t('admin.accounts.dataImportFailed')) } finally { importing.value = false } diff --git a/frontend/src/components/admin/usage/UsageFilters.vue b/frontend/src/components/admin/usage/UsageFilters.vue index bb63d9b8ee..aa981ac187 100644 --- a/frontend/src/components/admin/usage/UsageFilters.vue +++ b/frontend/src/components/admin/usage/UsageFilters.vue @@ -1,5 +1,5 @@