diff --git a/.github/audit-exceptions.yml b/.github/audit-exceptions.yml index 4e05aae66b..2a89dd0761 100644 --- a/.github/audit-exceptions.yml +++ b/.github/audit-exceptions.yml @@ -5,14 +5,14 @@ exceptions: severity: high reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2023-30533)" mitigation: "Load only on export; restrict export permissions and data scope" - expires_on: "2026-07-06" + expires_on: "2026-10-06" owner: "security@your-domain" - package: xlsx advisory: "GHSA-5pgg-2g8v-p4x9" severity: high reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2024-22363)" mitigation: "Load only on export; restrict export permissions and data scope" - expires_on: "2026-07-06" + expires_on: "2026-10-06" owner: "security@your-domain" - package: lodash advisory: "GHSA-r5fr-rjxr-66jc" diff --git a/Dockerfile b/Dockerfile index bae531ac6e..13a6b8700d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,3 +1,4 @@ +# syntax=docker/dockerfile:1.7 # ============================================================================= # Sub2API Multi-Stage Dockerfile # ============================================================================= @@ -12,11 +13,13 @@ 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 9aafad6264..a2b1901422 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,11 @@ Please read the following carefully before using this project: 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.FUN +Thanks 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! @@ -145,6 +150,14 @@ Model authenticity: no content intervention or secondary filtering — experienc + +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 @@ -523,20 +536,20 @@ Additional security-related options are available in `config.yaml`: **⚠️ Security Warning: HTTP URL Configuration** -When `security.url_allowlist.enabled=false`, the system performs minimal URL validation by default, **rejecting HTTP URLs** and only allowing HTTPS. To allow HTTP URLs (e.g., for development or internal testing), you must explicitly set: +When `security.url_allowlist.enabled=false`, the system performs minimal URL validation and **allows HTTP URLs by default** (dev-friendly mode; Docker Compose deployments use the same default). For production, explicitly tighten this to HTTPS-only: ```yaml security: url_allowlist: enabled: false # Disable allowlist checks - allow_insecure_http: true # Allow HTTP URLs (⚠️ INSECURE) + allow_insecure_http: false # HTTPS only (recommended for production) ``` **Or via environment variable:** ```bash SECURITY_URL_ALLOWLIST_ENABLED=false -SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true +SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=false ``` **Risks of allowing HTTP:** @@ -550,7 +563,7 @@ SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true - ✅ Testing account connectivity before obtaining HTTPS - ❌ Production environments (use HTTPS only) -**Example error without this setting:** +**Example error for HTTP URLs when `allow_insecure_http: false` is set:** ``` Invalid base URL: invalid url scheme: http ``` diff --git a/README_CN.md b/README_CN.md index 1f8cac8b03..a65f25f64c 100644 --- a/README_CN.md +++ b/README_CN.md @@ -42,6 +42,11 @@ 感谢 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% 赠送额度！ @@ -148,6 +153,14 @@ + +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 小时真人技术支持（非机器人），快速响应开发者需求。 + + + @@ -568,20 +581,20 @@ gateway: **⚠️ 安全警告：HTTP URL 配置** -当 `security.url_allowlist.enabled=false` 时，系统默认执行最小 URL 校验，**拒绝 HTTP URL**，仅允许 HTTPS。要允许 HTTP URL（例如用于开发或内网测试），必须显式设置： +当 `security.url_allowlist.enabled=false` 时，系统仅执行最小 URL 校验，且**默认允许 HTTP URL**（开发友好模式，Docker Compose 部署的默认值一致）。生产环境建议显式收紧为仅允许 HTTPS： ```yaml security: url_allowlist: enabled: false # 禁用白名单检查 - allow_insecure_http: true # 允许 HTTP URL（⚠️ 不安全） + allow_insecure_http: false # 仅允许 HTTPS（生产环境推荐） ``` **或通过环境变量：** ```bash SECURITY_URL_ALLOWLIST_ENABLED=false -SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true +SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=false ``` **允许 HTTP 的风险：** @@ -595,7 +608,7 @@ SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true - ✅ 获取 HTTPS 前测试账号连通性 - ❌ 生产环境（仅使用 HTTPS） -**未设置此项时的错误示例：** +**设置 `allow_insecure_http: false` 后，HTTP URL 会返回如下错误：** ``` Invalid base URL: invalid url scheme: http ``` diff --git a/README_JA.md b/README_JA.md index 684e685c8f..00fe27680c 100644 --- a/README_JA.md +++ b/README_JA.md @@ -41,6 +41,11 @@ 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 から登録すると、すべてのチャージで永久 5% 割引をご利用いただけます。 + + AIGoCode AIGoCode のご支援に感謝します！AIGoCode は Claude Code、Codex、最新の Gemini モデルを統合したオールインワンプラットフォームで、安定的かつ効率的でコストパフォーマンスに優れた AI コーディングサービスを提供します。柔軟なサブスクリプションプラン、アカウント停止リスクゼロ、VPN 不要の直接アクセス、超高速レスポンスが特長です。AIGoCode は sub2api ユーザー向けに特別特典を用意しています：こちらのリンクから登録すると、初回チャージ時に 10% のボーナスクレジットを追加プレゼント！ @@ -143,6 +148,14 @@ + +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 時間の有人テクニカルサポート（ボットではありません）により、開発者のニーズに迅速に対応します。 + + + ## 概要 @@ -521,20 +534,20 @@ default: **⚠️ セキュリティ警告: HTTP URL 設定** -`security.url_allowlist.enabled=false` の場合、システムはデフォルトで最小限の URL バリデーションを行い、**HTTP URL を拒否**して HTTPS のみを許可します。HTTP URL を許可するには（開発環境や内部テスト用など）、以下を明示的に設定する必要があります: +`security.url_allowlist.enabled=false` の場合、システムは最小限の URL バリデーションのみを行い、**デフォルトで HTTP URL を許可**します（開発フレンドリーモード。Docker Compose デプロイのデフォルトも同じです）。本番環境では、以下のように明示的に HTTPS のみに制限することを推奨します: ```yaml security: url_allowlist: enabled: false # 許可リストチェックを無効化 - allow_insecure_http: true # HTTP URL を許可（⚠️ セキュリティリスクあり） + allow_insecure_http: false # HTTPS のみ許可（本番環境推奨） ``` **または環境変数で設定:** ```bash SECURITY_URL_ALLOWLIST_ENABLED=false -SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true +SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=false ``` **HTTP を許可するリスク:** @@ -548,7 +561,7 @@ SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true - ✅ HTTPS 取得前のアカウント接続テスト - ❌ 本番環境（HTTPS のみを使用） -**この設定なしで表示されるエラー例:** +**`allow_insecure_http: false` 設定時に HTTP URL で表示されるエラー例:** ``` Invalid base URL: invalid url scheme: http ``` 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/cmd/server/VERSION b/backend/cmd/server/VERSION index 3ef481c326..22a9c16f5e 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -0.1.144 +0.1.146 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 a6fb5266aa..e148c5c363 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -67,7 +67,11 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { serviceUserPlatformQuotaRepository := repository.NewUserPlatformQuotaServiceAdapter(userPlatformQuotaRepository) billingCacheService := service.ProvideBillingCacheService(billingCache, userRepository, userSubscriptionRepository, apiKeyRepository, userRPMCache, userGroupRateRepository, configConfig, serviceUserPlatformQuotaRepository) apiKeyCache := repository.NewAPIKeyCache(redisClient) - apiKeyService := service.ProvideAPIKeyService(apiKeyRepository, userRepository, groupRepository, userSubscriptionRepository, userGroupRateRepository, apiKeyCache, configConfig, billingCacheService) + concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig) + schedulerCache := repository.ProvideSchedulerCache(redisClient, configConfig) + accountRepository := repository.NewAccountRepository(client, db, schedulerCache) + concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig) + apiKeyService := service.ProvideAPIKeyService(apiKeyRepository, userRepository, groupRepository, userSubscriptionRepository, userGroupRateRepository, apiKeyCache, configConfig, billingCacheService, concurrencyService) apiKeyAuthCacheInvalidator := service.ProvideAPIKeyAuthCacheInvalidator(apiKeyService) promoService := service.NewPromoService(promoCodeRepository, userRepository, billingCacheService, client, apiKeyAuthCacheInvalidator) subscriptionService := service.NewSubscriptionService(groupRepository, userSubscriptionRepository, billingCacheService, client, configConfig) @@ -92,10 +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) - schedulerCache := repository.ProvideSchedulerCache(redisClient, configConfig) - accountRepository := repository.NewAccountRepository(client, db, schedulerCache) - concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig) - concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig) + 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..2a0eb4d3ac 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,10 @@ 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"` // 是否仅允许 Claude Code 客户端 ClaudeCodeOnly bool `json:"claude_code_only,omitempty"` // 非 Claude Code 请求降级使用的分组 ID @@ -205,9 +211,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.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: 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 +361,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 +400,18 @@ 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.FieldClaudeCodeOnly: if value, ok := values[i].(*sql.NullBool); !ok { return fmt.Errorf("unexpected type %T for field claude_code_only", values[i]) @@ -631,6 +655,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 +679,12 @@ 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("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..540ce8f9f5 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,10 @@ 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" // 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 +194,14 @@ var Columns = []string{ FieldMonthlyLimitUsd, FieldDefaultValidityDays, FieldAllowImageGeneration, + FieldAllowBatchImageGeneration, FieldImageRateIndependent, FieldImageRateMultiplier, FieldImagePrice1k, FieldImagePrice2k, FieldImagePrice4k, + FieldBatchImageDiscountMultiplier, + FieldBatchImageHoldMultiplier, FieldClaudeCodeOnly, FieldFallbackGroupID, FieldFallbackGroupIDOnInvalidRequest, @@ -277,10 +286,16 @@ 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 // 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 +427,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 +457,16 @@ 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() +} + // 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..a76d3a8783 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,16 @@ 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)) +} + // 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 +1140,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 +1350,86 @@ 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)) +} + // 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..9c635847d0 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,34 @@ 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 +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_c *GroupCreate) SetClaudeCodeOnly(v bool) *GroupCreate { _c.mutation.SetClaudeCodeOnly(v) @@ -736,6 +778,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 +790,14 @@ 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.ClaudeCodeOnly(); !ok { v := group.DefaultClaudeCodeOnly _c.mutation.SetClaudeCodeOnly(v) @@ -869,12 +923,21 @@ 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.ClaudeCodeOnly(); !ok { return &ValidationError{Name: "claude_code_only", err: errors.New(`ent: missing required field "Group.claude_code_only"`)} } @@ -1019,6 +1082,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 +1106,14 @@ 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.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) _node.ClaudeCodeOnly = value @@ -1537,6 +1612,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 +1726,42 @@ 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 +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsert) SetClaudeCodeOnly(v bool) *GroupUpsert { u.Set(group.FieldClaudeCodeOnly, v) @@ -2235,6 +2358,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 +2491,48 @@ 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() + }) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsertOne) SetClaudeCodeOnly(v bool) *GroupUpsertOne { return u.Update(func(s *GroupUpsert) { @@ -3153,6 +3332,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 +3465,48 @@ 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() + }) +} + // 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..6f1831b1ea 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,48 @@ 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 +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_u *GroupUpdate) SetClaudeCodeOnly(v bool) *GroupUpdate { _u.mutation.SetClaudeCodeOnly(v) @@ -1116,6 +1172,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 +1211,18 @@ 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.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) } @@ -1853,6 +1924,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 +2054,48 @@ 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 +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_u *GroupUpdateOne) SetClaudeCodeOnly(v bool) *GroupUpdateOne { _u.mutation.SetClaudeCodeOnly(v) @@ -2647,6 +2774,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 +2813,18 @@ 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.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..a584cbe39d 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,14 @@ 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: "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 +910,7 @@ var ( { Name: "group_sort_order", Unique: false, - Columns: []*schema.Column{GroupsColumns[32]}, + Columns: []*schema.Column{GroupsColumns[35]}, }, }, } @@ -1500,6 +1685,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 +1712,7 @@ var ( { Name: "user_status", Unique: false, - Columns: []*schema.Column{UsersColumns[9]}, + Columns: []*schema.Column{UsersColumns[10]}, }, { Name: "user_deleted_at", @@ -1799,6 +1985,9 @@ var ( AnnouncementReadsTable, AuthIdentitiesTable, AuthIdentityChannelsTable, + BatchImageEventsTable, + BatchImageItemsTable, + BatchImageJobsTable, ChannelMonitorsTable, ChannelMonitorDailyRollupsTable, ChannelMonitorHistoriesTable, @@ -1862,6 +2051,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..987ec4146b 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,10 @@ 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 claude_code_only *bool fallback_group_id *int64 addfallback_group_id *int64 @@ -16366,6 +21847,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 +22185,118 @@ 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 +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (m *GroupMutation) SetClaudeCodeOnly(b bool) { m.claude_code_only = &b @@ -17702,7 +23331,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, 42) if m.created_at != nil { fields = append(fields, group.FieldCreatedAt) } @@ -17760,6 +23389,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 +23407,12 @@ 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.claude_code_only != nil { fields = append(fields, group.FieldClaudeCodeOnly) } @@ -17866,6 +23504,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 +23516,10 @@ 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.FieldClaudeCodeOnly: return m.ClaudeCodeOnly() case group.FieldFallbackGroupID: @@ -17953,6 +23597,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 +23609,10 @@ 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.FieldClaudeCodeOnly: return m.OldClaudeCodeOnly(ctx) case group.FieldFallbackGroupID: @@ -18135,6 +23785,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 +23827,20 @@ 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.FieldClaudeCodeOnly: v, ok := value.(bool) if !ok { @@ -18313,6 +23984,12 @@ 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.addfallback_group_id != nil { fields = append(fields, group.FieldFallbackGroupID) } @@ -18353,6 +24030,10 @@ 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.FieldFallbackGroupID: return m.AddedFallbackGroupID() case group.FieldFallbackGroupIDOnInvalidRequest: @@ -18440,6 +24121,20 @@ 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.FieldFallbackGroupID: v, ok := value.(int64) if !ok { @@ -18621,6 +24316,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 +24334,12 @@ 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.FieldClaudeCodeOnly: m.ResetClaudeCodeOnly() return nil @@ -39204,6 +44908,8 @@ type UserMutation struct { role *string balance *float64 addbalance *float64 + frozen_balance *float64 + addfrozen_balance *float64 concurrency *int addconcurrency *int status *string @@ -39652,6 +45358,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 +46872,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 +46894,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 +46967,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 +47024,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 +47116,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 +47246,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 +47271,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 +47297,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 +47419,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..2c05c01c52 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,74 @@ 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) // groupDescClaudeCodeOnly is the schema descriptor for claude_code_only field. - groupDescClaudeCodeOnly := groupFields[21].Descriptor() + groupDescClaudeCodeOnly := groupFields[24].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[28].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[29].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[30].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[31].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[32].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[33].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[34].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[35].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[36].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[37].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[38].Descriptor() // group.DefaultRpmLimit holds the default value on creation for the rpm_limit field. group.DefaultRpmLimit = groupDescRpmLimit.Default.(int) idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin() @@ -1860,54 +2041,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..d675ca52f1 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,14 @@ 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("批量图片生成冻结价格比例，按普通生图原价乘以该比例冻结，结算后释放差额"), // Claude Code 客户端限制 (added by migration 029) field.Bool("claude_code_only"). 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/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/internal/config/config.go b/backend/internal/config/config.go index b8b880eab0..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"` @@ -969,6 +1020,9 @@ type GatewayOpenAIWSSchedulerScoreWeights struct { Reset float64 `mapstructure:"reset"` // QuotaHeadroom 倾向 7d 剩余额度更健康的账号；默认 0（关闭，不改变原有行为）。 QuotaHeadroom float64 `mapstructure:"quota_headroom"` + // PreviousResponse/SessionSticky 仅在开启 OpenAI 高级调度的粘性加权时生效。 + PreviousResponse float64 `mapstructure:"previous_response"` + SessionSticky float64 `mapstructure:"session_sticky"` } // GatewayOpenAISchedulerConfig OpenAI 高级调度器配置。 @@ -1729,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) @@ -1891,6 +1992,8 @@ func setDefaults() { viper.SetDefault("gateway.openai_ws.scheduler_score_weights.ttft", 0.5) viper.SetDefault("gateway.openai_ws.scheduler_score_weights.reset", 0.0) viper.SetDefault("gateway.openai_ws.scheduler_score_weights.quota_headroom", 0.0) + viper.SetDefault("gateway.openai_ws.scheduler_score_weights.previous_response", 5.0) + viper.SetDefault("gateway.openai_ws.scheduler_score_weights.session_sticky", 3.0) // OpenAI HTTP upstream protocol strategy viper.SetDefault("gateway.openai_http2.enabled", true) viper.SetDefault("gateway.openai_http2.allow_proxy_fallback_to_http1", true) @@ -2328,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") @@ -2674,7 +2832,9 @@ func (c *Config) Validate() error { c.Gateway.OpenAIWS.SchedulerScoreWeights.Queue < 0 || c.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate < 0 || c.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT < 0 || - c.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom < 0 { + c.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom < 0 || + c.Gateway.OpenAIWS.SchedulerScoreWeights.PreviousResponse < 0 || + c.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky < 0 { return fmt.Errorf("gateway.openai_ws.scheduler_score_weights.* must be non-negative") } weightSum := c.Gateway.OpenAIWS.SchedulerScoreWeights.Priority + 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 8d4a5dfea6..8c91245fbf 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -171,13 +171,29 @@ type CheckMixedChannelRequest struct { // AccountWithConcurrency extends Account with real-time concurrency info type AccountWithConcurrency struct { *dto.Account - CurrentConcurrency int `json:"current_concurrency"` + CurrentConcurrency int `json:"current_concurrency"` + SchedulerScore *AccountSchedulerScore `json:"scheduler_score,omitempty"` + SchedulerScores []AccountSchedulerGroupScore `json:"scheduler_scores,omitempty"` // 以下字段仅对 Anthropic OAuth/SetupToken 账号有效，且仅在启用相应功能时返回 CurrentWindowCost *float64 `json:"current_window_cost,omitempty"` // 当前窗口费用 ActiveSessions *int `json:"active_sessions,omitempty"` // 当前活跃会话数 CurrentRPM *int `json:"current_rpm,omitempty"` // 当前分钟 RPM 计数 } +type AccountSchedulerScore struct { + BaseScore float64 `json:"base_score"` + StickyScore float64 `json:"sticky_score"` + StickyScoreInfinity bool `json:"sticky_score_infinity"` + StickyWeightedEnabled bool `json:"sticky_weighted_enabled"` +} + +type AccountSchedulerGroupScore struct { + GroupID *int64 `json:"group_id"` + GroupName string `json:"group_name,omitempty"` + GroupPriority *int `json:"group_priority,omitempty"` + AccountSchedulerScore +} + const accountListGroupUngroupedQueryValue = "ungrouped" func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, account *service.Account) AccountWithConcurrency { @@ -226,6 +242,232 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac return item } +// scoreOpenAIAccountSchedulerPool 对池内 OpenAI 账号计算调度分数快照。 +// loadMap 为共享的账号负载数据（含池内全部账号即可，多余条目无害）；传 nil 时自行批查。 +func (h *AccountHandler) scoreOpenAIAccountSchedulerPool(ctx context.Context, accounts []service.Account, loadMap map[int64]*service.AccountLoadInfo) map[int64]AccountSchedulerScore { + if len(accounts) == 0 { + return nil + } + + openAIAccounts := make([]*service.Account, 0, len(accounts)) + for i := range accounts { + account := &accounts[i] + if account.Platform != service.PlatformOpenAI { + continue + } + openAIAccounts = append(openAIAccounts, account) + } + if len(openAIAccounts) == 0 { + return nil + } + + if loadMap == nil { + loadMap = h.fetchOpenAIAccountLoadMap(ctx, openAIAccounts) + } + + var scores map[int64]service.OpenAIAccountSchedulerScoreSnapshot + if h.rateLimitService != nil { + scores = h.rateLimitService.BuildOpenAIAccountSchedulerScoreSnapshot(ctx, openAIAccounts, loadMap) + } else { + scores = service.BuildOpenAIAccountSchedulerScoreSnapshot(openAIAccounts, loadMap) + } + result := make(map[int64]AccountSchedulerScore, len(scores)) + for accountID, score := range scores { + result[accountID] = AccountSchedulerScore{ + BaseScore: score.BaseScore, + StickyScore: score.StickyScore, + StickyScoreInfinity: score.StickyScoreInfinity, + StickyWeightedEnabled: score.StickyWeightedEnabled, + } + } + return result +} + +// fetchOpenAIAccountLoadMap 一次性批查给定 OpenAI 账号的负载数据； +// 失败时记录日志并返回空表（分数按零负载计算，属可接受降级）。 +func (h *AccountHandler) fetchOpenAIAccountLoadMap(ctx context.Context, openAIAccounts []*service.Account) map[int64]*service.AccountLoadInfo { + loadMap := map[int64]*service.AccountLoadInfo{} + if h.concurrencyService == nil || len(openAIAccounts) == 0 { + return loadMap + } + seen := make(map[int64]struct{}, len(openAIAccounts)) + loadReq := make([]service.AccountWithConcurrency, 0, len(openAIAccounts)) + for _, account := range openAIAccounts { + if account == nil { + continue + } + if _, ok := seen[account.ID]; ok { + continue + } + seen[account.ID] = struct{}{} + loadReq = append(loadReq, service.AccountWithConcurrency{ + ID: account.ID, + MaxConcurrency: account.EffectiveLoadFactor(), + }) + } + if batchLoad, err := h.concurrencyService.GetAccountsLoadBatch(ctx, loadReq); err != nil { + slog.Warn("openai_scheduler_score_load_batch_failed", "error", err) + } else if batchLoad != nil { + loadMap = batchLoad + } + return loadMap +} + +func (h *AccountHandler) buildOpenAIAccountSchedulerScores( + ctx context.Context, + accounts []service.Account, + filterPool []service.Account, +) (map[int64]*AccountSchedulerScore, map[int64][]AccountSchedulerGroupScore) { + if len(accounts) == 0 { + return nil, nil + } + if len(filterPool) == 0 { + filterPool = accounts + } + + pageOpenAIAccountIDs := make(map[int64]struct{}) + groupIDs := make(map[int64]struct{}) + for i := range accounts { + account := &accounts[i] + if account.Platform != service.PlatformOpenAI { + continue + } + pageOpenAIAccountIDs[account.ID] = struct{}{} + if len(account.AccountGroups) == 0 && len(account.GroupIDs) == 0 { + continue + } + for _, accountGroup := range account.AccountGroups { + if accountGroup.GroupID > 0 { + groupIDs[accountGroup.GroupID] = struct{}{} + } + } + for _, groupID := range account.GroupIDs { + if groupID > 0 { + groupIDs[groupID] = struct{}{} + } + } + } + if len(pageOpenAIAccountIDs) == 0 { + return nil, nil + } + + // 先取各分组池，再对"过滤池 ∪ 分组池"的账号并集做一次负载批查， + // 避免每个池各查一次 Redis 的 N+1。 + groupIDList := make([]int64, 0, len(groupIDs)) + for groupID := range groupIDs { + groupIDList = append(groupIDList, groupID) + } + sort.Slice(groupIDList, func(i, j int) bool { return groupIDList[i] < groupIDList[j] }) + + groupPools := make(map[int64][]service.Account, len(groupIDList)) + if h.adminService != nil { + for _, groupID := range groupIDList { + gid := groupID + pool, err := h.adminService.ListOpenAISchedulableAccountsForSchedulerScore(ctx, &gid) + if err != nil { + slog.Warn("openai_scheduler_group_score_pool_failed", "group_id", gid, "error", err) + continue + } + groupPools[gid] = pool + } + } + + loadUnion := make([]*service.Account, 0, len(filterPool)) + collectOpenAIAccounts := func(pool []service.Account) { + for i := range pool { + if pool[i].Platform == service.PlatformOpenAI { + loadUnion = append(loadUnion, &pool[i]) + } + } + } + collectOpenAIAccounts(filterPool) + for _, pool := range groupPools { + collectOpenAIAccounts(pool) + } + loadMap := h.fetchOpenAIAccountLoadMap(ctx, loadUnion) + + baseScores := make(map[int64]*AccountSchedulerScore) + for accountID, score := range h.scoreOpenAIAccountSchedulerPool(ctx, filterPool, loadMap) { + copiedScore := score + baseScores[accountID] = &copiedScore + } + + groupScoresByAccount := make(map[int64][]AccountSchedulerGroupScore) + scoreGroupPool := func(groupID *int64, groupNameByID map[int64]string, groupPriorityByAccount map[int64]int, pool []service.Account) { + if len(pool) == 0 { + return + } + scores := h.scoreOpenAIAccountSchedulerPool(ctx, pool, loadMap) + for accountID, schedulerScore := range scores { + if _, ok := pageOpenAIAccountIDs[accountID]; !ok { + continue + } + groupScore := AccountSchedulerGroupScore{ + GroupID: groupID, + AccountSchedulerScore: schedulerScore, + } + if groupID != nil { + groupScore.GroupName = groupNameByID[*groupID] + if priority, ok := groupPriorityByAccount[accountID]; ok { + groupScore.GroupPriority = &priority + } + } + groupScoresByAccount[accountID] = append(groupScoresByAccount[accountID], groupScore) + } + } + + for _, groupID := range groupIDList { + gid := groupID + pool, ok := groupPools[gid] + if !ok { + continue + } + groupNameByID := make(map[int64]string) + groupPriorityByAccount := make(map[int64]int) + for i := range pool { + account := &pool[i] + for _, accountGroup := range account.AccountGroups { + if accountGroup.GroupID != gid { + continue + } + groupPriorityByAccount[account.ID] = accountGroup.Priority + if accountGroup.Group != nil { + groupNameByID[gid] = accountGroup.Group.Name + } + } + } + scoreGroupPool(&gid, groupNameByID, groupPriorityByAccount, pool) + } + + for accountID := range groupScoresByAccount { + sort.SliceStable(groupScoresByAccount[accountID], func(i, j int) bool { + left := groupScoresByAccount[accountID][i] + right := groupScoresByAccount[accountID][j] + return *left.GroupID < *right.GroupID + }) + } + return baseScores, groupScoresByAccount +} + +func (h *AccountHandler) listAccountSchedulerScoreFilterPool( + ctx context.Context, + platform, accountType, status, search string, + groupID int64, + privacyMode string, +) []service.Account { + if h.adminService == nil || (platform != "" && platform != service.PlatformOpenAI) { + return nil + } + // 池只用于 OpenAI 分数计算（非 OpenAI 账号会在打分时被丢弃）， + // 无论列表页平台过滤为何，查询一律限定 openai，避免无过滤时全表扫描。 + accounts, err := h.adminService.ListAccountsForSchedulerScoreFilter(ctx, service.PlatformOpenAI, accountType, status, search, groupID, privacyMode) + if err != nil { + slog.Warn("openai_scheduler_filter_score_pool_failed", "error", err) + return nil + } + return accounts +} + // List handles listing all accounts with pagination // GET /api/v1/admin/accounts func (h *AccountHandler) List(c *gin.Context) { @@ -278,6 +520,20 @@ func (h *AccountHandler) List(c *gin.Context) { var windowCosts map[int64]float64 var activeSessions map[int64]int var rpmCounts map[int64]int + // 仅当前页存在 OpenAI 账号时才计算调度分数，避免为空结果付出池查询开销。 + var schedulerScores map[int64]*AccountSchedulerScore + var schedulerGroupScores map[int64][]AccountSchedulerGroupScore + pageHasOpenAIAccounts := false + for i := range accounts { + if accounts[i].Platform == service.PlatformOpenAI { + pageHasOpenAIAccounts = true + break + } + } + if pageHasOpenAIAccounts { + schedulerFilterPool := h.listAccountSchedulerScoreFilterPool(c.Request.Context(), platform, accountType, status, search, groupID, privacyMode) + schedulerScores, schedulerGroupScores = h.buildOpenAIAccountSchedulerScores(c.Request.Context(), accounts, schedulerFilterPool) + } // 始终获取并发数（Redis ZCARD，极低开销） if h.concurrencyService != nil { @@ -358,6 +614,8 @@ func (h *AccountHandler) List(c *gin.Context) { item := AccountWithConcurrency{ Account: dto.AccountFromService(acc), CurrentConcurrency: concurrencyCounts[acc.ID], + SchedulerScore: schedulerScores[acc.ID], + SchedulerScores: schedulerGroupScores[acc.ID], } // 添加窗口费用（仅当启用时） diff --git a/backend/internal/handler/admin/account_handler_list_test.go b/backend/internal/handler/admin/account_handler_list_test.go index 4d628365df..29e36ad865 100644 --- a/backend/internal/handler/admin/account_handler_list_test.go +++ b/backend/internal/handler/admin/account_handler_list_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -50,3 +51,222 @@ func TestAccountHandlerListIncludesCreatedAt(t *testing.T) { _, offset := parsed.Zone() require.Equal(t, 0, offset) } + +func TestAccountHandlerListReturnsSchedulerScoresPerGroup(t *testing.T) { + router, adminSvc := setupAccountListRouter() + now := time.Now().UTC() + groupID := int64(41) + adminSvc.accounts = []service.Account{ + { + ID: 101, + Name: "account-high-priority", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 1, + AccountGroups: []service.AccountGroup{ + {AccountID: 101, GroupID: groupID, Priority: 100, Group: &service.Group{ID: groupID, Name: "openai"}}, + }, + GroupIDs: []int64{groupID}, + CreatedAt: now, + UpdatedAt: now, + }, + { + ID: 102, + Name: "account-low-priority", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 100000, + AccountGroups: []service.AccountGroup{ + {AccountID: 102, GroupID: groupID, Priority: 1, Group: &service.Group{ID: groupID, Name: "openai"}}, + }, + GroupIDs: []int64{groupID}, + 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) + var payload struct { + Data struct { + Items []struct { + ID int64 `json:"id"` + SchedulerScore struct { + BaseScore float64 `json:"base_score"` + } `json:"scheduler_score"` + SchedulerScores []struct { + GroupID *int64 `json:"group_id"` + GroupName string `json:"group_name"` + GroupPriority *int `json:"group_priority"` + BaseScore float64 `json:"base_score"` + } `json:"scheduler_scores"` + } `json:"items"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload)) + require.Len(t, payload.Data.Items, 2) + + var high, low *struct { + ID int64 `json:"id"` + SchedulerScore struct { + BaseScore float64 `json:"base_score"` + } `json:"scheduler_score"` + SchedulerScores []struct { + GroupID *int64 `json:"group_id"` + GroupName string `json:"group_name"` + GroupPriority *int `json:"group_priority"` + BaseScore float64 `json:"base_score"` + } `json:"scheduler_scores"` + } + for i := range payload.Data.Items { + item := &payload.Data.Items[i] + switch item.ID { + case 101: + high = item + case 102: + low = item + } + } + require.NotNil(t, high) + require.NotNil(t, low) + require.Len(t, high.SchedulerScores, 1) + require.Len(t, low.SchedulerScores, 1) + require.Equal(t, groupID, *high.SchedulerScores[0].GroupID) + require.Equal(t, "openai", high.SchedulerScores[0].GroupName) + require.Equal(t, 100, *high.SchedulerScores[0].GroupPriority) + require.Equal(t, 1, *low.SchedulerScores[0].GroupPriority) + require.Greater(t, high.SchedulerScores[0].BaseScore, low.SchedulerScores[0].BaseScore) +} + +func TestAccountHandlerListKeepsSchedulerScoreScopedToFilter(t *testing.T) { + router, adminSvc := setupAccountListRouter() + now := time.Now().UTC() + groupID := int64(42) + visibleAccount := service.Account{ + ID: 201, + Name: "visible-low-priority", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 100000, + AccountGroups: []service.AccountGroup{ + {AccountID: 201, GroupID: groupID, Priority: 1, Group: &service.Group{ID: groupID, Name: "openai"}}, + }, + GroupIDs: []int64{groupID}, + CreatedAt: now, + UpdatedAt: now, + } + hiddenGroupPeer := service.Account{ + ID: 202, + Name: "hidden-high-priority", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 1, + AccountGroups: []service.AccountGroup{ + {AccountID: 202, GroupID: groupID, Priority: 2, Group: &service.Group{ID: groupID, Name: "openai"}}, + }, + GroupIDs: []int64{groupID}, + CreatedAt: now, + UpdatedAt: now, + } + adminSvc.accounts = []service.Account{visibleAccount} + adminSvc.accountSchedulerScoreFilterAccounts = []service.Account{visibleAccount, hiddenGroupPeer} + 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) + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + var payload struct { + Data struct { + Items []struct { + ID int64 `json:"id"` + SchedulerScore struct { + BaseScore float64 `json:"base_score"` + } `json:"scheduler_score"` + SchedulerScores []struct { + GroupID *int64 `json:"group_id"` + BaseScore float64 `json:"base_score"` + } `json:"scheduler_scores"` + } `json:"items"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload)) + require.Len(t, payload.Data.Items, 1) + item := payload.Data.Items[0] + require.Equal(t, int64(201), item.ID) + require.Len(t, item.SchedulerScores, 1) + require.Equal(t, groupID, *item.SchedulerScores[0].GroupID) + require.Equal(t, item.SchedulerScores[0].BaseScore, item.SchedulerScore.BaseScore) +} + +func TestAccountHandlerListSchedulerScoreIgnoresPagination(t *testing.T) { + router, adminSvc := setupAccountListRouter() + now := time.Now().UTC() + visibleAccount := service.Account{ + ID: 301, + Name: "visible-low-priority", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 100000, + CreatedAt: now, + UpdatedAt: now, + } + hiddenFilterPeer := service.Account{ + ID: 302, + Name: "hidden-high-priority", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 10, + Priority: 1, + CreatedAt: now, + UpdatedAt: now, + } + adminSvc.accounts = []service.Account{visibleAccount} + 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) + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + var payload struct { + Data struct { + Items []struct { + ID int64 `json:"id"` + SchedulerScore struct { + BaseScore float64 `json:"base_score"` + } `json:"scheduler_score"` + SchedulerScores []struct { + GroupID *int64 `json:"group_id"` + BaseScore float64 `json:"base_score"` + } `json:"scheduler_scores"` + } `json:"items"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload)) + require.Len(t, payload.Data.Items, 1) + require.Equal(t, int64(301), payload.Data.Items[0].ID) + require.Less(t, payload.Data.Items[0].SchedulerScore.BaseScore, 3.75) + require.Empty(t, payload.Data.Items[0].SchedulerScores) +} diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index f7bab13253..187925e33a 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -10,27 +10,29 @@ import ( ) type stubAdminService struct { - users []service.User - apiKeys []service.APIKey - groups []service.Group - accounts []service.Account - proxies []service.Proxy - proxyCounts []service.ProxyWithAccountCount - redeems []service.RedeemCode - boundAuthIdentity *service.AdminBindAuthIdentityInput - boundAuthIdentityFor int64 - createdAccounts []*service.CreateAccountInput - createdProxies []*service.CreateProxyInput - updatedProxyIDs []int64 - updatedProxies []*service.UpdateProxyInput - testedProxyIDs []int64 - getUserErr error - createAccountErr error - createSparkShadowErr error - updateAccountErr error - bulkUpdateAccountErr error - checkMixedErr error - lastMixedCheck struct { + users []service.User + apiKeys []service.APIKey + groups []service.Group + accounts []service.Account + accountSchedulerScoreFilterAccounts []service.Account + openAISchedulerScorePoolAccounts []service.Account + proxies []service.Proxy + proxyCounts []service.ProxyWithAccountCount + redeems []service.RedeemCode + boundAuthIdentity *service.AdminBindAuthIdentityInput + boundAuthIdentityFor int64 + createdAccounts []*service.CreateAccountInput + createdProxies []*service.CreateProxyInput + updatedProxyIDs []int64 + updatedProxies []*service.UpdateProxyInput + testedProxyIDs []int64 + getUserErr error + createAccountErr error + createSparkShadowErr error + updateAccountErr error + bulkUpdateAccountErr error + checkMixedErr error + lastMixedCheck struct { accountID int64 platform string groupIDs []int64 @@ -329,7 +331,56 @@ func (s *stubAdminService) ListAccounts(ctx context.Context, page, pageSize int, s.lastListAccounts.sortBy = sortBy s.lastListAccounts.sortOrder = sortOrder s.lastListAccounts.calls++ - return s.accounts, int64(len(s.accounts)), nil + accounts := s.accounts + total := len(accounts) + if page < 1 { + page = 1 + } + if pageSize < 1 { + pageSize = total + } + start := (page - 1) * pageSize + if start >= total { + return []service.Account{}, int64(total), nil + } + end := start + pageSize + if end > total { + end = total + } + return accounts[start:end], int64(total), nil +} + +func (s *stubAdminService) ListAccountsForSchedulerScoreFilter(_ context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, error) { + if s.accountSchedulerScoreFilterAccounts != nil { + return s.accountSchedulerScoreFilterAccounts, nil + } + return s.accounts, nil +} + +func (s *stubAdminService) ListOpenAISchedulableAccountsForSchedulerScore(_ context.Context, groupID *int64) ([]service.Account, error) { + accounts := s.openAISchedulerScorePoolAccounts + if accounts == nil { + accounts = s.accounts + } + out := make([]service.Account, 0, len(accounts)) + for _, account := range accounts { + if account.Platform != service.PlatformOpenAI || !account.IsSchedulable() { + continue + } + if groupID == nil { + if len(account.AccountGroups) == 0 && len(account.GroupIDs) == 0 { + out = append(out, account) + } + continue + } + for _, accountGroup := range account.AccountGroups { + if accountGroup.GroupID == *groupID { + out = append(out, account) + break + } + } + } + return out, nil } func (s *stubAdminService) GetAccount(ctx context.Context, id int64) (*service.Account, error) { diff --git a/backend/internal/handler/admin/group_handler.go b/backend/internal/handler/admin/group_handler.go index 0a98ad6784..4595adeb24 100644 --- a/backend/internal/handler/admin/group_handler.go +++ b/backend/internal/handler/admin/group_handler.go @@ -93,8 +93,11 @@ 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"` PeakRateEnabled bool `json:"peak_rate_enabled"` PeakStart string `json:"peak_start"` PeakEnd string `json:"peak_end"` @@ -138,8 +141,11 @@ 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"` PeakRateEnabled *bool `json:"peak_rate_enabled"` PeakStart *string `json:"peak_start"` PeakEnd *string `json:"peak_end"` @@ -301,8 +307,11 @@ 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, PeakRateEnabled: req.PeakRateEnabled, PeakStart: req.PeakStart, PeakEnd: req.PeakEnd, @@ -361,8 +370,11 @@ 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, PeakRateEnabled: req.PeakRateEnabled, PeakStart: req.PeakStart, PeakEnd: req.PeakEnd, diff --git a/backend/internal/handler/admin/setting_handler.go b/backend/internal/handler/admin/setting_handler.go index 0f50ab723b..529e46c575 100644 --- a/backend/internal/handler/admin/setting_handler.go +++ b/backend/internal/handler/admin/setting_handler.go @@ -119,188 +119,211 @@ func (h *SettingHandler) GetSettings(c *gin.Context) { } payload := dto.SystemSettings{ - RegistrationEnabled: settings.RegistrationEnabled, - EmailVerifyEnabled: settings.EmailVerifyEnabled, - RegistrationEmailSuffixWhitelist: settings.RegistrationEmailSuffixWhitelist, - PromoCodeEnabled: settings.PromoCodeEnabled, - PasswordResetEnabled: settings.PasswordResetEnabled, - FrontendURL: settings.FrontendURL, - InvitationCodeEnabled: settings.InvitationCodeEnabled, - TotpEnabled: settings.TotpEnabled, - TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(), - LoginAgreementEnabled: settings.LoginAgreementEnabled, - LoginAgreementMode: settings.LoginAgreementMode, - LoginAgreementUpdatedAt: settings.LoginAgreementUpdatedAt, - LoginAgreementDocuments: loginAgreementDocumentsToDTO(settings.LoginAgreementDocuments), - SMTPHost: settings.SMTPHost, - SMTPPort: settings.SMTPPort, - SMTPUsername: settings.SMTPUsername, - SMTPPasswordConfigured: settings.SMTPPasswordConfigured, - SMTPFrom: settings.SMTPFrom, - SMTPFromName: settings.SMTPFromName, - SMTPUseTLS: settings.SMTPUseTLS, - TurnstileEnabled: settings.TurnstileEnabled, - TurnstileSiteKey: settings.TurnstileSiteKey, - TurnstileSecretKeyConfigured: settings.TurnstileSecretKeyConfigured, - APIKeyACLTrustForwardedIP: settings.APIKeyACLTrustForwardedIP, - LinuxDoConnectEnabled: settings.LinuxDoConnectEnabled, - LinuxDoConnectClientID: settings.LinuxDoConnectClientID, - LinuxDoConnectClientSecretConfigured: settings.LinuxDoConnectClientSecretConfigured, - LinuxDoConnectRedirectURL: settings.LinuxDoConnectRedirectURL, - DingTalkConnectEnabled: settings.DingTalkConnectEnabled, - DingTalkConnectClientID: settings.DingTalkConnectClientID, - DingTalkConnectClientSecretConfigured: settings.DingTalkConnectClientSecretConfigured, - DingTalkConnectRedirectURL: settings.DingTalkConnectRedirectURL, - DingTalkConnectCorpRestrictionPolicy: settings.DingTalkConnectCorpRestrictionPolicy, - DingTalkConnectInternalCorpID: settings.DingTalkConnectInternalCorpID, - DingTalkConnectBypassRegistration: settings.DingTalkConnectBypassRegistration, - DingTalkConnectSyncCorpEmail: settings.DingTalkConnectSyncCorpEmail, - DingTalkConnectSyncDisplayName: settings.DingTalkConnectSyncDisplayName, - DingTalkConnectSyncDept: settings.DingTalkConnectSyncDept, - DingTalkConnectSyncCorpEmailAttrKey: settings.DingTalkConnectSyncCorpEmailAttrKey, - DingTalkConnectSyncDisplayNameAttrKey: settings.DingTalkConnectSyncDisplayNameAttrKey, - DingTalkConnectSyncDeptAttrKey: settings.DingTalkConnectSyncDeptAttrKey, - DingTalkConnectSyncCorpEmailAttrName: settings.DingTalkConnectSyncCorpEmailAttrName, - DingTalkConnectSyncDisplayNameAttrName: settings.DingTalkConnectSyncDisplayNameAttrName, - DingTalkConnectSyncDeptAttrName: settings.DingTalkConnectSyncDeptAttrName, - WeChatConnectEnabled: settings.WeChatConnectEnabled, - WeChatConnectAppID: settings.WeChatConnectAppID, - WeChatConnectAppSecretConfigured: settings.WeChatConnectAppSecretConfigured, - WeChatConnectOpenAppID: settings.WeChatConnectOpenAppID, - WeChatConnectOpenAppSecretConfigured: settings.WeChatConnectOpenAppSecretConfigured, - WeChatConnectMPAppID: settings.WeChatConnectMPAppID, - WeChatConnectMPAppSecretConfigured: settings.WeChatConnectMPAppSecretConfigured, - WeChatConnectMobileAppID: settings.WeChatConnectMobileAppID, - WeChatConnectMobileAppSecretConfigured: settings.WeChatConnectMobileAppSecretConfigured, - WeChatConnectOpenEnabled: settings.WeChatConnectOpenEnabled, - WeChatConnectMPEnabled: settings.WeChatConnectMPEnabled, - WeChatConnectMobileEnabled: settings.WeChatConnectMobileEnabled, - WeChatConnectMode: settings.WeChatConnectMode, - WeChatConnectScopes: settings.WeChatConnectScopes, - WeChatConnectRedirectURL: settings.WeChatConnectRedirectURL, - WeChatConnectFrontendRedirectURL: settings.WeChatConnectFrontendRedirectURL, - OIDCConnectEnabled: settings.OIDCConnectEnabled, - OIDCConnectProviderName: settings.OIDCConnectProviderName, - OIDCConnectClientID: settings.OIDCConnectClientID, - OIDCConnectClientSecretConfigured: settings.OIDCConnectClientSecretConfigured, - OIDCConnectIssuerURL: settings.OIDCConnectIssuerURL, - OIDCConnectDiscoveryURL: settings.OIDCConnectDiscoveryURL, - OIDCConnectAuthorizeURL: settings.OIDCConnectAuthorizeURL, - OIDCConnectTokenURL: settings.OIDCConnectTokenURL, - OIDCConnectUserInfoURL: settings.OIDCConnectUserInfoURL, - OIDCConnectJWKSURL: settings.OIDCConnectJWKSURL, - OIDCConnectScopes: settings.OIDCConnectScopes, - OIDCConnectRedirectURL: settings.OIDCConnectRedirectURL, - OIDCConnectFrontendRedirectURL: settings.OIDCConnectFrontendRedirectURL, - OIDCConnectTokenAuthMethod: settings.OIDCConnectTokenAuthMethod, - OIDCConnectUsePKCE: settings.OIDCConnectUsePKCE, - OIDCConnectValidateIDToken: settings.OIDCConnectValidateIDToken, - OIDCConnectAllowedSigningAlgs: settings.OIDCConnectAllowedSigningAlgs, - OIDCConnectClockSkewSeconds: settings.OIDCConnectClockSkewSeconds, - OIDCConnectRequireEmailVerified: settings.OIDCConnectRequireEmailVerified, - OIDCConnectUserInfoEmailPath: settings.OIDCConnectUserInfoEmailPath, - OIDCConnectUserInfoIDPath: settings.OIDCConnectUserInfoIDPath, - OIDCConnectUserInfoUsernamePath: settings.OIDCConnectUserInfoUsernamePath, - GitHubOAuthEnabled: settings.GitHubOAuthEnabled, - GitHubOAuthClientID: settings.GitHubOAuthClientID, - GitHubOAuthClientSecretConfigured: settings.GitHubOAuthClientSecretConfigured, - GitHubOAuthRedirectURL: settings.GitHubOAuthRedirectURL, - GitHubOAuthFrontendRedirectURL: settings.GitHubOAuthFrontendRedirectURL, - GoogleOAuthEnabled: settings.GoogleOAuthEnabled, - GoogleOAuthClientID: settings.GoogleOAuthClientID, - GoogleOAuthClientSecretConfigured: settings.GoogleOAuthClientSecretConfigured, - GoogleOAuthRedirectURL: settings.GoogleOAuthRedirectURL, - GoogleOAuthFrontendRedirectURL: settings.GoogleOAuthFrontendRedirectURL, - SiteName: settings.SiteName, - SiteLogo: settings.SiteLogo, - SiteSubtitle: settings.SiteSubtitle, - APIBaseURL: settings.APIBaseURL, - ContactInfo: settings.ContactInfo, - DocURL: settings.DocURL, - HomeContent: settings.HomeContent, - HideCcsImportButton: settings.HideCcsImportButton, - PurchaseSubscriptionEnabled: settings.PurchaseSubscriptionEnabled, - PurchaseSubscriptionURL: settings.PurchaseSubscriptionURL, - TableDefaultPageSize: settings.TableDefaultPageSize, - TablePageSizeOptions: settings.TablePageSizeOptions, - CustomMenuItems: dto.ParseCustomMenuItems(settings.CustomMenuItems), - CustomEndpoints: dto.ParseCustomEndpoints(settings.CustomEndpoints), - DefaultConcurrency: settings.DefaultConcurrency, - DefaultBalance: settings.DefaultBalance, - RiskControlEnabled: settings.RiskControlEnabled, - CyberSessionBlockEnabled: settings.CyberSessionBlockEnabled, - CyberSessionBlockTTLSeconds: settings.CyberSessionBlockTTLSeconds, - AffiliateRebateRate: settings.AffiliateRebateRate, - AffiliateRebateFreezeHours: settings.AffiliateRebateFreezeHours, - AffiliateRebateDurationDays: settings.AffiliateRebateDurationDays, - AffiliateRebatePerInviteeCap: settings.AffiliateRebatePerInviteeCap, - DefaultUserRPMLimit: settings.DefaultUserRPMLimit, - DefaultSubscriptions: defaultSubscriptions, - EnableModelFallback: settings.EnableModelFallback, - FallbackModelAnthropic: settings.FallbackModelAnthropic, - FallbackModelOpenAI: settings.FallbackModelOpenAI, - FallbackModelGemini: settings.FallbackModelGemini, - FallbackModelAntigravity: settings.FallbackModelAntigravity, - EnableIdentityPatch: settings.EnableIdentityPatch, - IdentityPatchPrompt: settings.IdentityPatchPrompt, - OpsMonitoringEnabled: opsEnabled && settings.OpsMonitoringEnabled, - OpsRealtimeMonitoringEnabled: settings.OpsRealtimeMonitoringEnabled, - OpsQueryModeDefault: settings.OpsQueryModeDefault, - OpsMetricsIntervalSeconds: settings.OpsMetricsIntervalSeconds, - MinClaudeCodeVersion: settings.MinClaudeCodeVersion, - MaxClaudeCodeVersion: settings.MaxClaudeCodeVersion, - AllowUngroupedKeyScheduling: settings.AllowUngroupedKeyScheduling, - BackendModeEnabled: settings.BackendModeEnabled, - EnableFingerprintUnification: settings.EnableFingerprintUnification, - EnableMetadataPassthrough: settings.EnableMetadataPassthrough, - EnableCCHSigning: settings.EnableCCHSigning, - EnableClaudeOAuthSystemPromptInjection: settings.EnableClaudeOAuthSystemPromptInjection, - ClaudeOAuthSystemPrompt: settings.ClaudeOAuthSystemPrompt, - ClaudeOAuthSystemPromptBlocks: settings.ClaudeOAuthSystemPromptBlocks, - EnableAnthropicCacheTTL1hInjection: settings.EnableAnthropicCacheTTL1hInjection, - RewriteMessageCacheControl: settings.RewriteMessageCacheControl, - EnableClientDatelineNormalization: settings.EnableClientDatelineNormalization, - AntigravityUserAgentVersion: settings.AntigravityUserAgentVersion, - OpenAICodexUserAgent: settings.OpenAICodexUserAgent, - MinCodexVersion: settings.MinCodexVersion, - MaxCodexVersion: settings.MaxCodexVersion, - CodexCLIOnlyBlacklist: settings.CodexCLIOnlyBlacklist, - CodexCLIOnlyWhitelist: settings.CodexCLIOnlyWhitelist, - CodexCLIOnlyAllowAppServerClients: settings.CodexCLIOnlyAllowAppServerClients, - CodexCLIOnlyEngineFingerprintSignals: settings.CodexCLIOnlyEngineFingerprintSignals, - WebSearchEmulationEnabled: settings.WebSearchEmulationEnabled, - PaymentVisibleMethodAlipaySource: settings.PaymentVisibleMethodAlipaySource, - PaymentVisibleMethodWxpaySource: settings.PaymentVisibleMethodWxpaySource, - PaymentVisibleMethodAlipayEnabled: settings.PaymentVisibleMethodAlipayEnabled, - PaymentVisibleMethodWxpayEnabled: settings.PaymentVisibleMethodWxpayEnabled, - OpenAIAdvancedSchedulerEnabled: settings.OpenAIAdvancedSchedulerEnabled, - BalanceLowNotifyEnabled: settings.BalanceLowNotifyEnabled, - BalanceLowNotifyThreshold: settings.BalanceLowNotifyThreshold, - BalanceLowNotifyRechargeURL: settings.BalanceLowNotifyRechargeURL, - SubscriptionExpiryNotifyEnabled: settings.SubscriptionExpiryNotifyEnabled, - AccountQuotaNotifyEnabled: settings.AccountQuotaNotifyEnabled, - AccountQuotaNotifyEmails: dto.NotifyEmailEntriesFromService(settings.AccountQuotaNotifyEmails), - PaymentEnabled: paymentCfg.Enabled, - PaymentMinAmount: paymentCfg.MinAmount, - PaymentMaxAmount: paymentCfg.MaxAmount, - PaymentDailyLimit: paymentCfg.DailyLimit, - PaymentOrderTimeoutMin: paymentCfg.OrderTimeoutMin, - PaymentMaxPendingOrders: paymentCfg.MaxPendingOrders, - PaymentEnabledTypes: paymentCfg.EnabledTypes, - PaymentBalanceDisabled: paymentCfg.BalanceDisabled, - PaymentBalanceRechargeMultiplier: paymentCfg.BalanceRechargeMultiplier, - PaymentRechargeFeeRate: paymentCfg.RechargeFeeRate, - PaymentLoadBalanceStrat: paymentCfg.LoadBalanceStrategy, - PaymentProductNamePrefix: paymentCfg.ProductNamePrefix, - PaymentProductNameSuffix: paymentCfg.ProductNameSuffix, - PaymentHelpImageURL: paymentCfg.HelpImageURL, - PaymentHelpText: paymentCfg.HelpText, - PaymentCancelRateLimitEnabled: paymentCfg.CancelRateLimitEnabled, - PaymentCancelRateLimitMax: paymentCfg.CancelRateLimitMax, - PaymentCancelRateLimitWindow: paymentCfg.CancelRateLimitWindow, - PaymentCancelRateLimitUnit: paymentCfg.CancelRateLimitUnit, - PaymentCancelRateLimitMode: paymentCfg.CancelRateLimitMode, - PaymentAlipayForceQRCode: paymentCfg.AlipayForceQRCode, + RegistrationEnabled: settings.RegistrationEnabled, + EmailVerifyEnabled: settings.EmailVerifyEnabled, + RegistrationEmailSuffixWhitelist: settings.RegistrationEmailSuffixWhitelist, + PromoCodeEnabled: settings.PromoCodeEnabled, + PasswordResetEnabled: settings.PasswordResetEnabled, + FrontendURL: settings.FrontendURL, + InvitationCodeEnabled: settings.InvitationCodeEnabled, + TotpEnabled: settings.TotpEnabled, + TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(), + LoginAgreementEnabled: settings.LoginAgreementEnabled, + LoginAgreementMode: settings.LoginAgreementMode, + LoginAgreementUpdatedAt: settings.LoginAgreementUpdatedAt, + LoginAgreementDocuments: loginAgreementDocumentsToDTO(settings.LoginAgreementDocuments), + SMTPHost: settings.SMTPHost, + SMTPPort: settings.SMTPPort, + SMTPUsername: settings.SMTPUsername, + SMTPPasswordConfigured: settings.SMTPPasswordConfigured, + SMTPFrom: settings.SMTPFrom, + SMTPFromName: settings.SMTPFromName, + SMTPUseTLS: settings.SMTPUseTLS, + TurnstileEnabled: settings.TurnstileEnabled, + TurnstileSiteKey: settings.TurnstileSiteKey, + TurnstileSecretKeyConfigured: settings.TurnstileSecretKeyConfigured, + APIKeyACLTrustForwardedIP: settings.APIKeyACLTrustForwardedIP, + LinuxDoConnectEnabled: settings.LinuxDoConnectEnabled, + LinuxDoConnectClientID: settings.LinuxDoConnectClientID, + LinuxDoConnectClientSecretConfigured: settings.LinuxDoConnectClientSecretConfigured, + LinuxDoConnectRedirectURL: settings.LinuxDoConnectRedirectURL, + DingTalkConnectEnabled: settings.DingTalkConnectEnabled, + DingTalkConnectClientID: settings.DingTalkConnectClientID, + DingTalkConnectClientSecretConfigured: settings.DingTalkConnectClientSecretConfigured, + DingTalkConnectRedirectURL: settings.DingTalkConnectRedirectURL, + DingTalkConnectCorpRestrictionPolicy: settings.DingTalkConnectCorpRestrictionPolicy, + DingTalkConnectInternalCorpID: settings.DingTalkConnectInternalCorpID, + DingTalkConnectBypassRegistration: settings.DingTalkConnectBypassRegistration, + DingTalkConnectSyncCorpEmail: settings.DingTalkConnectSyncCorpEmail, + DingTalkConnectSyncDisplayName: settings.DingTalkConnectSyncDisplayName, + DingTalkConnectSyncDept: settings.DingTalkConnectSyncDept, + DingTalkConnectSyncCorpEmailAttrKey: settings.DingTalkConnectSyncCorpEmailAttrKey, + DingTalkConnectSyncDisplayNameAttrKey: settings.DingTalkConnectSyncDisplayNameAttrKey, + DingTalkConnectSyncDeptAttrKey: settings.DingTalkConnectSyncDeptAttrKey, + DingTalkConnectSyncCorpEmailAttrName: settings.DingTalkConnectSyncCorpEmailAttrName, + DingTalkConnectSyncDisplayNameAttrName: settings.DingTalkConnectSyncDisplayNameAttrName, + DingTalkConnectSyncDeptAttrName: settings.DingTalkConnectSyncDeptAttrName, + WeChatConnectEnabled: settings.WeChatConnectEnabled, + WeChatConnectAppID: settings.WeChatConnectAppID, + WeChatConnectAppSecretConfigured: settings.WeChatConnectAppSecretConfigured, + WeChatConnectOpenAppID: settings.WeChatConnectOpenAppID, + WeChatConnectOpenAppSecretConfigured: settings.WeChatConnectOpenAppSecretConfigured, + WeChatConnectMPAppID: settings.WeChatConnectMPAppID, + WeChatConnectMPAppSecretConfigured: settings.WeChatConnectMPAppSecretConfigured, + WeChatConnectMobileAppID: settings.WeChatConnectMobileAppID, + WeChatConnectMobileAppSecretConfigured: settings.WeChatConnectMobileAppSecretConfigured, + WeChatConnectOpenEnabled: settings.WeChatConnectOpenEnabled, + WeChatConnectMPEnabled: settings.WeChatConnectMPEnabled, + WeChatConnectMobileEnabled: settings.WeChatConnectMobileEnabled, + WeChatConnectMode: settings.WeChatConnectMode, + WeChatConnectScopes: settings.WeChatConnectScopes, + WeChatConnectRedirectURL: settings.WeChatConnectRedirectURL, + WeChatConnectFrontendRedirectURL: settings.WeChatConnectFrontendRedirectURL, + OIDCConnectEnabled: settings.OIDCConnectEnabled, + OIDCConnectProviderName: settings.OIDCConnectProviderName, + OIDCConnectClientID: settings.OIDCConnectClientID, + OIDCConnectClientSecretConfigured: settings.OIDCConnectClientSecretConfigured, + OIDCConnectIssuerURL: settings.OIDCConnectIssuerURL, + OIDCConnectDiscoveryURL: settings.OIDCConnectDiscoveryURL, + OIDCConnectAuthorizeURL: settings.OIDCConnectAuthorizeURL, + OIDCConnectTokenURL: settings.OIDCConnectTokenURL, + OIDCConnectUserInfoURL: settings.OIDCConnectUserInfoURL, + OIDCConnectJWKSURL: settings.OIDCConnectJWKSURL, + OIDCConnectScopes: settings.OIDCConnectScopes, + OIDCConnectRedirectURL: settings.OIDCConnectRedirectURL, + OIDCConnectFrontendRedirectURL: settings.OIDCConnectFrontendRedirectURL, + OIDCConnectTokenAuthMethod: settings.OIDCConnectTokenAuthMethod, + OIDCConnectUsePKCE: settings.OIDCConnectUsePKCE, + OIDCConnectValidateIDToken: settings.OIDCConnectValidateIDToken, + OIDCConnectAllowedSigningAlgs: settings.OIDCConnectAllowedSigningAlgs, + OIDCConnectClockSkewSeconds: settings.OIDCConnectClockSkewSeconds, + OIDCConnectRequireEmailVerified: settings.OIDCConnectRequireEmailVerified, + OIDCConnectUserInfoEmailPath: settings.OIDCConnectUserInfoEmailPath, + OIDCConnectUserInfoIDPath: settings.OIDCConnectUserInfoIDPath, + OIDCConnectUserInfoUsernamePath: settings.OIDCConnectUserInfoUsernamePath, + GitHubOAuthEnabled: settings.GitHubOAuthEnabled, + GitHubOAuthClientID: settings.GitHubOAuthClientID, + GitHubOAuthClientSecretConfigured: settings.GitHubOAuthClientSecretConfigured, + GitHubOAuthRedirectURL: settings.GitHubOAuthRedirectURL, + GitHubOAuthFrontendRedirectURL: settings.GitHubOAuthFrontendRedirectURL, + GoogleOAuthEnabled: settings.GoogleOAuthEnabled, + GoogleOAuthClientID: settings.GoogleOAuthClientID, + GoogleOAuthClientSecretConfigured: settings.GoogleOAuthClientSecretConfigured, + GoogleOAuthRedirectURL: settings.GoogleOAuthRedirectURL, + GoogleOAuthFrontendRedirectURL: settings.GoogleOAuthFrontendRedirectURL, + SiteName: settings.SiteName, + SiteLogo: settings.SiteLogo, + SiteSubtitle: settings.SiteSubtitle, + APIBaseURL: settings.APIBaseURL, + ContactInfo: settings.ContactInfo, + DocURL: settings.DocURL, + HomeContent: settings.HomeContent, + HideCcsImportButton: settings.HideCcsImportButton, + PurchaseSubscriptionEnabled: settings.PurchaseSubscriptionEnabled, + PurchaseSubscriptionURL: settings.PurchaseSubscriptionURL, + TableDefaultPageSize: settings.TableDefaultPageSize, + TablePageSizeOptions: settings.TablePageSizeOptions, + CustomMenuItems: dto.ParseCustomMenuItems(settings.CustomMenuItems), + CustomEndpoints: dto.ParseCustomEndpoints(settings.CustomEndpoints), + DefaultConcurrency: settings.DefaultConcurrency, + DefaultBalance: settings.DefaultBalance, + RiskControlEnabled: settings.RiskControlEnabled, + CyberSessionBlockEnabled: settings.CyberSessionBlockEnabled, + CyberSessionBlockTTLSeconds: settings.CyberSessionBlockTTLSeconds, + AffiliateRebateRate: settings.AffiliateRebateRate, + AffiliateRebateFreezeHours: settings.AffiliateRebateFreezeHours, + AffiliateRebateDurationDays: settings.AffiliateRebateDurationDays, + AffiliateRebatePerInviteeCap: settings.AffiliateRebatePerInviteeCap, + DefaultUserRPMLimit: settings.DefaultUserRPMLimit, + DefaultSubscriptions: defaultSubscriptions, + EnableModelFallback: settings.EnableModelFallback, + FallbackModelAnthropic: settings.FallbackModelAnthropic, + FallbackModelOpenAI: settings.FallbackModelOpenAI, + FallbackModelGemini: settings.FallbackModelGemini, + FallbackModelAntigravity: settings.FallbackModelAntigravity, + EnableIdentityPatch: settings.EnableIdentityPatch, + IdentityPatchPrompt: settings.IdentityPatchPrompt, + OpsMonitoringEnabled: opsEnabled && settings.OpsMonitoringEnabled, + OpsRealtimeMonitoringEnabled: settings.OpsRealtimeMonitoringEnabled, + OpsQueryModeDefault: settings.OpsQueryModeDefault, + OpsMetricsIntervalSeconds: settings.OpsMetricsIntervalSeconds, + MinClaudeCodeVersion: settings.MinClaudeCodeVersion, + MaxClaudeCodeVersion: settings.MaxClaudeCodeVersion, + AllowUngroupedKeyScheduling: settings.AllowUngroupedKeyScheduling, + BackendModeEnabled: settings.BackendModeEnabled, + EnableFingerprintUnification: settings.EnableFingerprintUnification, + EnableMetadataPassthrough: settings.EnableMetadataPassthrough, + EnableCCHSigning: settings.EnableCCHSigning, + EnableClaudeOAuthSystemPromptInjection: settings.EnableClaudeOAuthSystemPromptInjection, + ClaudeOAuthSystemPrompt: settings.ClaudeOAuthSystemPrompt, + ClaudeOAuthSystemPromptBlocks: settings.ClaudeOAuthSystemPromptBlocks, + EnableAnthropicCacheTTL1hInjection: settings.EnableAnthropicCacheTTL1hInjection, + RewriteMessageCacheControl: settings.RewriteMessageCacheControl, + EnableClientDatelineNormalization: settings.EnableClientDatelineNormalization, + AntigravityUserAgentVersion: settings.AntigravityUserAgentVersion, + OpenAICodexUserAgent: settings.OpenAICodexUserAgent, + MinCodexVersion: settings.MinCodexVersion, + MaxCodexVersion: settings.MaxCodexVersion, + CodexCLIOnlyBlacklist: settings.CodexCLIOnlyBlacklist, + CodexCLIOnlyWhitelist: settings.CodexCLIOnlyWhitelist, + CodexCLIOnlyAllowAppServerClients: settings.CodexCLIOnlyAllowAppServerClients, + CodexCLIOnlyEngineFingerprintSignals: settings.CodexCLIOnlyEngineFingerprintSignals, + WebSearchEmulationEnabled: settings.WebSearchEmulationEnabled, + PaymentVisibleMethodAlipaySource: settings.PaymentVisibleMethodAlipaySource, + PaymentVisibleMethodWxpaySource: settings.PaymentVisibleMethodWxpaySource, + PaymentVisibleMethodAlipayEnabled: settings.PaymentVisibleMethodAlipayEnabled, + PaymentVisibleMethodWxpayEnabled: settings.PaymentVisibleMethodWxpayEnabled, + OpenAIAdvancedSchedulerEnabled: settings.OpenAIAdvancedSchedulerEnabled, + OpenAIAdvancedSchedulerStickyWeightedEnabled: settings.OpenAIAdvancedSchedulerStickyWeightedEnabled, + OpenAIAdvancedSchedulerSubscriptionPriorityEnabled: settings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled, + OpenAIAdvancedSchedulerLBTopK: settings.OpenAIAdvancedSchedulerLBTopK, + OpenAIAdvancedSchedulerWeightPriority: settings.OpenAIAdvancedSchedulerWeightPriority, + OpenAIAdvancedSchedulerWeightLoad: settings.OpenAIAdvancedSchedulerWeightLoad, + OpenAIAdvancedSchedulerWeightQueue: settings.OpenAIAdvancedSchedulerWeightQueue, + OpenAIAdvancedSchedulerWeightErrorRate: settings.OpenAIAdvancedSchedulerWeightErrorRate, + OpenAIAdvancedSchedulerWeightTTFT: settings.OpenAIAdvancedSchedulerWeightTTFT, + OpenAIAdvancedSchedulerWeightReset: settings.OpenAIAdvancedSchedulerWeightReset, + OpenAIAdvancedSchedulerWeightQuotaHeadroom: settings.OpenAIAdvancedSchedulerWeightQuotaHeadroom, + OpenAIAdvancedSchedulerWeightPreviousResponse: settings.OpenAIAdvancedSchedulerWeightPreviousResponse, + OpenAIAdvancedSchedulerWeightSessionSticky: settings.OpenAIAdvancedSchedulerWeightSessionSticky, + OpenAIAdvancedSchedulerEffectiveLBTopK: settings.OpenAIAdvancedSchedulerEffectiveLBTopK, + OpenAIAdvancedSchedulerEffectiveWeightPriority: settings.OpenAIAdvancedSchedulerEffectiveWeightPriority, + OpenAIAdvancedSchedulerEffectiveWeightLoad: settings.OpenAIAdvancedSchedulerEffectiveWeightLoad, + OpenAIAdvancedSchedulerEffectiveWeightQueue: settings.OpenAIAdvancedSchedulerEffectiveWeightQueue, + OpenAIAdvancedSchedulerEffectiveWeightErrorRate: settings.OpenAIAdvancedSchedulerEffectiveWeightErrorRate, + OpenAIAdvancedSchedulerEffectiveWeightTTFT: settings.OpenAIAdvancedSchedulerEffectiveWeightTTFT, + OpenAIAdvancedSchedulerEffectiveWeightReset: settings.OpenAIAdvancedSchedulerEffectiveWeightReset, + OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom: settings.OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom, + OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse: settings.OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse, + OpenAIAdvancedSchedulerEffectiveWeightSessionSticky: settings.OpenAIAdvancedSchedulerEffectiveWeightSessionSticky, + BalanceLowNotifyEnabled: settings.BalanceLowNotifyEnabled, + BalanceLowNotifyThreshold: settings.BalanceLowNotifyThreshold, + BalanceLowNotifyRechargeURL: settings.BalanceLowNotifyRechargeURL, + SubscriptionExpiryNotifyEnabled: settings.SubscriptionExpiryNotifyEnabled, + AccountQuotaNotifyEnabled: settings.AccountQuotaNotifyEnabled, + AccountQuotaNotifyEmails: dto.NotifyEmailEntriesFromService(settings.AccountQuotaNotifyEmails), + PaymentEnabled: paymentCfg.Enabled, + PaymentMinAmount: paymentCfg.MinAmount, + PaymentMaxAmount: paymentCfg.MaxAmount, + PaymentDailyLimit: paymentCfg.DailyLimit, + PaymentOrderTimeoutMin: paymentCfg.OrderTimeoutMin, + PaymentMaxPendingOrders: paymentCfg.MaxPendingOrders, + PaymentEnabledTypes: paymentCfg.EnabledTypes, + PaymentBalanceDisabled: paymentCfg.BalanceDisabled, + PaymentBalanceRechargeMultiplier: paymentCfg.BalanceRechargeMultiplier, + PaymentSubscriptionUSDToCNYRate: paymentCfg.SubscriptionUSDToCNYRate, + PaymentRechargeFeeRate: paymentCfg.RechargeFeeRate, + PaymentLoadBalanceStrat: paymentCfg.LoadBalanceStrategy, + PaymentProductNamePrefix: paymentCfg.ProductNamePrefix, + PaymentProductNameSuffix: paymentCfg.ProductNameSuffix, + PaymentHelpImageURL: paymentCfg.HelpImageURL, + PaymentHelpText: paymentCfg.HelpText, + PaymentCancelRateLimitEnabled: paymentCfg.CancelRateLimitEnabled, + PaymentCancelRateLimitMax: paymentCfg.CancelRateLimitMax, + PaymentCancelRateLimitWindow: paymentCfg.CancelRateLimitWindow, + PaymentCancelRateLimitUnit: paymentCfg.CancelRateLimitUnit, + PaymentCancelRateLimitMode: paymentCfg.CancelRateLimitMode, + PaymentAlipayForceQRCode: paymentCfg.AlipayForceQRCode, ChannelMonitorEnabled: settings.ChannelMonitorEnabled, ChannelMonitorDefaultIntervalSeconds: settings.ChannelMonitorDefaultIntervalSeconds, @@ -618,7 +641,19 @@ type UpdateSettingsRequest struct { PaymentVisibleMethodWxpayEnabled *bool `json:"payment_visible_method_wxpay_enabled"` // OpenAI account scheduling - OpenAIAdvancedSchedulerEnabled *bool `json:"openai_advanced_scheduler_enabled"` + 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"` @@ -638,6 +673,7 @@ type UpdateSettingsRequest struct { 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"` @@ -1792,6 +1828,28 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { } 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 @@ -1959,6 +2017,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { EnabledTypes: req.PaymentEnabledTypes, BalanceDisabled: req.PaymentBalanceDisabled, BalanceRechargeMultiplier: req.PaymentBalanceRechargeMultiplier, + SubscriptionUSDToCNYRate: req.PaymentSubscriptionUSDToCNYRate, RechargeFeeRate: req.PaymentRechargeFeeRate, LoadBalanceStrategy: req.PaymentLoadBalanceStrat, ProductNamePrefix: req.PaymentProductNamePrefix, @@ -2014,184 +2073,207 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { } 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, - 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, - 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, + 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, @@ -2238,7 +2320,8 @@ func hasPaymentFields(req UpdateSettingsRequest) bool { req.PaymentMaxAmount != nil || req.PaymentDailyLimit != nil || req.PaymentOrderTimeoutMin != nil || req.PaymentMaxPendingOrders != nil || req.PaymentEnabledTypes != nil || req.PaymentBalanceDisabled != nil || - req.PaymentBalanceRechargeMultiplier != nil || req.PaymentRechargeFeeRate != 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 || @@ -2677,6 +2760,42 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings, 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") @@ -3829,3 +3948,10 @@ func equalPlatformQuotaSettings(before, after map[string]*service.DefaultPlatfor } 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_auth_source_defaults_test.go b/backend/internal/handler/admin/setting_handler_auth_source_defaults_test.go index f953f76760..1626007f19 100644 --- a/backend/internal/handler/admin/setting_handler_auth_source_defaults_test.go +++ b/backend/internal/handler/admin/setting_handler_auth_source_defaults_test.go @@ -217,12 +217,13 @@ func TestSettingHandler_UpdateSettings_PersistsPaymentVisibleMethodsAndAdvancedS handler := NewSettingHandler(svc, nil, nil, nil, nil, nil, nil) body := map[string]any{ - "promo_code_enabled": true, - "payment_visible_method_alipay_source": "easypay", - "payment_visible_method_wxpay_source": "wxpay", - "payment_visible_method_alipay_enabled": true, - "payment_visible_method_wxpay_enabled": false, - "openai_advanced_scheduler_enabled": true, + "promo_code_enabled": true, + "payment_visible_method_alipay_source": "easypay", + "payment_visible_method_wxpay_source": "wxpay", + "payment_visible_method_alipay_enabled": true, + "payment_visible_method_wxpay_enabled": false, + "openai_advanced_scheduler_enabled": true, + "openai_advanced_scheduler_subscription_priority_enabled": true, } rawBody, err := json.Marshal(body) require.NoError(t, err) @@ -240,6 +241,7 @@ func TestSettingHandler_UpdateSettings_PersistsPaymentVisibleMethodsAndAdvancedS require.Equal(t, "true", repo.values[service.SettingPaymentVisibleMethodAlipayEnabled]) require.Equal(t, "false", repo.values[service.SettingPaymentVisibleMethodWxpayEnabled]) require.Equal(t, "true", repo.values["openai_advanced_scheduler_enabled"]) + require.Equal(t, "true", repo.values[service.SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]) var resp response.Response require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) @@ -250,6 +252,7 @@ func TestSettingHandler_UpdateSettings_PersistsPaymentVisibleMethodsAndAdvancedS require.Equal(t, true, data["payment_visible_method_alipay_enabled"]) require.Equal(t, false, data["payment_visible_method_wxpay_enabled"]) require.Equal(t, true, data["openai_advanced_scheduler_enabled"]) + require.Equal(t, true, data["openai_advanced_scheduler_subscription_priority_enabled"]) } func TestSettingHandler_UpdateSettings_PreservesLegacyBlankPaymentVisibleMethodSource(t *testing.T) { diff --git a/backend/internal/handler/batch_image_handler.go b/backend/internal/handler/batch_image_handler.go new file mode 100644 index 0000000000..3a36634618 --- /dev/null +++ b/backend/internal/handler/batch_image_handler.go @@ -0,0 +1,257 @@ +package handler + +import ( + "errors" + "io" + "net/http" + "strconv" + "strings" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + + "github.com/gin-gonic/gin" +) + +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.service.MarkDownloaded(c.Request.Context(), owner, c.Param("id")) +} + +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.service.MarkDownloaded(c.Request.Context(), owner, c.Param("id")) +} + +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 99644ced7f..d63baba91a 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 @@ -11,18 +11,20 @@ import ( func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) { lastUsed := time.Now().UTC().Truncate(time.Second) src := &service.APIKey{ - ID: 1, - UserID: 2, - Key: "sk-map-last-used", - Name: "Mapper", - Status: service.StatusActive, - LastUsedAt: &lastUsed, + ID: 1, + UserID: 2, + Key: "sk-map-last-used", + Name: "Mapper", + Status: service.StatusActive, + LastUsedAt: &lastUsed, + CurrentConcurrency: 3, } out := APIKeyFromService(src) require.NotNil(t, out) require.NotNil(t, out.LastUsedAt) require.WithinDuration(t, lastUsed, *out.LastUsedAt, time.Second) + require.Equal(t, 3, out.CurrentConcurrency) } func TestAPIKeyFromService_MapsNilLastUsedAt(t *testing.T) { diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 5bbab4d45f..03e4c97309 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, @@ -79,31 +80,32 @@ func APIKeyFromService(k *service.APIKey) *APIKey { return nil } out := &APIKey{ - ID: k.ID, - UserID: k.UserID, - Key: k.Key, - Name: k.Name, - GroupID: k.GroupID, - Status: k.Status, - IPWhitelist: k.IPWhitelist, - IPBlacklist: k.IPBlacklist, - LastUsedAt: k.LastUsedAt, - Quota: k.Quota, - QuotaUsed: k.QuotaUsed, - ExpiresAt: k.ExpiresAt, - CreatedAt: k.CreatedAt, - UpdatedAt: k.UpdatedAt, - RateLimit5h: k.RateLimit5h, - RateLimit1d: k.RateLimit1d, - RateLimit7d: k.RateLimit7d, - Usage5h: k.EffectiveUsage5h(), - Usage1d: k.EffectiveUsage1d(), - Usage7d: k.EffectiveUsage7d(), - Window5hStart: k.Window5hStart, - Window1dStart: k.Window1dStart, - Window7dStart: k.Window7dStart, - User: UserFromServiceShallow(k.User), - Group: GroupFromServiceShallow(k.Group), + ID: k.ID, + UserID: k.UserID, + Key: k.Key, + Name: k.Name, + GroupID: k.GroupID, + Status: k.Status, + IPWhitelist: k.IPWhitelist, + IPBlacklist: k.IPBlacklist, + LastUsedAt: k.LastUsedAt, + Quota: k.Quota, + QuotaUsed: k.QuotaUsed, + ExpiresAt: k.ExpiresAt, + CreatedAt: k.CreatedAt, + UpdatedAt: k.UpdatedAt, + CurrentConcurrency: k.CurrentConcurrency, + RateLimit5h: k.RateLimit5h, + RateLimit1d: k.RateLimit1d, + RateLimit7d: k.RateLimit7d, + Usage5h: k.EffectiveUsage5h(), + Usage1d: k.EffectiveUsage1d(), + Usage7d: k.EffectiveUsage7d(), + Window5hStart: k.Window5hStart, + Window1dStart: k.Window1dStart, + Window7dStart: k.Window7dStart, + User: UserFromServiceShallow(k.User), + Group: GroupFromServiceShallow(k.Group), } if k.Window5hStart != nil && !service.IsWindowExpired(k.Window5hStart, service.RateLimitWindow5h) { t := k.Window5hStart.Add(service.RateLimitWindow5h) @@ -179,8 +181,11 @@ 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, PeakRateEnabled: g.PeakRateEnabled, PeakStart: g.PeakStart, PeakEnd: g.PeakEnd, diff --git a/backend/internal/handler/dto/settings.go b/backend/internal/handler/dto/settings.go index 1d6f73bb99..99fba54980 100644 --- a/backend/internal/handler/dto/settings.go +++ b/backend/internal/handler/dto/settings.go @@ -208,7 +208,29 @@ type SystemSettings struct { PaymentVisibleMethodWxpayEnabled bool `json:"payment_visible_method_wxpay_enabled"` // OpenAI account scheduling - OpenAIAdvancedSchedulerEnabled bool `json:"openai_advanced_scheduler_enabled"` + 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"` + OpenAIAdvancedSchedulerEffectiveLBTopK string `json:"openai_advanced_scheduler_effective_lb_top_k"` + OpenAIAdvancedSchedulerEffectiveWeightPriority string `json:"openai_advanced_scheduler_effective_weight_priority"` + OpenAIAdvancedSchedulerEffectiveWeightLoad string `json:"openai_advanced_scheduler_effective_weight_load"` + OpenAIAdvancedSchedulerEffectiveWeightQueue string `json:"openai_advanced_scheduler_effective_weight_queue"` + OpenAIAdvancedSchedulerEffectiveWeightErrorRate string `json:"openai_advanced_scheduler_effective_weight_error_rate"` + OpenAIAdvancedSchedulerEffectiveWeightTTFT string `json:"openai_advanced_scheduler_effective_weight_ttft"` + OpenAIAdvancedSchedulerEffectiveWeightReset string `json:"openai_advanced_scheduler_effective_weight_reset"` + OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom string `json:"openai_advanced_scheduler_effective_weight_quota_headroom"` + OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse string `json:"openai_advanced_scheduler_effective_weight_previous_response"` + OpenAIAdvancedSchedulerEffectiveWeightSessionSticky string `json:"openai_advanced_scheduler_effective_weight_session_sticky"` // Payment configuration PaymentEnabled bool `json:"payment_enabled"` @@ -220,6 +242,7 @@ type SystemSettings struct { 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"` diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index b08dea5680..286d2d5459 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"` @@ -63,6 +64,8 @@ type APIKey struct { ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = never expires) CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` + // CurrentConcurrency is the real-time active request count for this API key. + CurrentConcurrency int `json:"current_concurrency"` // Rate limit fields RateLimit5h float64 `json:"rate_limit_5h"` @@ -97,9 +100,12 @@ 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"` // 高峰时段倍率配置 PeakRateEnabled bool `json:"peak_rate_enabled"` PeakStart string `json:"peak_start"` 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 b65dedf9cc..0caa7f718b 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -158,6 +158,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 } @@ -1006,7 +1007,8 @@ func (h *GatewayHandler) Models(c *gin.Context) { // Get available models from account configurations for the selected group platform. availableModels := h.gatewayService.GetAvailableModels(c.Request.Context(), groupID, platform) if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() { - availableModels = filterModelsByCustomList(availableModels, defaultModelIDsForPlatform(platform), apiKey.Group.ModelsListConfig.Models) + fallbackModels := defaultModelIDsForPlatform(platform) + availableModels = filterModelsByCustomList(customModelsListSource(platform, availableModels, fallbackModels), fallbackModels, apiKey.Group.ModelsListConfig.Models) writeCustomModelsList(c, platform, availableModels) return } @@ -1090,6 +1092,13 @@ func writeOpenAIModelsList(c *gin.Context, modelIDs []string) { }) } +func customModelsListSource(platform string, availableModels, fallbackModels []string) []string { + if platform == service.PlatformAnthropic && len(availableModels) > 0 { + return mergeModelIDs(availableModels, fallbackModels) + } + return availableModels +} + func filterModelsByCustomList(availableModels, fallbackModels, selectedModels []string) []string { if len(selectedModels) == 0 { return availableModels @@ -1158,6 +1167,15 @@ func defaultModelIDsForPlatform(platform string) []string { ids = append(ids, model.ID) } return ids + case service.PlatformAnthropic: + ids := make([]string, 0, len(claude.DefaultModels)+len(antigravity.DefaultModels())) + for _, model := range claude.DefaultModels { + ids = append(ids, model.ID) + } + for _, model := range antigravity.DefaultModels() { + ids = append(ids, model.ID) + } + return mergeModelIDs(ids, nil) case service.PlatformGrok: return xai.DefaultModelIDs() default: @@ -1169,6 +1187,25 @@ func defaultModelIDsForPlatform(platform string) []string { } } +func mergeModelIDs(primary, secondary []string) []string { + seen := make(map[string]struct{}, len(primary)+len(secondary)) + merged := make([]string, 0, len(primary)+len(secondary)) + for _, models := range [][]string{primary, secondary} { + for _, model := range models { + model = strings.TrimSpace(model) + if model == "" { + continue + } + if _, ok := seen[model]; ok { + continue + } + seen[model] = struct{}{} + merged = append(merged, model) + } + } + return merged +} + // AntigravityModels 返回 Antigravity 支持的全部模型 // GET /antigravity/models func (h *GatewayHandler) AntigravityModels(c *gin.Context) { @@ -1760,6 +1797,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..03ceb0d952 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -64,6 +64,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..f5ee18b722 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -64,6 +64,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/gateway_helper.go b/backend/internal/handler/gateway_helper.go index b948ac8fc7..48110da93f 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -10,6 +10,7 @@ import ( "sync" "time" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" @@ -211,6 +212,14 @@ func (h *ConcurrencyHelper) TryAcquireUserSlot(ctx context.Context, userID int64 return result.ReleaseFunc, true, nil } +func (h *ConcurrencyHelper) TryAcquireUserSlotForAPIKey(ctx context.Context, userID int64, maxConcurrency int, apiKeyID int64) (func(), bool, error) { + releaseFunc, acquired, err := h.TryAcquireUserSlot(ctx, userID, maxConcurrency) + if err != nil || !acquired { + return releaseFunc, acquired, err + } + return h.withAPIKeySlot(ctx, apiKeyID, releaseFunc), true, nil +} + // TryAcquireAccountSlot 尝试立即获取账号并发槽位。 // 返回值: (releaseFunc, acquired, error) func (h *ConcurrencyHelper) TryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (func(), bool, error) { @@ -241,7 +250,7 @@ func (h *ConcurrencyHelper) acquireUserSlotWithWaitTimeout(c *gin.Context, userI } if acquired { - return releaseFunc, nil + return h.withAPIKeySlotFromGin(c, releaseFunc), nil } queueLimit := service.CalculateMaxWait(maxConcurrency) - maxConcurrency @@ -258,7 +267,37 @@ func (h *ConcurrencyHelper) acquireUserSlotWithWaitTimeout(c *gin.Context, userI defer h.DecrementWaitCount(ctx, userID) // Need to wait - handle streaming ping if needed - return h.waitForSlotWithPingTimeout(c, "user", userID, maxConcurrency, timeout, isStream, streamStarted, false) + releaseFunc, err = h.waitForSlotWithPingTimeout(c, "user", userID, maxConcurrency, timeout, isStream, streamStarted, false) + if err != nil { + return nil, err + } + return h.withAPIKeySlotFromGin(c, releaseFunc), nil +} + +func (h *ConcurrencyHelper) withAPIKeySlotFromGin(c *gin.Context, releaseFunc func()) func() { + if c == nil { + return releaseFunc + } + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok || apiKey == nil { + return releaseFunc + } + return h.withAPIKeySlot(c.Request.Context(), apiKey.ID, releaseFunc) +} + +func (h *ConcurrencyHelper) withAPIKeySlot(ctx context.Context, apiKeyID int64, releaseFunc func()) func() { + if h == nil || h.concurrencyService == nil || apiKeyID <= 0 { + return releaseFunc + } + apiKeyReleaseFunc := h.concurrencyService.TrackAPIKeySlot(ctx, apiKeyID) + return func() { + if releaseFunc != nil { + releaseFunc() + } + if apiKeyReleaseFunc != nil { + apiKeyReleaseFunc() + } + } } // AcquireAccountSlotWithWait acquires an account concurrency slot, waiting if necessary. diff --git a/backend/internal/handler/gateway_helper_hotpath_test.go b/backend/internal/handler/gateway_helper_hotpath_test.go index fb17481f1c..5e0697f083 100644 --- a/backend/internal/handler/gateway_helper_hotpath_test.go +++ b/backend/internal/handler/gateway_helper_hotpath_test.go @@ -9,6 +9,7 @@ import ( "testing" "time" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -29,6 +30,9 @@ type helperConcurrencyCacheStub struct { waitDecrementCalls int waitMaxWait int waitIncrementHook func() + apiKeyTrackCalls int + apiKeyReleaseCalls int + apiKeyTrackIDs []int64 } func (s *helperConcurrencyCacheStub) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) { @@ -97,6 +101,29 @@ func (s *helperConcurrencyCacheStub) GetUserConcurrency(ctx context.Context, use return 0, nil } +func (s *helperConcurrencyCacheStub) TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.apiKeyTrackCalls++ + s.apiKeyTrackIDs = append(s.apiKeyTrackIDs, apiKeyID) + return nil +} + +func (s *helperConcurrencyCacheStub) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.apiKeyReleaseCalls++ + return nil +} + +func (s *helperConcurrencyCacheStub) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) { + out := make(map[int64]int, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + out[apiKeyID] = 0 + } + return out, nil +} + func (s *helperConcurrencyCacheStub) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) { s.mu.Lock() s.waitIncrementCalls++ @@ -270,6 +297,48 @@ func TestAcquireUserSlotWithWait_ImmediateAcquireSkipsWaitQueue(t *testing.T) { require.Equal(t, 1, cache.userReleaseCalls) } +func TestAcquireUserSlotWithWait_TracksAPIKeySlot(t *testing.T) { + cache := &helperConcurrencyCacheStub{ + userSeq: []bool{true}, + } + concurrency := service.NewConcurrencyService(cache) + helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond) + c, _ := newHelperTestContext(http.MethodPost, "/v1/messages") + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 77}) + streamStarted := false + + release, err := helper.acquireUserSlotWithWaitTimeout(c, 202, 3, time.Second, false, &streamStarted) + require.NoError(t, err) + require.NotNil(t, release) + require.Equal(t, 1, cache.apiKeyTrackCalls) + require.Equal(t, []int64{77}, cache.apiKeyTrackIDs) + + release() + + require.Equal(t, 1, cache.userReleaseCalls) + require.Equal(t, 1, cache.apiKeyReleaseCalls) +} + +func TestTryAcquireUserSlotForAPIKey_TracksAPIKeySlot(t *testing.T) { + cache := &helperConcurrencyCacheStub{ + userSeq: []bool{true}, + } + concurrency := service.NewConcurrencyService(cache) + helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond) + + release, acquired, err := helper.TryAcquireUserSlotForAPIKey(context.Background(), 202, 3, 77) + require.NoError(t, err) + require.True(t, acquired) + require.NotNil(t, release) + require.Equal(t, 1, cache.apiKeyTrackCalls) + require.Equal(t, []int64{77}, cache.apiKeyTrackIDs) + + release() + + require.Equal(t, 1, cache.userReleaseCalls) + require.Equal(t, 1, cache.apiKeyReleaseCalls) +} + func TestAcquireUserSlotWithWait_WaitSuccessDecrementsBeforeReturn(t *testing.T) { cache := &helperConcurrencyCacheStub{ userSeq: []bool{false, true}, diff --git a/backend/internal/handler/gateway_models_test.go b/backend/internal/handler/gateway_models_test.go index c5238f2a2f..6011e13027 100644 --- a/backend/internal/handler/gateway_models_test.go +++ b/backend/internal/handler/gateway_models_test.go @@ -269,6 +269,149 @@ func TestGatewayModels_CustomModelsListKeepsConcreteModelAllowedByWildcardMappin require.Equal(t, []string{"claude-sonnet-4-6"}, modelIDsForTest(got.Data)) } +func TestGatewayModels_AnthropicCustomModelsListIncludesOAuthClaudeAndMappedDeepSeek(t *testing.T) { + gin.SetMode(gin.TestMode) + + groupID := int64(28) + h := newGatewayModelsHandlerForTest( + &gatewayModelsAccountRepoStub{ + byGroup: map[int64][]service.Account{ + groupID: { + { + ID: 1, + Platform: service.PlatformAnthropic, + Type: service.AccountTypeOAuth, + }, + { + ID: 2, + Platform: service.PlatformAnthropic, + Type: service.AccountTypeAPIKey, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "deepseek-v4-pro": "deepseek-v4-pro", + }, + }, + }, + }, + }, + }, + ) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ + Group: &service.Group{ + ID: groupID, + Platform: service.PlatformAnthropic, + ModelsListConfig: service.GroupModelsListConfig{ + Enabled: true, + Models: []string{"claude-fable-5", "claude-opus-4-8", "deepseek-v4-pro"}, + }, + }, + }) + + h.Models(c) + + require.Equal(t, http.StatusOK, rec.Code) + + var got gatewayModelsResponseForTest + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Equal(t, []string{"claude-fable-5", "claude-opus-4-8", "deepseek-v4-pro"}, modelIDsForTest(got.Data)) +} + +func TestGatewayModels_AnthropicCustomModelsListDisabledKeepsMappedModelList(t *testing.T) { + gin.SetMode(gin.TestMode) + + groupID := int64(29) + h := newGatewayModelsHandlerForTest( + &gatewayModelsAccountRepoStub{ + byGroup: map[int64][]service.Account{ + groupID: { + { + ID: 1, + Platform: service.PlatformAnthropic, + Type: service.AccountTypeOAuth, + }, + { + ID: 2, + Platform: service.PlatformAnthropic, + Type: service.AccountTypeAPIKey, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "deepseek-v4-pro": "deepseek-v4-pro", + }, + }, + }, + }, + }, + }, + ) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ + Group: &service.Group{ + ID: groupID, + Platform: service.PlatformAnthropic, + ModelsListConfig: service.GroupModelsListConfig{ + Enabled: false, + Models: []string{"claude-fable-5", "deepseek-v4-pro"}, + }, + }, + }) + + h.Models(c) + + require.Equal(t, http.StatusOK, rec.Code) + + var got gatewayModelsResponseForTest + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Equal(t, []string{"deepseek-v4-pro"}, modelIDsForTest(got.Data)) +} + +func TestGatewayModels_AnthropicCustomModelsListIncludesOAuthClaudeWithoutMappings(t *testing.T) { + gin.SetMode(gin.TestMode) + + groupID := int64(30) + h := newGatewayModelsHandlerForTest( + &gatewayModelsAccountRepoStub{ + byGroup: map[int64][]service.Account{ + groupID: { + { + ID: 1, + Platform: service.PlatformAnthropic, + Type: service.AccountTypeOAuth, + }, + }, + }, + }, + ) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ + Group: &service.Group{ + ID: groupID, + Platform: service.PlatformAnthropic, + ModelsListConfig: service.GroupModelsListConfig{ + Enabled: true, + Models: []string{"claude-opus-4-6-thinking", "claude-sonnet-4-5"}, + }, + }, + }) + + h.Models(c) + + require.Equal(t, http.StatusOK, rec.Code) + + var got gatewayModelsResponseForTest + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Equal(t, []string{"claude-opus-4-6-thinking", "claude-sonnet-4-5"}, modelIDsForTest(got.Data)) +} + func TestGatewayModels_CustomModelsListCanReturnEmptyWhenSelectionsUnavailable(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 8e236ea49f..4fd1411b23 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -174,6 +174,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. service.OpenAIUpstreamTransportHTTPSSE, "", false, + false, service.PlatformGrok, ) if err != nil { 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 ca43a2cff3..847d386cde 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -64,6 +64,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 } @@ -145,6 +146,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { service.OpenAIUpstreamTransportAny, service.OpenAIEndpointCapabilityChatCompletions, false, + false, requestPlatform, ) if err != nil { diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index a80c7f7d96..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 } @@ -117,6 +118,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { service.OpenAIUpstreamTransportHTTPSSE, service.OpenAIEndpointCapabilityEmbeddings, false, + false, ) if err != nil { reqLog.Warn("openai_embeddings.account_select_failed", diff --git a/backend/internal/handler/openai_gateway_count_tokens.go b/backend/internal/handler/openai_gateway_count_tokens.go index ec530e8ab6..9a6709cc4f 100644 --- a/backend/internal/handler/openai_gateway_count_tokens.go +++ b/backend/internal/handler/openai_gateway_count_tokens.go @@ -64,6 +64,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 } @@ -110,6 +111,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) { service.OpenAIUpstreamTransportAny, service.OpenAIEndpointCapabilityChatCompletions, false, + false, openAICompatibleRequestPlatform(apiKey), ) service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 8eec667fa6..551de2dd61 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -218,6 +218,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { // 校验请求体 JSON 合法性 if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } @@ -350,6 +351,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { service.OpenAIUpstreamTransportAny, service.OpenAIEndpointCapabilityChatCompletions, requireCompact, + false, requestPlatform, ) if err != nil { @@ -696,6 +698,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 } @@ -783,6 +786,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { service.OpenAIUpstreamTransportAny, service.OpenAIEndpointCapabilityChatCompletions, false, + false, requestPlatform, ) if err != nil { @@ -1266,6 +1270,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "previous_response_id must be a response.id (resp_*), not a message id") return } + firstMessageToolCoverage := service.AnalyzeToolCallOutputContextCoverageBytes(firstMessage) + previousResponseCanMove := !firstMessageToolCoverage.HasFunctionCallOutput || firstMessageToolCoverage.ContextCoversAllCallIDs reqLog = reqLog.With( zap.Bool("ws_ingress", true), zap.String("model", reqModel), @@ -1318,7 +1324,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // 必须尽早注册，确保任何 early return 都能释放已获取的并发槽位。 defer releaseTurnSlots() - userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, subject.UserID, subject.Concurrency) + userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, apiKey.ID) if err != nil { reqLog.Warn("openai.websocket_user_slot_acquire_failed", zap.Error(err)) closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to acquire user concurrency slot") @@ -1333,7 +1339,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if currentUserRelease != nil { return true } - userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, subject.UserID, subject.Concurrency) + userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, apiKey.ID) if err != nil { reqLog.Warn("openai.websocket_user_slot_reacquire_failed", zap.Error(err)) closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to acquire user concurrency slot") @@ -1381,6 +1387,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { requiredTransport, service.OpenAIEndpointCapabilityChatCompletions, false, + previousResponseCanMove, requestPlatform, ) if err != nil { @@ -1484,7 +1491,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // 防御式清理：避免异常路径下旧槽位覆盖导致泄漏。 releaseTurnSlots() // 非首轮 turn 需要重新抢占并发槽位，避免长连接空闲占槽。 - userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, subject.UserID, subject.Concurrency) + userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, apiKey.ID) if err != nil { return service.NewOpenAIWSClientCloseError(coderws.StatusInternalError, "failed to acquire user concurrency slot", err) } @@ -1581,8 +1588,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // 说明该会话链不属于本次调度到的账号，原样转发会触发上游会话链鉴权失败（“鉴权失败，请检查 API Key”）。 // 故剥离首包里的 previous_response_id，改用首包内 input 重建上下文；带 function_call_output 的 // 工具续链无法重建，保持原样。仅作用于首轮首包，后续 turn 的续链由 WS 转发层既有逻辑处理。 - if previousResponseID != "" && !scheduleDecision.StickyPreviousHit && - !service.ValidateFunctionCallOutputContextBytes(wsFirstMessage).HasFunctionCallOutput { + if previousResponseID != "" && !scheduleDecision.StickyPreviousHit && previousResponseCanMove { wsFirstMessage = service.RemovePreviousResponseIDFromBody(wsFirstMessage) reqLog.Debug("openai.websocket_previous_response_id_stripped_cross_group", zap.Int64("account_id", account.ID), diff --git a/backend/internal/handler/payment_handler.go b/backend/internal/handler/payment_handler.go index 7cdf73cd3d..a267d73724 100644 --- a/backend/internal/handler/payment_handler.go +++ b/backend/internal/handler/payment_handler.go @@ -150,6 +150,7 @@ func (h *PaymentHandler) GetCheckoutInfo(c *gin.Context) { Plans: planList, BalanceDisabled: cfg.BalanceDisabled, BalanceRechargeMultiplier: cfg.BalanceRechargeMultiplier, + SubscriptionUSDToCNYRate: cfg.SubscriptionUSDToCNYRate, RechargeFeeRate: cfg.RechargeFeeRate, HelpText: cfg.HelpText, HelpImageURL: cfg.HelpImageURL, @@ -165,6 +166,7 @@ type checkoutInfoResponse struct { Plans []checkoutPlan `json:"plans"` BalanceDisabled bool `json:"balance_disabled"` BalanceRechargeMultiplier float64 `json:"balance_recharge_multiplier"` + SubscriptionUSDToCNYRate float64 `json:"subscription_usd_to_cny_rate"` RechargeFeeRate float64 `json:"recharge_fee_rate"` HelpText string `json:"help_text"` HelpImageURL string `json:"help_image_url"` 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/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/payment/provider/easypay.go b/backend/internal/payment/provider/easypay.go index 32d6b7bebf..f1c17427ad 100644 --- a/backend/internal/payment/provider/easypay.go +++ b/backend/internal/payment/provider/easypay.go @@ -39,6 +39,12 @@ type EasyPay struct { httpClient *http.Client } +type easyPayCustomMethod struct { + Type string `json:"type"` + UpstreamType string `json:"upstreamType"` + DisplayName string `json:"displayName"` +} + // NewEasyPay creates a new EasyPay provider. // config keys: pid, pkey, apiBase, notifyUrl, returnUrl, cid, cidAlipay, cidWxpay func NewEasyPay(instanceID string, config map[string]string) (*EasyPay, error) { @@ -95,7 +101,13 @@ func (e *EasyPay) apiBase() string { func (e *EasyPay) Name() string { return "EasyPay" } func (e *EasyPay) ProviderKey() string { return payment.TypeEasyPay } func (e *EasyPay) SupportedTypes() []payment.PaymentType { - return []payment.PaymentType{payment.TypeAlipay, payment.TypeWxpay} + types := []payment.PaymentType{payment.TypeAlipay, payment.TypeWxpay} + for _, method := range e.customMethods() { + if method.Type != "" { + types = append(types, method.Type) + } + } + return types } func (e *EasyPay) MerchantIdentityMetadata() map[string]string { @@ -124,13 +136,14 @@ func (e *EasyPay) CreatePayment(ctx context.Context, req payment.CreatePaymentRe // TradeNo is empty; it arrives via the notify callback after payment. func (e *EasyPay) createRedirectPayment(req payment.CreatePaymentRequest) (*payment.CreatePaymentResponse, error) { notifyURL, returnURL := e.resolveURLs(req) + paymentType := e.upstreamPaymentType(req.PaymentType) params := map[string]string{ - "pid": e.config["pid"], "type": req.PaymentType, + "pid": e.config["pid"], "type": paymentType, "out_trade_no": req.OrderID, "notify_url": notifyURL, "return_url": returnURL, "name": req.Subject, "money": req.Amount, } - if cid := e.resolveCID(req.PaymentType); cid != "" { + if cid := e.resolveCID(paymentType); cid != "" { params["cid"] = cid } if req.IsMobile { @@ -150,13 +163,14 @@ func (e *EasyPay) createRedirectPayment(req payment.CreatePaymentRequest) (*paym // createAPIPayment calls mapi.php to get payurl/qrcode (existing behavior). func (e *EasyPay) createAPIPayment(ctx context.Context, req payment.CreatePaymentRequest) (*payment.CreatePaymentResponse, error) { notifyURL, returnURL := e.resolveURLs(req) + paymentType := e.upstreamPaymentType(req.PaymentType) params := map[string]string{ - "pid": e.config["pid"], "type": req.PaymentType, + "pid": e.config["pid"], "type": paymentType, "out_trade_no": req.OrderID, "notify_url": notifyURL, "return_url": returnURL, "name": req.Subject, "money": req.Amount, "clientip": req.ClientIP, } - if cid := e.resolveCID(req.PaymentType); cid != "" { + if cid := e.resolveCID(paymentType); cid != "" { params["cid"] = cid } if req.IsMobile { @@ -204,6 +218,41 @@ func (e *EasyPay) resolveURLs(req payment.CreatePaymentRequest) (string, string) return notifyURL, returnURL } +func (e *EasyPay) customMethods() []easyPayCustomMethod { + if e == nil { + return nil + } + raw := strings.TrimSpace(e.config["customMethods"]) + if raw == "" { + return nil + } + var methods []easyPayCustomMethod + if err := json.Unmarshal([]byte(raw), &methods); err != nil { + return nil + } + result := make([]easyPayCustomMethod, 0, len(methods)) + for _, method := range methods { + method.Type = strings.TrimSpace(method.Type) + method.UpstreamType = strings.TrimSpace(method.UpstreamType) + method.DisplayName = strings.TrimSpace(method.DisplayName) + if method.Type == "" || method.UpstreamType == "" { + continue + } + result = append(result, method) + } + return result +} + +func (e *EasyPay) upstreamPaymentType(paymentType string) string { + paymentType = strings.TrimSpace(paymentType) + for _, method := range e.customMethods() { + if paymentType == method.Type { + return method.UpstreamType + } + } + return paymentType +} + func (e *EasyPay) QueryOrder(ctx context.Context, tradeNo string) (*payment.QueryOrderResponse, error) { params := map[string]string{ "act": "order", "pid": e.config["pid"], diff --git a/backend/internal/payment/provider/easypay_refund_test.go b/backend/internal/payment/provider/easypay_refund_test.go index 9e0e4942c2..3b76329870 100644 --- a/backend/internal/payment/provider/easypay_refund_test.go +++ b/backend/internal/payment/provider/easypay_refund_test.go @@ -179,6 +179,102 @@ func TestEasyPayRefundResponseErrors(t *testing.T) { } } +func TestEasyPayCustomMethodsUseConfiguredUpstreamType(t *testing.T) { + t.Parallel() + + provider, err := NewEasyPay("test-instance", map[string]string{ + "pid": "pid-1", + "pkey": "pkey-1", + "apiBase": "https://pay.example.com", + "notifyUrl": "https://example.com/notify", + "returnUrl": "https://example.com/return", + "paymentMode": paymentModePopup, + "customMethods": `[{"type":"ldc","upstreamType":"epay","displayName":"LDC"},{"type":"usdt_trc20","upstreamType":"usdt","displayName":"USDT-TRC20"}]`, + }) + if err != nil { + t.Fatalf("NewEasyPay: %v", err) + } + + resp, err := provider.CreatePayment(context.Background(), payment.CreatePaymentRequest{ + OrderID: "sub2-custom-1", + Amount: "1.00", + PaymentType: "usdt_trc20", + Subject: "Custom EasyPay", + }) + if err != nil { + t.Fatalf("CreatePayment: %v", err) + } + payURL, err := url.Parse(resp.PayURL) + if err != nil { + t.Fatalf("parse pay url: %v", err) + } + if got := payURL.Query().Get("type"); got != "usdt" { + t.Fatalf("pay url type = %q, want usdt (%s)", got, resp.PayURL) + } +} + +func TestEasyPayCustomMethodsResolveCIDFromConfiguredUpstreamType(t *testing.T) { + t.Parallel() + + provider, err := NewEasyPay("test-instance", map[string]string{ + "pid": "pid-1", + "pkey": "pkey-1", + "apiBase": "https://pay.example.com", + "notifyUrl": "https://example.com/notify", + "returnUrl": "https://example.com/return", + "paymentMode": paymentModePopup, + "cidAlipay": "cid-alipay", + "cidWxpay": "cid-wxpay", + "customMethods": `[{"type":"ldc","upstreamType":"alipay","displayName":"LDC"}]`, + }) + if err != nil { + t.Fatalf("NewEasyPay: %v", err) + } + + resp, err := provider.CreatePayment(context.Background(), payment.CreatePaymentRequest{ + OrderID: "sub2-custom-cid", + Amount: "1.00", + PaymentType: "ldc", + Subject: "Custom EasyPay CID", + }) + if err != nil { + t.Fatalf("CreatePayment: %v", err) + } + payURL, err := url.Parse(resp.PayURL) + if err != nil { + t.Fatalf("parse pay url: %v", err) + } + if got := payURL.Query().Get("type"); got != "alipay" { + t.Fatalf("pay url type = %q, want alipay (%s)", got, resp.PayURL) + } + if got := payURL.Query().Get("cid"); got != "cid-alipay" { + t.Fatalf("pay url cid = %q, want cid-alipay (%s)", got, resp.PayURL) + } +} + +func TestEasyPaySupportedTypesIncludeCustomMethods(t *testing.T) { + t.Parallel() + + provider, err := NewEasyPay("test-instance", map[string]string{ + "pid": "pid-1", + "pkey": "pkey-1", + "apiBase": "https://pay.example.com", + "notifyUrl": "https://example.com/notify", + "returnUrl": "https://example.com/return", + "customMethods": `[{"type":"ldc","upstreamType":"epay","displayName":"LDC"},{"type":"usdt_trc20","upstreamType":"usdt","displayName":"USDT-TRC20"}]`, + }) + if err != nil { + t.Fatalf("NewEasyPay: %v", err) + } + + got := strings.Join(provider.SupportedTypes(), ",") + for _, want := range []string{"alipay", "wxpay", "ldc", "usdt_trc20"} { + if !strings.Contains(got, want) { + t.Fatalf("SupportedTypes() = %q, want it to include %q", got, want) + } + } +} + func newTestEasyPay(t *testing.T, apiBase string) *EasyPay { t.Helper() diff --git a/backend/internal/pkg/openai/constants.go b/backend/internal/pkg/openai/constants.go index f658cf0675..c9d391df4e 100644 --- a/backend/internal/pkg/openai/constants.go +++ b/backend/internal/pkg/openai/constants.go @@ -18,6 +18,9 @@ type Model struct { // DefaultModels OpenAI models list var DefaultModels = []Model{ + {ID: "gpt-5.6-sol", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Sol"}, + {ID: "gpt-5.6-terra", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Terra"}, + {ID: "gpt-5.6-luna", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Luna"}, {ID: "gpt-5.5", Object: "model", Created: 1776873600, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.5"}, {ID: "gpt-5.4", Object: "model", Created: 1738368000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.4"}, {ID: "gpt-5.4-mini", Object: "model", Created: 1738368000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.4 Mini"}, diff --git a/backend/internal/pkg/xai/models.go b/backend/internal/pkg/xai/models.go index 4902fcb94f..a5b800cf2c 100644 --- a/backend/internal/pkg/xai/models.go +++ b/backend/internal/pkg/xai/models.go @@ -12,6 +12,7 @@ type Model struct { var defaultModels = []Model{ {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"}, @@ -46,6 +47,7 @@ func DefaultModelMapping() map[string]string { mapping["grok"] = "grok-4.3" mapping["grok-latest"] = "grok-4.3" mapping["grok-build"] = "grok-build-0.1" + mapping["grok-composer"] = "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..28609a08fa 100644 --- a/backend/internal/pkg/xai/oauth_test.go +++ b/backend/internal/pkg/xai/oauth_test.go @@ -210,6 +210,7 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) { require.Equal(t, "grok-4.3", mapping["grok"]) require.Equal(t, "grok-4.3", mapping["grok-latest"]) require.Equal(t, "grok-build-0.1", mapping["grok-build"]) + require.Equal(t, "grok-composer-2.5-fast", mapping["grok-composer"]) 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/account_repo.go b/backend/internal/repository/account_repo.go index 1a21b5893c..a8015ddba0 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -482,7 +482,7 @@ func (r *accountRepository) List(ctx context.Context, params pagination.Paginati return r.ListWithFilters(ctx, params, "", "", "", "", 0, "") } -func (r *accountRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, *pagination.PaginationResult, error) { +func (r *accountRepository) accountListFilteredQuery(platform, accountType, status, search string, groupID int64, privacyMode string) *dbent.AccountQuery { q := r.client.Account.Query() if platform != "" { @@ -575,6 +575,11 @@ func (r *accountRepository) ListWithFilters(ctx context.Context, params paginati })) } + return q +} + +func (r *accountRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, *pagination.PaginationResult, error) { + q := r.accountListFilteredQuery(platform, accountType, status, search, groupID, privacyMode) // Clone before Count so interceptor-appended predicates (SoftDeleteMixin's // deleted_at IS NULL) don't accumulate on the shared builder and pollute the // subsequent list query. Same pattern used in group_repo/promo_code_repo/user_repo @@ -603,6 +608,14 @@ func (r *accountRepository) ListWithFilters(ctx context.Context, params paginati return outAccounts, paginationResultFromTotal(int64(total), params), nil } +func (r *accountRepository) ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, error) { + accounts, err := r.accountListFilteredQuery(platform, accountType, status, search, groupID, privacyMode).All(ctx) + if err != nil { + return nil, err + } + return r.accountsToService(ctx, accounts) +} + func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platformFilter string, groupIDFilter *int64) ([]service.Account, error) { if r == nil || r.client == nil { return []service.Account{}, nil diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index 76cad809c3..877fc90353 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -177,6 +177,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldAllowImageGeneration, + group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldImageRateMultiplier, group.FieldImagePrice1k, @@ -755,6 +756,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 +799,14 @@ 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, DefaultValidityDays: g.DefaultValidityDays, ClaudeCodeOnly: g.ClaudeCodeOnly, FallbackGroupID: g.FallbackGroupID, 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..16de6652e3 --- /dev/null +++ b/backend/internal/repository/batch_image_queue.go @@ -0,0 +1,280 @@ +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 +) + +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 +`) + +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 + } + + ok, err := q.rdb.SetNX(ctx, q.inflightKey(batchID), batchID, q.inflightTTL).Result() + if err != nil { + return err + } + if !ok { + return service.ErrBatchImageAlreadyQueued + } + if err := q.rdb.LPush(ctx, q.readyKey, batchID).Err(); err != nil { + _ = q.rdb.Del(ctx, q.inflightKey(batchID)).Err() + return err + } + return nil +} + +func (q *batchImageQueue) Reserve(ctx context.Context, blockTimeout time.Duration) (service.ReservedBatchImageJob, error) { + result, err := q.rdb.BRPop(ctx, blockTimeout, q.readyKey).Result() + if errors.Is(err, redis.Nil) { + return service.ReservedBatchImageJob{}, service.ErrBatchImageQueueEmpty + } + if err != nil { + return service.ReservedBatchImageJob{}, err + } + if len(result) != 2 || !service.IsValidBatchImageID(result[1]) { + return service.ReservedBatchImageJob{}, service.ErrInvalidBatchImageQueuePayload + } + + batchID := result[1] + if err := q.rdb.ZAdd(ctx, q.activeKey, redis.Z{ + Score: float64(time.Now().UnixMilli()), + Member: batchID, + }).Err(); err != nil { + return service.ReservedBatchImageJob{}, err + } + return service.ReservedBatchImageJob{BatchID: 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 + } + return q.rdb.ZAdd(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() +} + +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..5188e67d07 --- /dev/null +++ b/backend/internal/repository/batch_image_queue_test.go @@ -0,0 +1,123 @@ +//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 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..cd4bf70339 --- /dev/null +++ b/backend/internal/repository/batch_image_repo.go @@ -0,0 +1,941 @@ +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) 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 + if err := sqlq.QueryRowContext(ctx, `SELECT id FROM batch_image_jobs WHERE batch_id = $1 FOR UPDATE`, batchID).Scan(&id); err != nil { + return translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil) + } + 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..174e0657e3 --- /dev/null +++ b/backend/internal/repository/batch_image_repo_integration_test.go @@ -0,0 +1,371 @@ +//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) + + 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 b530c30b90..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" ) @@ -27,6 +28,8 @@ const ( accountSlotKeyPrefix = "concurrency:account:" // 格式: concurrency:user:{userID} userSlotKeyPrefix = "concurrency:user:" + // 格式: concurrency:api_key:{apiKeyID} + apiKeySlotKeyPrefix = "concurrency:api_key:" // 等待队列计数器格式: concurrency:wait:{userID} waitQueueKeyPrefix = "concurrency:wait:" // 账号级等待队列计数器格式: wait:account:{accountID} @@ -34,6 +37,19 @@ const ( // 默认槽位过期时间（分钟），可通过配置覆盖 defaultSlotTTLMinutes = 15 + + // 活跃索引用来替代后台任务全量 SCAN 槽位键。 + // member 是账号/用户 ID，score 是“预计仍需关注到”的 Redis Unix 秒时间戳。 + accountActiveIndexKey = "concurrency:account:active_index" // ZSET member=accountID, score=expireAtUnixSeconds + userActiveIndexKey = "concurrency:user:active_index" // ZSET member=userID, score=expireAtUnixSeconds + + // 后台清理只按批处理索引候选，避免单次任务占用 Redis 太久。 + activeIndexCleanupBatchSize = 1000 + activeIndexPipelineChunkSize = 500 + + // 一次性迁移 marker：活跃索引机制上线前遗留的等待计数键无法被索引发现， + // 且有流量时 TTL 会被不断刷新，必须清扫一次。marker 存在即代表已完成。 + legacyWaitSweepMarkerKey = "concurrency:startup:legacy_wait_sweep:v1" ) var ( @@ -43,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). @@ -65,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 -- 检查是否达到并发上限 @@ -73,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 统计有序集合中的槽位数量并清理过期条目 @@ -99,50 +116,82 @@ var ( return redis.call('ZCARD', key) `) + // trackSlotScript 记录 stats-only 槽位，不做并发上限判断。 + // KEYS[1] = 有序集合键 + // ARGV[1] = TTL（秒） + // ARGV[2] = requestID + trackSlotScript = 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 key = KEYS[1] + local ttl = tonumber(ARGV[1]) + local requestID = ARGV[2] + + local timeResult = redis.call('TIME') + local now = tonumber(timeResult[1]) + local expireBefore = now - ttl + + redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore) + redis.call('ZADD', key, now, requestID) + redis.call('EXPIRE', key, ttl) + return 1 + `) + // incrementWaitScript - refreshes TTL on each increment to keep queue depth accurate // 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(` @@ -174,51 +223,27 @@ var ( return 1 `) - // cleanupExpiredSlotKeysScript 批量清理实际存在的账号槽位键，避免后台任务从数据库加载全量账号。 - // KEYS = 有序集合键列表，ARGV[1] = TTL（秒）。 - cleanupExpiredSlotKeysScript = 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 ttl = tonumber(ARGV[1]) - local timeResult = redis.call('TIME') - local now = tonumber(timeResult[1]) - local expireBefore = now - ttl - local removed = 0 - for i = 1, #KEYS do - local key = KEYS[i] - removed = removed + redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore) - if redis.call('ZCARD', key) == 0 then - redis.call('DEL', key) - else - redis.call('EXPIRE', key, ttl) - end - end - return removed - `) - - // startupCleanupScript 清理非当前进程前缀的槽位成员。 - // KEYS 是有序集合键列表，ARGV[1] 是当前进程前缀，ARGV[2] 是槽位 TTL。 - // 遍历每个 KEYS[i]，移除前缀不匹配的成员，清空后删 key，否则刷新 EXPIRE。 - startupCleanupScript = redis.NewScript(` + // 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] local slotTTL = tonumber(ARGV[2]) local removed = 0 - for i = 1, #KEYS do - local key = KEYS[i] - 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) + 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 - return removed + local remaining = redis.call('ZCARD', key) + if remaining == 0 then + redis.call('DEL', key) + else + redis.call('EXPIRE', key, slotTTL) + end + return {removed, remaining} `) ) @@ -254,6 +279,10 @@ func userSlotKey(userID int64) string { return fmt.Sprintf("%s%d", userSlotKeyPrefix, userID) } +func apiKeySlotKey(apiKeyID int64) string { + return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID) +} + func waitQueueKey(userID int64) string { return fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) } @@ -262,21 +291,242 @@ func accountWaitKey(accountID int64) string { return fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID) } +// redisUnixSeconds 统一使用 Redis 服务器时间，避免多实例本地时钟漂移导致索引提前/延后过期。 +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 +} + +// slotIndexSpec 描述一个活跃索引及其对应的槽位/等待键构造方式。 +// 用具名字段避免把 slotKey/waitKey 两个同签名函数按位置传参时写反。 +type slotIndexSpec struct { + indexKey string + slotKey func(int64) string + waitKey func(int64) string +} + +var ( + accountSlotIndex = slotIndexSpec{indexKey: accountActiveIndexKey, slotKey: accountSlotKey, waitKey: accountWaitKey} + userSlotIndex = slotIndexSpec{indexKey: userActiveIndexKey, slotKey: userSlotKey, waitKey: waitQueueKey} +) + +// 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 + } + if err := c.rdb.ZAdd(ctx, indexKey, redis.Z{ + Score: float64(expireAt), + Member: strconv.FormatInt(id, 10), + }).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) { + c.refreshActiveIndex(ctx, accountActiveIndexKey, accountID, accountSlotKey(accountID), accountWaitKey(accountID)) +} + +func (c *concurrencyCache) refreshUserActiveIndex(ctx context.Context, userID int64) { + c.refreshActiveIndex(ctx, userActiveIndexKey, userID, userSlotKey(userID), waitQueueKey(userID)) +} + +// 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 { + 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 + } + + ttlSeconds := c.activeIndexTTL(load.slotCount, load.waitCount) + if ttlSeconds <= 0 { + return + } + c.touchActiveIndexAt(ctx, indexKey, id, now+int64(ttlSeconds)) +} + +type activeIndexLoad struct { + id int64 + member string + slotCount int + waitCount int +} + +// activeIndexTTL 取槽位 TTL 与等待队列 TTL 中仍然需要关注的较大值。 +// 只要并发槽位或等待计数还有负载，就保留索引；两者都为 0 时调用方会删除索引。 +func (c *concurrencyCache) activeIndexTTL(slotCount int, waitCount int) int { + ttlSeconds := 0 + if slotCount > 0 { + ttlSeconds = c.slotTTLSeconds + } + if waitCount > 0 && c.waitQueueTTLSeconds > ttlSeconds { + ttlSeconds = c.waitQueueTTLSeconds + } + return ttlSeconds +} + +// readActiveLoadForKey 读取单个 ID 的当前负载，并顺手清理该槽位集合中的过期成员。 +func (c *concurrencyCache) readActiveLoadForKey(ctx context.Context, id int64, slotKey, waitKey string, now int64) (activeIndexLoad, error) { + cutoffTime := now - int64(c.slotTTLSeconds) + pipe := c.rdb.Pipeline() + pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10)) + zcardCmd := pipe.ZCard(ctx, slotKey) + getCmd := pipe.Get(ctx, waitKey) + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return activeIndexLoad{}, fmt.Errorf("pipeline exec: %w", err) + } + + waitCount := 0 + if v, err := getCmd.Int(); err == nil && v > 0 { + waitCount = v + } + return activeIndexLoad{ + id: id, + member: strconv.FormatInt(id, 10), + slotCount: int(zcardCmd.Val()), + waitCount: waitCount, + }, nil +} + +// readIndexLoads 批量读取索引候选的真实负载（账号/用户通用）。 +// 分块 Pipeline 可以减少 Redis 往返，同时避免一次 Pipeline 塞入过多命令。 +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)) + for _, member := range members { + id, err := strconv.ParseInt(member, 10, 64) + if err != nil || id <= 0 { + staleMembers = append(staleMembers, member) + continue + } + candidates = append(candidates, activeIndexLoad{id: id, member: member}) + } + + cutoffTime := now - int64(c.slotTTLSeconds) + for start := 0; start < len(candidates); start += activeIndexPipelineChunkSize { + end := start + activeIndexPipelineChunkSize + if end > len(candidates) { + end = len(candidates) + } + chunk := candidates[start:end] + + pipe := c.rdb.Pipeline() + type loadCmd struct { + activeIndexLoad + zcardCmd *redis.IntCmd + getCmd *redis.StringCmd + } + cmds := make([]loadCmd, 0, len(chunk)) + for _, candidate := range chunk { + slotKey := spec.slotKey(candidate.id) + waitKey := spec.waitKey(candidate.id) + pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10)) + cmds = append(cmds, loadCmd{ + activeIndexLoad: candidate, + zcardCmd: pipe.ZCard(ctx, slotKey), + getCmd: pipe.Get(ctx, waitKey), + }) + } + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, nil, fmt.Errorf("pipeline exec: %w", err) + } + for _, cmd := range cmds { + waitCount := 0 + if v, err := cmd.getCmd.Int(); err == nil && v > 0 { + waitCount = v + } + loads = append(loads, activeIndexLoad{ + id: cmd.id, + member: cmd.member, + slotCount: int(cmd.zcardCmd.Val()), + waitCount: waitCount, + }) + } + } + + return loads, staleMembers, nil +} + +// removeActiveIndexMembers 清理无效 member；这是辅助索引的维护动作，调用方无需因为失败中断主流程。 +func (c *concurrencyCache) removeActiveIndexMembers(ctx context.Context, indexKey string, members []string) { + if len(members) == 0 { + return + } + args := make([]any, 0, len(members)) + for _, member := range members { + args = append(args, member) + } + 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) + } +} + +// 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 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.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.slotTTLSeconds)) + } return result == 1, nil } func (c *concurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error { key := accountSlotKey(accountID) - return c.rdb.ZRem(ctx, key, requestID).Err() + if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil { + return err + } + // 释放后用真实负载刷新索引；若没有槽位和等待计数，会移除索引 member。 + c.refreshAccountActiveIndex(ctx, accountID) + return nil } func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID int64) (int, error) { @@ -331,16 +581,25 @@ 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.touchActiveIndexAt(ctx, userActiveIndexKey, userID, now+int64(c.slotTTLSeconds)) + } return result == 1, nil } func (c *concurrencyCache) ReleaseUserSlot(ctx context.Context, userID int64, requestID string) error { key := userSlotKey(userID) - return c.rdb.ZRem(ctx, key, requestID).Err() + if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil { + return err + } + // 释放后按 Redis 中剩余负载修正索引状态。 + c.refreshUserActiveIndex(ctx, userID) + return nil } func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) (int, error) { @@ -353,20 +612,76 @@ func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) return result, nil } +func (c *concurrencyCache) TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + key := apiKeySlotKey(apiKeyID) + _, err := trackSlotScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds, requestID).Result() + return err +} + +func (c *concurrencyCache) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + key := apiKeySlotKey(apiKeyID) + return c.rdb.ZRem(ctx, key, requestID).Err() +} + +func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) { + if len(apiKeyIDs) == 0 { + return map[int64]int{}, nil + } + + now, err := c.rdb.Time(ctx).Result() + if err != nil { + return nil, fmt.Errorf("redis TIME: %w", err) + } + cutoffTime := now.Unix() - int64(c.slotTTLSeconds) + + pipe := c.rdb.Pipeline() + type apiKeyCmd struct { + apiKeyID int64 + zcardCmd *redis.IntCmd + } + cmds := make([]apiKeyCmd, 0, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + slotKey := apiKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10) + pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10)) + cmds = append(cmds, apiKeyCmd{ + apiKeyID: apiKeyID, + zcardCmd: pipe.ZCard(ctx, slotKey), + }) + } + + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, fmt.Errorf("pipeline exec: %w", err) + } + + result := make(map[int64]int, len(apiKeyIDs)) + for _, cmd := range cmds { + result[cmd.apiKeyID] = int(cmd.zcardCmd.Val()) + } + return result, nil +} + // Wait queue operations 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.touchActiveIndexAt(ctx, userActiveIndexKey, userID, now+int64(c.waitQueueTTLSeconds)) + } return result == 1, nil } func (c *concurrencyCache) DecrementWaitCount(ctx context.Context, userID int64) error { key := waitQueueKey(userID) _, err := decrementWaitScript.Run(ctx, c.rdb, []string{key}).Result() + if err == nil { + // 等待数减少后重新判断是否还需要保留索引。 + c.refreshUserActiveIndex(ctx, userID) + } return err } @@ -374,16 +689,24 @@ 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.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.waitQueueTTLSeconds)) + } return result == 1, nil } func (c *concurrencyCache) DecrementAccountWaitCount(ctx context.Context, accountID int64) error { key := accountWaitKey(accountID) _, err := decrementWaitScript.Run(ctx, c.rdb, []string{key}).Result() + if err == nil { + // 等待计数归零后索引需要同步删除，避免后台任务反复处理空账号。 + c.refreshAccountActiveIndex(ctx, accountID) + } return err } @@ -523,101 +846,183 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic func (c *concurrencyCache) CleanupExpiredAccountSlots(ctx context.Context, accountID int64) error { key := accountSlotKey(accountID) _, err := cleanupExpiredSlotsScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Result() + if err == nil { + // 单账号清理后同步索引，保持后台批量清理的候选集准确。 + c.refreshAccountActiveIndex(ctx, accountID) + } return err } +// CleanupExpiredAccountSlotKeys 处理账号与用户两个活跃索引中已到期的候选。 +// （方法名中的 Account 是历史遗留，保留以避免接口变更；实际同时回收两个索引， +// 否则 user 索引的过期成员没有任何清理路径，会无界累积。） func (c *concurrencyCache) CleanupExpiredAccountSlotKeys(ctx context.Context) error { - return c.cleanupExpiredSlotKeysByPattern(ctx, accountSlotKeyPrefix+"*") + if err := c.reconcileExpiredIndexCandidates(ctx, accountSlotIndex); err != nil { + return err + } + return c.reconcileExpiredIndexCandidates(ctx, userSlotIndex) } +// 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, spec.indexKey, &redis.ZRangeBy{ + Min: "-inf", + Max: strconv.FormatInt(now, 10), + Count: activeIndexCleanupBatchSize, + }).Result() + if err != nil { + return fmt.Errorf("read expired index %s: %w", spec.indexKey, err) + } + + 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 + } + refreshed = append(refreshed, redis.Z{ + Score: float64(now + int64(c.activeIndexTTL(load.slotCount, load.waitCount))), + Member: load.member, + }) + } + 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 启动时清理非当前进程前缀的槽位。 +// 清理范围来自活跃索引（含 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 } - - // 1. 清理有序集合中非当前进程前缀的成员 - slotPatterns := []string{accountSlotKeyPrefix + "*", userSlotKeyPrefix + "*"} - for _, pattern := range slotPatterns { - if err := c.cleanupSlotsByPattern(ctx, pattern, activeRequestPrefix); err != nil { - return err - } + if err := c.sweepLegacyWaitKeysOnce(ctx); err != nil { + return err + } + now, err := c.redisUnixSeconds(ctx) + if err != nil { + return err } - // 2. 删除所有等待队列计数器（重启后计数器失效） - waitPatterns := []string{accountWaitKeyPrefix + "*", waitQueueKeyPrefix + "*"} - for _, pattern := range waitPatterns { - if err := c.deleteKeysByPattern(ctx, pattern); err != nil { - return err - } + accountMembers, err := c.allIndexMembers(ctx, accountActiveIndexKey) + if err != nil { + return err + } + if err := c.cleanupStaleProcessSlotsForIndex(ctx, accountSlotIndex, accountMembers, activeRequestPrefix, now); err != nil { + return err } - return nil + userMembers, err := c.allIndexMembers(ctx, userActiveIndexKey) + if err != nil { + return err + } + return c.cleanupStaleProcessSlotsForIndex(ctx, userSlotIndex, userMembers, activeRequestPrefix, now) } -// cleanupExpiredSlotKeysByPattern 扫描实际存在的账号槽位键并批量清理过期成员。 -func (c *concurrencyCache) cleanupExpiredSlotKeysByPattern(ctx context.Context, pattern string) error { - const scanCount = 200 - var cursor uint64 - for { - keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, scanCount).Result() - if err != nil { - return fmt.Errorf("scan %s: %w", pattern, err) - } - if len(keys) > 0 { - _, err := cleanupExpiredSlotKeysScript.Run(ctx, c.rdb, keys, c.slotTTLSeconds).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("cleanup expired slots %s: %w", pattern, err) + 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 } } - cursor = nextCursor - 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 } -// cleanupSlotsByPattern 扫描匹配 pattern 的有序集合键，批量调用 Lua 脚本清理非当前进程成员。 -func (c *concurrencyCache) cleanupSlotsByPattern(ctx context.Context, pattern, activePrefix string) error { - const scanCount = 200 - var cursor uint64 - for { - keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, scanCount).Result() - if err != nil { - return fmt.Errorf("scan %s: %w", pattern, err) - } - if len(keys) > 0 { - _, err := startupCleanupScript.Run(ctx, c.rdb, keys, activePrefix, c.slotTTLSeconds).Result() - if err != nil { - return fmt.Errorf("cleanup slots %s: %w", pattern, err) - } - } - cursor = nextCursor - if cursor == 0 { - break - } +// 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) } - return nil + return members, nil } -// deleteKeysByPattern 扫描匹配 pattern 的键并删除。 -func (c *concurrencyCache) deleteKeysByPattern(ctx context.Context, pattern string) error { - const scanCount = 200 - var cursor uint64 - for { - keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, scanCount).Result() +// cleanupStaleProcessSlotsForIndex 逐个处理索引中的账号/用户。 +// Lua 脚本一次只碰一个槽位 key，兼容 Redis Cluster，随后删除重启后已失效的等待计数； +// 索引 member 的去留由脚本返回的剩余槽位数决定，最后批量写回。 +func (c *concurrencyCache) cleanupStaleProcessSlotsForIndex( + ctx context.Context, + spec slotIndexSpec, + members []string, + activeRequestPrefix string, + 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 { + staleMembers = append(staleMembers, member) + continue + } + + _, remaining, err := runScriptInt64Pair(ctx, c.rdb, startupCleanupSlotScript, []string{spec.slotKey(id)}, activeRequestPrefix, c.slotTTLSeconds) if err != nil { - return fmt.Errorf("scan %s: %w", pattern, err) + return fmt.Errorf("cleanup stale process slots %s: %w", spec.slotKey(id), err) } - if len(keys) > 0 { - if err := c.rdb.Del(ctx, keys...).Err(); err != nil { - return fmt.Errorf("del %s: %w", pattern, 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) } - cursor = nextCursor - if cursor == 0 { - break + if remaining > 0 { + refreshed = append(refreshed, redis.Z{ + Score: float64(now + int64(c.slotTTLSeconds)), + Member: member, + }) + } else { + staleMembers = append(staleMembers, member) } } + 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 4aa172b252..f7e27d1118 100644 --- a/backend/internal/repository/concurrency_cache_integration_test.go +++ b/backend/internal/repository/concurrency_cache_integration_test.go @@ -3,8 +3,10 @@ package repository import ( + "context" "errors" "fmt" + "strconv" "testing" "time" @@ -22,7 +24,8 @@ var testSlotTTL = time.Duration(testSlotTTLMinutes) * time.Minute type ConcurrencyCacheSuite struct { IntegrationRedisSuite - cache service.ConcurrencyCache + cache service.ConcurrencyCache + rawCache *concurrencyCache } func TestConcurrencyCacheSuite(t *testing.T) { @@ -31,7 +34,20 @@ func TestConcurrencyCacheSuite(t *testing.T) { func (s *ConcurrencyCacheSuite) SetupTest() { s.IntegrationRedisSuite.SetupTest() - s.cache = NewConcurrencyCache(s.rdb, testSlotTTLMinutes, int(testSlotTTL.Seconds())) + s.rawCache = NewConcurrencyCache(s.rdb, testSlotTTLMinutes, int(testSlotTTL.Seconds())).(*concurrencyCache) + s.cache = s.rawCache +} + +type apiKeyConcurrencyCacheForTest interface { + TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) +} + +func (s *ConcurrencyCacheSuite) apiKeyConcurrencyCache() apiKeyConcurrencyCacheForTest { + cache, ok := s.cache.(apiKeyConcurrencyCacheForTest) + require.True(s.T(), ok) + return cache } func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() { @@ -61,6 +77,63 @@ func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() { require.Equal(s.T(), 1, cur, "expected 1 after release") } +func (s *ConcurrencyCacheSuite) TestAccountActiveIndex_AcquireAndRelease() { + accountID := int64(610) + 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) + + score, err := s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result() + require.NoError(s.T(), err) + 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)) + + _, 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) 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) + + _, 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)) + + _, 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) TestUserActiveIndex_AcquireAndRelease() { + userID := int64(612) + member := strconv.FormatInt(userID, 10) + reqID := "user-active-index-req" + + ok, err := s.cache.AcquireUserSlot(s.ctx, userID, 2, reqID) + require.NoError(s.T(), err) + require.True(s.T(), ok) + + _, 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() { accountID := int64(11) reqID := "req_ttl_test" @@ -160,6 +233,34 @@ func (s *ConcurrencyCacheSuite) TestUserSlot_TTL() { s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL) } +func (s *ConcurrencyCacheSuite) TestAPIKeySlot_TrackReleaseAndBatchCount() { + cache := s.apiKeyConcurrencyCache() + apiKeyID := int64(300) + emptyAPIKeyID := int64(301) + slotKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID) + + require.NoError(s.T(), cache.TrackAPIKeySlot(s.ctx, apiKeyID, "req1")) + require.NoError(s.T(), cache.TrackAPIKeySlot(s.ctx, apiKeyID, "req2")) + + counts, err := cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID, emptyAPIKeyID}) + require.NoError(s.T(), err) + require.Equal(s.T(), map[int64]int{apiKeyID: 2, emptyAPIKeyID: 0}, counts) + + ttl, err := s.rdb.TTL(s.ctx, slotKey).Result() + require.NoError(s.T(), err, "TTL") + s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL) + + require.NoError(s.T(), cache.ReleaseAPIKeySlot(s.ctx, apiKeyID, "req1")) + counts, err = cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID}) + require.NoError(s.T(), err) + require.Equal(s.T(), 1, counts[apiKeyID]) + + require.NoError(s.T(), cache.ReleaseAPIKeySlot(s.ctx, apiKeyID, "req2")) + counts, err = cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID}) + require.NoError(s.T(), err) + require.Equal(s.T(), 0, counts[apiKeyID]) +} + func (s *ConcurrencyCacheSuite) TestWaitQueue_IncrementAndDecrement() { userID := int64(20) waitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) @@ -252,14 +353,22 @@ 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) + unindexedAccountID := int64(1901) accountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID) userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID) + apiKeyKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID) + unindexedAccountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, unindexedAccountID) userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID) + unindexedAccountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, unindexedAccountID) - now := time.Now().Unix() + 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"}, redis.Z{Score: float64(now), Member: "keep-1"}, @@ -268,8 +377,24 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() { redis.Z{Score: float64(now), Member: "oldproc-2"}, redis.Z{Score: float64(now), Member: "keep-2"}, ).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, unindexedAccountKey, + redis.Z{Score: float64(now), Member: "oldproc-unindexed"}, + ).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, apiKeyKey, + redis.Z{Score: float64(now), Member: "oldproc-3"}, + redis.Z{Score: float64(now), Member: "keep-3"}, + ).Err()) require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, time.Minute).Err()) require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, time.Minute).Err()) + require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 2, time.Minute).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{ + Score: float64(now + 60), + Member: strconv.FormatInt(accountID, 10), + }).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{ + Score: float64(now + 60), + Member: strconv.FormatInt(userID, 10), + }).Err()) require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-")) @@ -281,11 +406,22 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() { require.NoError(s.T(), err) require.Equal(s.T(), []string{"keep-2"}, userMembers) + // API Key 槽位（stats-only）不在启动清理范围内，靠分数裁剪与 key TTL 自愈。 + apiKeyMembers, err := s.rdb.ZRange(s.ctx, apiKeyKey, 0, -1).Result() + require.NoError(s.T(), err) + require.ElementsMatch(s.T(), []string{"keep-3", "oldproc-3"}, apiKeyMembers) + _, err = s.rdb.Get(s.ctx, userWaitKey).Result() require.True(s.T(), errors.Is(err, redis.Nil)) _, err = s.rdb.Get(s.ctx, accountWaitKey).Result() require.True(s.T(), errors.Is(err, redis.Nil)) + + unindexedMembers, err := s.rdb.ZRange(s.ctx, unindexedAccountKey, 0, -1).Result() + require.NoError(s.T(), err) + require.Equal(s.T(), []string{"oldproc-unindexed"}, unindexedMembers) + _, err = s.rdb.Get(s.ctx, unindexedAccountWaitKey).Result() + require.NoError(s.T(), err) } func (s *ConcurrencyCacheSuite) TestGetAccountConcurrency_Missing() { @@ -436,11 +572,13 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlots_NoExpired() { } func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() { - now := time.Now().Unix() + now, err := s.rawCache.redisUnixSeconds(s.ctx) + require.NoError(s.T(), err) expiredTime := now - int64(testSlotTTL.Seconds()) - 10 accountKeyWithFresh := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 301) accountKeyExpiredOnly := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 302) userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, 303) + unindexedAccountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 304) require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKeyWithFresh, redis.Z{Score: float64(expiredTime), Member: "expired"}, @@ -452,6 +590,13 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() { require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKey, redis.Z{Score: float64(expiredTime), Member: "user-expired"}, ).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, unindexedAccountKey, + redis.Z{Score: float64(expiredTime), Member: "unindexed-expired"}, + ).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, + redis.Z{Score: float64(now), Member: "301"}, + redis.Z{Score: float64(now), Member: "302"}, + ).Err()) require.NoError(s.T(), s.cache.CleanupExpiredAccountSlotKeys(s.ctx)) @@ -466,9 +611,125 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() { userMembers, err := s.rdb.ZRange(s.ctx, userKey, 0, -1).Result() require.NoError(s.T(), err) require.Equal(s.T(), []string{"user-expired"}, userMembers) + + unindexedMembers, err := s.rdb.ZRange(s.ctx, unindexedAccountKey, 0, -1).Result() + require.NoError(s.T(), err) + require.Equal(s.T(), []string{"unindexed-expired"}, unindexedMembers) + + score, err := s.rdb.ZScore(s.ctx, accountActiveIndexKey, "301").Result() + require.NoError(s.T(), err) + require.Greater(s.T(), int64(score), now) + _, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, "302").Result() + 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) @@ -476,19 +737,28 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesA userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID) - now := float64(time.Now().Unix()) + now, err := s.rawCache.redisUnixSeconds(s.ctx) + require.NoError(s.T(), err) require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey, - redis.Z{Score: now, Member: "oldproc-1"}, - redis.Z{Score: now, Member: "activeproc-1"}, + redis.Z{Score: float64(now), Member: "oldproc-1"}, + redis.Z{Score: float64(now), Member: "activeproc-1"}, ).Err()) require.NoError(s.T(), s.rdb.Expire(s.ctx, accountSlotKey, testSlotTTL).Err()) require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userSlotKey, - redis.Z{Score: now, Member: "oldproc-2"}, - redis.Z{Score: now, Member: "activeproc-2"}, + redis.Z{Score: float64(now), Member: "oldproc-2"}, + redis.Z{Score: float64(now), Member: "activeproc-2"}, ).Err()) require.NoError(s.T(), s.rdb.Expire(s.ctx, userSlotKey, testSlotTTL).Err()) require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, testSlotTTL).Err()) require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, testSlotTTL).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{ + Score: float64(now + 60), + Member: strconv.FormatInt(accountID, 10), + }).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{ + Score: float64(now + 60), + Member: strconv.FormatInt(userID, 10), + }).Err()) require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "activeproc-")) @@ -509,8 +779,14 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesA func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_DeletesEmptySlotKeys() { accountID := int64(903) accountSlotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID) - require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey, redis.Z{Score: float64(time.Now().Unix()), Member: "oldproc-1"}).Err()) + now, err := s.rawCache.redisUnixSeconds(s.ctx) + require.NoError(s.T(), err) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey, redis.Z{Score: float64(now), Member: "oldproc-1"}).Err()) require.NoError(s.T(), s.rdb.Expire(s.ctx, accountSlotKey, testSlotTTL).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{ + Score: float64(now + 60), + Member: strconv.FormatInt(accountID, 10), + }).Err()) require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "activeproc-")) diff --git a/backend/internal/repository/group_repo.go b/backend/internal/repository/group_repo.go index a4e173006e..0aab6b0c05 100644 --- a/backend/internal/repository/group_repo.go +++ b/backend/internal/repository/group_repo.go @@ -50,11 +50,14 @@ 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). SetDefaultValidityDays(groupIn.DefaultValidityDays). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetNillableFallbackGroupID(groupIn.FallbackGroupID). @@ -132,11 +135,14 @@ 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). SetDefaultValidityDays(groupIn.DefaultValidityDays). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetModelRoutingEnabled(groupIn.ModelRoutingEnabled). 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/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index c508f09c71..c8e1fe14e0 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -518,7 +518,7 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any { if len(credentials) == 0 { return nil } - keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type"} + keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type"} filtered := make(map[string]any) for _, key := range keys { if value, ok := credentials[key]; ok && value != nil { diff --git a/backend/internal/repository/scheduler_cache_test.go b/backend/internal/repository/scheduler_cache_test.go new file mode 100644 index 0000000000..f438b2f456 --- /dev/null +++ b/backend/internal/repository/scheduler_cache_test.go @@ -0,0 +1,37 @@ +package repository + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestFilterSchedulerCredentialsKeepsSubscriptionPlanType(t *testing.T) { + filtered := filterSchedulerCredentials(map[string]any{ + "plan_type": "plus", + "access_token": "secret-access-token", + "refresh_token": "secret-refresh-token", + }) + + require.Equal(t, "plus", filtered["plan_type"]) + require.NotContains(t, filtered, "access_token") + require.NotContains(t, filtered, "refresh_token") +} + +func TestSchedulerMetadataAccountKeepsOpenAISubscriptionIdentity(t *testing.T) { + account := service.Account{ + ID: 24, + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Credentials: map[string]any{ + "plan_type": "plus", + "access_token": "secret-access-token", + }, + } + + metadata := buildSchedulerMetadataAccount(account) + + require.True(t, metadata.IsOpenAIChatGPTSubscription()) + require.Empty(t, metadata.GetCredential("access_token")) +} diff --git a/backend/internal/repository/usage_billing_repo.go b/backend/internal/repository/usage_billing_repo.go index 91ac536eee..f7e675439f 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,108 @@ 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 + } + 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") +} + +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..0c469db899 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,111 @@ 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(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, 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()) +} diff --git a/backend/internal/repository/user_msg_queue_cache.go b/backend/internal/repository/user_msg_queue_cache.go index 2b5b96bdc3..9b7707614c 100644 --- a/backend/internal/repository/user_msg_queue_cache.go +++ b/backend/internal/repository/user_msg_queue_cache.go @@ -5,9 +5,8 @@ import ( "errors" "fmt" "strconv" - "strings" - "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" ) @@ -18,18 +17,41 @@ const ( umqKeyPrefix = "umq:" umqLockSuffix = ":lock" // STRING (requestID), PX lockTtlMs umqLastSuffix = ":last" // STRING (毫秒时间戳), EX 60s + + // 锁索引用来替代后台清理对 umq:*:lock 的全量 SCAN。 + // member 是 accountID，score 是锁预计过期的 Redis Unix 毫秒时间戳。 + umqLockIndexKey = "umq:lock:index" // ZSET member=accountID, score=lockExpireAtUnixMs + umqLockIndexCleanupBatchSize = 1000 ) // 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]) +local ttl = tonumber(ARGV[2]) if cur == ARGV[1] then - redis.call('PEXPIRE', KEYS[1], tonumber(ARGV[2])) - return 1 + 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 end -redis.call('SET', KEYS[1], ARGV[1], 'PX', tonumber(ARGV[2])) -return 1 +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) +return {1, ms + ttl} `) // Lua 脚本：原子释放锁 + 记录完成时间（使用 Redis TIME 避免时钟偏差） @@ -48,14 +70,18 @@ end return 0 `) -// Lua 脚本：原子清理孤儿锁（仅在 PTTL == -1 时删除，避免 TOCTOU 竞态误删合法锁） -var forceReleaseLockScript = redis.NewScript(` +// Lua 脚本：校验锁 TTL 状态，PTTL == -1 时原子删除异常锁。 +// 返回状态: -2=锁不存在，-1=无 TTL 的异常锁已删除，1=锁仍存活并返回剩余 PTTL。 +var reconcileLockScript = redis.NewScript(` 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 + return {-1, 0} end -return 0 +return {1, pttl} `) type userMsgQueueCache struct { @@ -77,22 +103,36 @@ func umqLastKey(accountID int64) string { return umqKeyPrefix + "{" + strconv.FormatInt(accountID, 10) + "}" + umqLastSuffix } -// umqScanPattern 用于 SCAN 扫描锁 key -func umqScanPattern() string { - return umqKeyPrefix + "{*}" + umqLockSuffix -} - // AcquireLock 尝试获取账号级串行锁 +// 无论成功与否都尽力写入锁索引：成功时登记自己的锁，失败时回填观测到的持有者锁， +// 保证任何被争用的锁都能被后台 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).Int() + result, err := acquireLockScript.Run(ctx, c.rdb, []string{key}, requestID, lockTtlMs).Result() if err != nil { return false, fmt.Errorf("umq acquire lock: %w", err) } - return result == 1, nil + acquired, err := redisScriptInt64At(result, 0) + if err != nil { + return false, fmt.Errorf("umq parse acquire lock result: %w", err) + } + expireAtMs, err := redisScriptInt64At(result, 1) + if err != nil { + return false, fmt.Errorf("umq parse acquire lock expire: %w", err) + } + if expireAtMs > 0 { + if err := c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{ + Score: float64(expireAtMs), + Member: strconv.FormatInt(accountID, 10), + }).Err(); err != nil { + logger.LegacyPrintf("repository.umq", "Warning: update lock index for account %d failed: %v", accountID, err) + } + } + return acquired == 1, nil } // ReleaseLock 释放锁并记录完成时间 +// 只有 requestID 匹配时才删除锁索引，避免误删其他请求重入后写入的新锁。 func (c *userMsgQueueCache) ReleaseLock(ctx context.Context, accountID int64, requestID string) (bool, error) { lockKey := umqLockKey(accountID) lastKey := umqLastKey(accountID) @@ -100,6 +140,13 @@ func (c *userMsgQueueCache) ReleaseLock(ctx context.Context, accountID int64, re if err != nil { return false, fmt.Errorf("umq release lock: %w", err) } + if result == 1 { + // 与下一个 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 } @@ -120,65 +167,6 @@ func (c *userMsgQueueCache) GetLastCompletedMs(ctx context.Context, accountID in return ms, nil } -// ForceReleaseLock 原子清理孤儿锁（仅在 PTTL == -1 时删除，防止 TOCTOU 竞态误删合法锁） -func (c *userMsgQueueCache) ForceReleaseLock(ctx context.Context, accountID int64) error { - key := umqLockKey(accountID) - _, err := forceReleaseLockScript.Run(ctx, c.rdb, []string{key}).Result() - if err != nil && !errors.Is(err, redis.Nil) { - return fmt.Errorf("umq force release lock: %w", err) - } - return nil -} - -// ScanLockKeys 扫描所有锁 key，仅返回 PTTL == -1（无过期时间）的孤儿锁 accountID 列表 -// 正常的锁都有 PX 过期时间，PTTL == -1 表示异常状态（如 Redis 故障恢复后丢失 TTL） -func (c *userMsgQueueCache) ScanLockKeys(ctx context.Context, maxCount int) ([]int64, error) { - var accountIDs []int64 - var cursor uint64 - pattern := umqScanPattern() - - for { - keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, 100).Result() - if err != nil { - return nil, fmt.Errorf("umq scan lock keys: %w", err) - } - for _, key := range keys { - // 检查 PTTL：只清理 PTTL == -1（无过期时间）的异常锁 - pttl, err := c.rdb.PTTL(ctx, key).Result() - if err != nil { - continue - } - // PTTL 返回值：-2 = key 不存在，-1 = 无过期时间，>0 = 剩余毫秒 - // go-redis 对哨兵值 -1/-2 不乘精度系数，直接返回 time.Duration(-1)/-2 - // 只删除 -1（无过期时间的异常锁），跳过正常持有的锁 - if pttl != time.Duration(-1) { - continue - } - - // 从 key 中提取 accountID: umq:{123}:lock → 提取 {} 内的数字 - openBrace := strings.IndexByte(key, '{') - closeBrace := strings.IndexByte(key, '}') - if openBrace < 0 || closeBrace <= openBrace+1 { - continue - } - idStr := key[openBrace+1 : closeBrace] - id, err := strconv.ParseInt(idStr, 10, 64) - if err != nil { - continue - } - accountIDs = append(accountIDs, id) - if len(accountIDs) >= maxCount { - return accountIDs, nil - } - } - cursor = nextCursor - if cursor == 0 { - break - } - } - return accountIDs, nil -} - // GetCurrentTimeMs 通过 Redis TIME 命令获取当前服务器时间（毫秒），确保与锁记录的时间源一致 func (c *userMsgQueueCache) GetCurrentTimeMs(ctx context.Context) (int64, error) { t, err := c.rdb.Time(ctx).Result() @@ -187,3 +175,94 @@ func (c *userMsgQueueCache) GetCurrentTimeMs(ctx context.Context) (int64, error) } return t.UnixMilli(), nil } + +// ReconcileExpiredLockCandidates 只处理索引里已经到期的候选锁。 +// 候选到期不等于锁一定失效：可能是续租后索引滞后，所以必须再用 PTTL 二次确认。 +func (c *userMsgQueueCache) ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (int, error) { + if maxCount <= 0 { + maxCount = umqLockIndexCleanupBatchSize + } + nowMs, err := c.GetCurrentTimeMs(ctx) + if err != nil { + return 0, err + } + members, err := c.rdb.ZRangeByScore(ctx, umqLockIndexKey, &redis.ZRangeBy{ + Min: "-inf", + Max: strconv.FormatInt(nowMs, 10), + Count: int64(maxCount), + }).Result() + if err != nil { + return 0, fmt.Errorf("umq read lock index: %w", err) + } + + cleaned := 0 + for _, member := range members { + accountID, err := strconv.ParseInt(member, 10, 64) + if err != nil || accountID <= 0 { + c.removeLockIndexMember(ctx, member) + continue + } + + result, err := reconcileLockScript.Run(ctx, c.rdb, []string{umqLockKey(accountID)}).Result() + if err != nil && !errors.Is(err, redis.Nil) { + return cleaned, fmt.Errorf("umq reconcile lock: %w", err) + } + status, err := redisScriptInt64At(result, 0) + if err != nil { + return cleaned, fmt.Errorf("umq parse reconcile status: %w", err) + } + pttl, err := redisScriptInt64At(result, 1) + if err != nil { + return cleaned, fmt.Errorf("umq parse reconcile pttl: %w", err) + } + + switch status { + case -2: + // 锁自然过期或已释放，只需移除索引残留。 + c.removeLockIndexMember(ctx, member) + case -1: + // 无 TTL 的锁会永久阻塞队列，Lua 已原子删除它，这里统计一次清理。 + c.removeLockIndexMember(ctx, member) + cleaned++ + case 1: + // 锁仍存活，说明索引过期时间滞后；按剩余 PTTL 重新排期。 + if err := c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{ + Score: float64(nowMs + pttl), + Member: member, + }).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) + if !ok { + return 0, fmt.Errorf("expected redis script array, got %T", result) + } + if index < 0 || index >= len(values) { + return 0, fmt.Errorf("redis script array missing index %d", index) + } + switch v := values[index].(type) { + case int64: + return v, nil + case int: + return int64(v), nil + case string: + return strconv.ParseInt(v, 10, 64) + case []byte: + return strconv.ParseInt(string(v), 10, 64) + default: + return 0, fmt.Errorf("unexpected redis script value %T", v) + } +} diff --git a/backend/internal/repository/user_msg_queue_cache_integration_test.go b/backend/internal/repository/user_msg_queue_cache_integration_test.go new file mode 100644 index 0000000000..e683b44aa3 --- /dev/null +++ b/backend/internal/repository/user_msg_queue_cache_integration_test.go @@ -0,0 +1,177 @@ +//go:build integration + +package repository + +import ( + "errors" + "testing" + "time" + + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" + "github.com/stretchr/testify/suite" +) + +type UserMsgQueueCacheSuite struct { + IntegrationRedisSuite + cache *userMsgQueueCache +} + +func TestUserMsgQueueCacheSuite(t *testing.T) { + suite.Run(t, new(UserMsgQueueCacheSuite)) +} + +func (s *UserMsgQueueCacheSuite) SetupTest() { + s.IntegrationRedisSuite.SetupTest() + s.cache = NewUserMsgQueueCache(s.rdb).(*userMsgQueueCache) +} + +func (s *UserMsgQueueCacheSuite) TestAcquireLockWritesIndexAndReleaseRemovesIt() { + accountID := int64(701) + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + require.NoError(s.T(), err) + + acquired, err := s.cache.AcquireLock(s.ctx, accountID, "req-701", 10_000) + require.NoError(s.T(), err) + require.True(s.T(), acquired) + + score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "701").Result() + require.NoError(s.T(), err) + require.Greater(s.T(), int64(score), nowMs) + + released, err := s.cache.ReleaseLock(s.ctx, accountID, "req-701") + require.NoError(s.T(), err) + require.True(s.T(), released) + + _, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "701").Result() + require.ErrorIs(s.T(), err, redis.Nil) +} + +func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRemovesNaturallyExpiredLockIndex() { + accountID := int64(702) + acquired, err := s.cache.AcquireLock(s.ctx, accountID, "req-702", 20) + require.NoError(s.T(), err) + require.True(s.T(), acquired) + + score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "702").Result() + require.NoError(s.T(), err) + require.Eventually(s.T(), func() bool { + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + return err == nil && nowMs >= int64(score) + }, time.Second, 10*time.Millisecond) + + cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000) + require.NoError(s.T(), err) + require.Equal(s.T(), 0, cleaned) + + _, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "702").Result() + require.ErrorIs(s.T(), err, redis.Nil) +} + +func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRefreshesLiveLockIndex() { + accountID := int64(703) + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + require.NoError(s.T(), err) + require.NoError(s.T(), s.rdb.Set(s.ctx, umqLockKey(accountID), "req-703", time.Minute).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, umqLockIndexKey, redis.Z{ + Score: float64(nowMs - 1), + Member: "703", + }).Err()) + + cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000) + require.NoError(s.T(), err) + require.Equal(s.T(), 0, cleaned) + + score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "703").Result() + require.NoError(s.T(), err) + require.Greater(s.T(), int64(score), nowMs) + exists, err := s.rdb.Exists(s.ctx, umqLockKey(accountID)).Result() + require.NoError(s.T(), err) + require.EqualValues(s.T(), 1, exists) +} + +func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesDeletesNoTTLLock() { + accountID := int64(704) + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + require.NoError(s.T(), err) + require.NoError(s.T(), s.rdb.Set(s.ctx, umqLockKey(accountID), "req-704", 0).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, umqLockIndexKey, redis.Z{ + Score: float64(nowMs), + Member: "704", + }).Err()) + + cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000) + require.NoError(s.T(), err) + require.Equal(s.T(), 1, cleaned) + + exists, err := s.rdb.Exists(s.ctx, umqLockKey(accountID)).Result() + require.NoError(s.T(), err) + require.EqualValues(s.T(), 0, exists) + _, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "704").Result() + require.ErrorIs(s.T(), err, redis.Nil) +} + +func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRemovesInvalidMember() { + nowMs, err := s.cache.GetCurrentTimeMs(s.ctx) + require.NoError(s.T(), err) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, umqLockIndexKey, redis.Z{ + Score: float64(nowMs), + Member: "not-an-account-id", + }).Err()) + + cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000) + require.NoError(s.T(), err) + require.Equal(s.T(), 0, cleaned) + + _, 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 26f976c87b..278e654834 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, + "current_concurrency": 0, "quota": 0, "quota_used": 0, "rate_limit_5h": 0, @@ -282,6 +284,7 @@ func TestAPIContracts(t *testing.T) { "ip_whitelist": null, "ip_blacklist": null, "last_used_at": null, + "current_concurrency": 0, "quota": 0, "quota_used": 0, "rate_limit_5h": 0, @@ -359,6 +362,9 @@ func TestAPIContracts(t *testing.T) { "image_price_2k": null, "image_price_4k": 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, "claude_code_only": false, @@ -661,15 +667,17 @@ func TestAPIContracts(t *testing.T) { service.SettingKeyTableDefaultPageSize: "20", service.SettingKeyTablePageSizeOptions: "[10,20,50,100]", - service.SettingKeyOpsMonitoringEnabled: "false", - service.SettingKeyOpsRealtimeMonitoringEnabled: "true", - service.SettingKeyOpsQueryModeDefault: "auto", - service.SettingKeyOpsMetricsIntervalSeconds: "60", - service.SettingPaymentVisibleMethodAlipaySource: service.VisibleMethodSourceEasyPayAlipay, - service.SettingPaymentVisibleMethodWxpaySource: service.VisibleMethodSourceOfficialWechat, - service.SettingPaymentVisibleMethodAlipayEnabled: "true", - service.SettingPaymentVisibleMethodWxpayEnabled: "false", - "openai_advanced_scheduler_enabled": "true", + service.SettingKeyOpsMonitoringEnabled: "false", + service.SettingKeyOpsRealtimeMonitoringEnabled: "true", + service.SettingKeyOpsQueryModeDefault: "auto", + service.SettingKeyOpsMetricsIntervalSeconds: "60", + service.SettingPaymentVisibleMethodAlipaySource: service.VisibleMethodSourceEasyPayAlipay, + service.SettingPaymentVisibleMethodWxpaySource: service.VisibleMethodSourceOfficialWechat, + service.SettingPaymentVisibleMethodAlipayEnabled: "true", + service.SettingPaymentVisibleMethodWxpayEnabled: "false", + "openai_advanced_scheduler_enabled": "true", + service.SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled: "false", + service.SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled: "false", }) }, method: http.MethodGet, @@ -861,6 +869,28 @@ func TestAPIContracts(t *testing.T) { "payment_visible_method_alipay_enabled": true, "payment_visible_method_wxpay_enabled": false, "openai_advanced_scheduler_enabled": true, + "openai_advanced_scheduler_sticky_weighted_enabled": false, + "openai_advanced_scheduler_subscription_priority_enabled": false, + "openai_advanced_scheduler_lb_top_k": "", + "openai_advanced_scheduler_weight_priority": "", + "openai_advanced_scheduler_weight_load": "", + "openai_advanced_scheduler_weight_queue": "", + "openai_advanced_scheduler_weight_error_rate": "", + "openai_advanced_scheduler_weight_ttft": "", + "openai_advanced_scheduler_weight_reset": "", + "openai_advanced_scheduler_weight_quota_headroom": "", + "openai_advanced_scheduler_weight_previous_response": "", + "openai_advanced_scheduler_weight_session_sticky": "", + "openai_advanced_scheduler_effective_lb_top_k": "7", + "openai_advanced_scheduler_effective_weight_priority": "1", + "openai_advanced_scheduler_effective_weight_load": "1", + "openai_advanced_scheduler_effective_weight_queue": "0.7", + "openai_advanced_scheduler_effective_weight_error_rate": "0.8", + "openai_advanced_scheduler_effective_weight_ttft": "0.5", + "openai_advanced_scheduler_effective_weight_reset": "0", + "openai_advanced_scheduler_effective_weight_quota_headroom": "0", + "openai_advanced_scheduler_effective_weight_previous_response": "5", + "openai_advanced_scheduler_effective_weight_session_sticky": "3", "openai_codex_user_agent": "", "openai_fast_policy_settings": { "rules": [] @@ -875,6 +905,7 @@ func TestAPIContracts(t *testing.T) { "payment_max_pending_orders": 0, "payment_balance_disabled": false, "payment_balance_recharge_multiplier": 0, + "payment_subscription_usd_to_cny_rate": 0, "payment_recharge_fee_rate": 0, "payment_load_balance_strategy": "", "payment_product_name_prefix": "", @@ -1110,6 +1141,28 @@ func TestAPIContracts(t *testing.T) { "payment_visible_method_alipay_enabled": false, "payment_visible_method_wxpay_enabled": false, "openai_advanced_scheduler_enabled": false, + "openai_advanced_scheduler_sticky_weighted_enabled": false, + "openai_advanced_scheduler_subscription_priority_enabled": false, + "openai_advanced_scheduler_lb_top_k": "", + "openai_advanced_scheduler_weight_priority": "", + "openai_advanced_scheduler_weight_load": "", + "openai_advanced_scheduler_weight_queue": "", + "openai_advanced_scheduler_weight_error_rate": "", + "openai_advanced_scheduler_weight_ttft": "", + "openai_advanced_scheduler_weight_reset": "", + "openai_advanced_scheduler_weight_quota_headroom": "", + "openai_advanced_scheduler_weight_previous_response": "", + "openai_advanced_scheduler_weight_session_sticky": "", + "openai_advanced_scheduler_effective_lb_top_k": "7", + "openai_advanced_scheduler_effective_weight_priority": "1", + "openai_advanced_scheduler_effective_weight_load": "1", + "openai_advanced_scheduler_effective_weight_queue": "0.7", + "openai_advanced_scheduler_effective_weight_error_rate": "0.8", + "openai_advanced_scheduler_effective_weight_ttft": "0.5", + "openai_advanced_scheduler_effective_weight_reset": "0", + "openai_advanced_scheduler_effective_weight_quota_headroom": "0", + "openai_advanced_scheduler_effective_weight_previous_response": "5", + "openai_advanced_scheduler_effective_weight_session_sticky": "3", "openai_codex_user_agent": "", "openai_fast_policy_settings": { "rules": [] @@ -1123,6 +1176,7 @@ func TestAPIContracts(t *testing.T) { "payment_enabled_types": null, "payment_balance_disabled": false, "payment_balance_recharge_multiplier": 0, + "payment_subscription_usd_to_cny_rate": 0, "payment_recharge_fee_rate": 0, "payment_load_balance_strategy": "", "payment_product_name_prefix": "", @@ -1690,6 +1744,10 @@ func (s *stubAccountRepo) List(ctx context.Context, params pagination.Pagination return nil, nil, errors.New("not implemented") } +func (s *stubAccountRepo) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]service.Account, error) { + return nil, nil +} + func (s *stubAccountRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, *pagination.PaginationResult, error) { return nil, nil, errors.New("not implemented") } diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index 2f0a3f1cf7..9bcb56d7fa 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,16 @@ func setGroupContext(c *gin.Context, group *service.Group) { c.Request = c.Request.WithContext(ctx) } +func apiKeyBalanceBelowAuthThreshold(balance float64, cfg *config.Config) bool { + if balance <= 0 { + return true + } + if cfg == nil || cfg.Billing.MinimumBalanceReserve <= 0 { + return false + } + return balance < cfg.Billing.MinimumBalanceReserve +} + 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..5c5ee147a4 100644 --- a/backend/internal/server/middleware/api_key_auth_google.go +++ b/backend/internal/server/middleware/api_key_auth_google.go @@ -109,7 +109,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..899cd8bbe9 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,42 @@ func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) { require.Equal(t, "PERMISSION_DENIED", resp.Error.Status) } +func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(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.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.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..04ab9410ac 100644 --- a/backend/internal/server/middleware/api_key_auth_test.go +++ b/backend/internal/server/middleware/api_key_auth_test.go @@ -1000,6 +1000,49 @@ func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) { require.Equal(t, 1, touchCalls) } +func TestAPIKeyAuthRejectsBalanceBelowMinimumReserve(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) + + 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/gateway.go b/backend/internal/server/routes/gateway.go index 9522578051..d22e339c75 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -164,6 +164,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) } diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 26af0a7281..ae25bf387d 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -70,6 +70,14 @@ type Account struct { modelMappingCacheRawPtr uintptr modelMappingCacheRawLen int modelMappingCacheRawSig uint64 + + // header_overrides 热路径缓存（非持久化字段，同 model_mapping 缓存先例） + headerOverrideCache map[string]string + headerOverrideCacheReady bool + headerOverrideCacheCredentialsPtr uintptr + headerOverrideCacheRawPtr uintptr + headerOverrideCacheRawLen int + headerOverrideCacheRawSig uint64 } type OpenAIEndpointCapability string @@ -1182,6 +1190,18 @@ func (a *Account) IsOpenAIOAuth() bool { return a.IsOpenAI() && a.Type == AccountTypeOAuth } +func (a *Account) IsOpenAIChatGPTSubscription() bool { + if !a.IsOpenAIOAuth() { + return false + } + switch strings.ToLower(strings.TrimSpace(a.GetCredential("plan_type"))) { + case "", "free", "abnormal": + return false + default: + return true + } +} + func (a *Account) IsOpenAIPersonalAccessToken() bool { if !a.IsOpenAIOAuth() { return false diff --git a/backend/internal/service/account_header_override.go b/backend/internal/service/account_header_override.go new file mode 100644 index 0000000000..8882bbef91 --- /dev/null +++ b/backend/internal/service/account_header_override.go @@ -0,0 +1,280 @@ +package service + +import ( + "net/http" + "strings" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + + "golang.org/x/net/http/httpguts" +) + +// 请求头覆写（header override）：仅对 Anthropic / OpenAI 平台的 api_key 账号生效。 +// 管理员在账号上配置一组 header name -> value，转发到上游前用配置值覆盖同名请求头 +// （匹配不区分大小写）；value 为空的条目视为"未填写"，不参与覆盖。 +const ( + credKeyHeaderOverrideEnabled = "header_override_enabled" + credKeyHeaderOverrides = "header_overrides" + + maxHeaderOverrideEntries = 64 + maxHeaderOverrideNameLength = 200 + maxHeaderOverrideValueLength = 8192 +) + +// headerOverrideBlockedNames 禁止覆写的请求头（小写）。 +// - 连接控制/逐跳头：由 HTTP 栈管理，覆写会破坏请求传输； +// - host/content-length：由 Go 的 Request.Host / ContentLength 字段管理，header 覆写不生效或产生冲突； +// - content-type：承载报文框架信息（multipart boundary 为每请求随机值），静态覆写必然与 body 不匹配； +// - authorization/x-api-key/cookie 等：上游认证头由账号凭据统一注入，禁止通过覆写篡改或重新引入； +// - accept-encoding：强制压缩会破坏网关对上游流式响应（SSE/usage）的解析； +// - sec-websocket-*：WebSocket 握手头由拨号器管理（OpenAI WS 模式）； +// - session_id/x-claude-code-session-id 等：逐请求会话隔离头，固定值会造成会话串扰。 +var headerOverrideBlockedNames = map[string]struct{}{ + "host": {}, + "content-length": {}, + "content-type": {}, + "transfer-encoding": {}, + "connection": {}, + "keep-alive": {}, + "proxy-authenticate": {}, + "proxy-authorization": {}, + "proxy-connection": {}, + "te": {}, + "trailer": {}, + "upgrade": {}, + "authorization": {}, + "x-api-key": {}, + "x-goog-api-key": {}, + "cookie": {}, + "accept-encoding": {}, + "sec-websocket-key": {}, + "sec-websocket-version": {}, + "sec-websocket-extensions": {}, + "sec-websocket-protocol": {}, + "sec-websocket-accept": {}, + "session_id": {}, + "conversation_id": {}, + "x-codex-turn-state": {}, + "x-codex-turn-metadata": {}, + "chatgpt-account-id": {}, + "x-claude-code-session-id": {}, + "x-client-request-id": {}, +} + +func isHeaderOverrideBlockedName(lowerName string) bool { + _, blocked := headerOverrideBlockedNames[lowerName] + return blocked +} + +// IsHeaderOverrideEligible 报告账号类型是否支持请求头覆写。 +// 目前仅开放 Anthropic / OpenAI 两个平台的 api_key 账号。 +func (a *Account) IsHeaderOverrideEligible() bool { + if a == nil || a.Type != AccountTypeAPIKey { + return false + } + return a.Platform == PlatformAnthropic || a.Platform == PlatformOpenAI +} + +// IsHeaderOverrideEnabled 报告账号是否启用了请求头覆写。 +func (a *Account) IsHeaderOverrideEnabled() bool { + if !a.IsHeaderOverrideEligible() || a.Credentials == nil { + return false + } + enabled, ok := a.Credentials[credKeyHeaderOverrideEnabled].(bool) + return ok && enabled +} + +// GetHeaderOverrides 返回生效的请求头覆写表（key 统一小写）。 +// 未启用、不符合平台/类型条件或配置为空时返回 nil。 +// 空 value 的条目（模板占位）与非法/禁止的 header 名会被跳过。 +// 结果带热路径缓存（同 GetModelMapping 先例）：同一 credentials 映射在 +// 一次请求 / 一条 WS 会话内的多次调用只做一次解析与校验。 +func (a *Account) GetHeaderOverrides() map[string]string { + if !a.IsHeaderOverrideEnabled() { + return nil + } + rawMapping, rawIsAnyMap := a.Credentials[credKeyHeaderOverrides].(map[string]any) + if !rawIsAnyMap { + // 非 JSON 反序列化产物（如直接注入的 map[string]string）：直接解析，不缓存 + return resolveHeaderOverrides(stringMappingFromRaw(a.Credentials[credKeyHeaderOverrides])) + } + + credentialsPtr := mapPtr(a.Credentials) + rawPtr := mapPtr(rawMapping) + rawLen := len(rawMapping) + rawSig := uint64(0) + rawSigReady := false + + if a.headerOverrideCacheReady && + a.headerOverrideCacheCredentialsPtr == credentialsPtr && + a.headerOverrideCacheRawPtr == rawPtr && + a.headerOverrideCacheRawLen == rawLen { + rawSig = modelMappingSignature(rawMapping) + rawSigReady = true + if a.headerOverrideCacheRawSig == rawSig { + return a.headerOverrideCache + } + } + + overrides := resolveHeaderOverrides(stringMappingFromRaw(rawMapping)) + if !rawSigReady { + rawSig = modelMappingSignature(rawMapping) + } + + a.headerOverrideCache = overrides + a.headerOverrideCacheReady = true + a.headerOverrideCacheCredentialsPtr = credentialsPtr + a.headerOverrideCacheRawPtr = rawPtr + a.headerOverrideCacheRawLen = rawLen + a.headerOverrideCacheRawSig = rawSig + return overrides +} + +// resolveHeaderOverrides 解析并防御性过滤原始覆写表：保存路径已做校验， +// 这里兜底未经 Normalize 落库的数据（含名单扩充前保存的旧配置），非法条目直接跳过。 +func resolveHeaderOverrides(raw map[string]string) map[string]string { + if len(raw) == 0 { + return nil + } + result := make(map[string]string, len(raw)) + for name, value := range raw { + lowerName, value, err := normalizeHeaderOverrideEntry(name, value) + if err != nil || lowerName == "" || value == "" { + continue + } + result[lowerName] = value + } + if len(result) == 0 { + return nil + } + return result +} + +// HeaderOverrideValue 返回指定 header（小写名）的生效覆写值。 +// 供转发链路在 header 写入前感知覆写结果（如 anthropic-beta 需要参与 body 净化）。 +func (a *Account) HeaderOverrideValue(lowerName string) (string, bool) { + value, ok := a.GetHeaderOverrides()[lowerName] + return value, ok +} + +// ApplyHeaderOverrides 将账号配置的请求头覆写应用到出站请求头。 +// 对每个覆写条目：先删除所有大小写变体（转发链路会以 wire casing 直接写入 map， +// 可能存在非 canonical key），再按已知 wire casing 写入，避免产生重复头。 +// 账号未启用或不符合条件时为 no-op，可安全地在 OAuth/api_key 共用的构建器中调用。 +func (a *Account) ApplyHeaderOverrides(h http.Header) { + if h == nil { + return + } + overrides := a.GetHeaderOverrides() + if len(overrides) == 0 { + return + } + // 覆写名两两不同（大小写不敏感）且各自只操作同名键，应用顺序不影响结果。 + // 全量 EqualFold 扫描兜底删除任意 casing 的既有键：透传链路可能保留客户端 + // 原始 casing，非 canonical/wire casing 的键 deleteHeaderAllForms 覆盖不到。 + for name, value := range overrides { + for existing := range h { + if strings.EqualFold(existing, name) { + delete(h, existing) + } + } + h[resolveWireCasing(name)] = []string{value} + } +} + +// NormalizeHeaderOverrideCredentials 校验并原地规范化 credentials 中的请求头覆写字段。 +// 供账号创建/更新/批量更新的保存路径调用；credentials 未携带相关字段时为 no-op。 +// 规范化内容：header 名转小写并去除首尾空白，value 去除首尾空白，丢弃名和值均为空的条目。 +func NormalizeHeaderOverrideCredentials(credentials map[string]any) error { + if credentials == nil { + return nil + } + if raw, ok := credentials[credKeyHeaderOverrideEnabled]; ok && raw != nil { + if _, isBool := raw.(bool); !isBool { + return infraerrors.New(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header_override_enabled must be a boolean") + } + } + raw, ok := credentials[credKeyHeaderOverrides] + if !ok || raw == nil { + return nil + } + + var entries map[string]any + switch m := raw.(type) { + case map[string]any: + entries = m + case map[string]string: + entries = make(map[string]any, len(m)) + for k, v := range m { + entries[k] = v + } + default: + return infraerrors.New(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header_overrides must be an object of header name to string value") + } + + if len(entries) > maxHeaderOverrideEntries { + return infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header_overrides supports at most %d entries", maxHeaderOverrideEntries) + } + + normalized := make(map[string]any, len(entries)) + for name, rawValue := range entries { + value, isString := rawValue.(string) + if !isString { + return infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header %q value must be a string", name) + } + lowerName, value, err := normalizeHeaderOverrideEntry(name, value) + if err != nil { + return err + } + if lowerName == "" { + continue // 丢弃完全为空的占位行 + } + if _, dup := normalized[lowerName]; dup { + return infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "duplicate header name %q (matching is case-insensitive)", lowerName) + } + normalized[lowerName] = value + } + credentials[credKeyHeaderOverrides] = normalized + return nil +} + +// normalizeHeaderOverrideEntry 校验并规范化单个覆写条目，保存路径（Normalize，err → 400） +// 与应用路径（resolveHeaderOverrides，err → 跳过）共用同一套规则，避免两处校验漂移。 +// 名和值均为空表示空占位行，返回 ("", "", nil)；空 value 的具名条目合法（模板占位）。 +func normalizeHeaderOverrideEntry(name, value string) (string, string, error) { + lowerName := strings.ToLower(strings.TrimSpace(name)) + value = strings.TrimSpace(value) + if lowerName == "" { + if value == "" { + return "", "", nil + } + return "", "", infraerrors.New(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header name must not be empty") + } + if len(lowerName) > maxHeaderOverrideNameLength { + return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header name %q exceeds %d characters", lowerName, maxHeaderOverrideNameLength) + } + if !httpguts.ValidHeaderFieldName(lowerName) { + return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "invalid header name %q", lowerName) + } + if isHeaderOverrideBlockedName(lowerName) { + return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header %q is not allowed to be overridden", lowerName) + } + if len(value) > maxHeaderOverrideValueLength { + return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header %q value exceeds %d characters", lowerName, maxHeaderOverrideValueLength) + } + if !httpguts.ValidHeaderFieldValue(value) { + return "", "", infraerrors.Newf(http.StatusBadRequest, "INVALID_HEADER_OVERRIDE", + "header %q has an invalid value", lowerName) + } + return lowerName, value, nil +} diff --git a/backend/internal/service/account_header_override_test.go b/backend/internal/service/account_header_override_test.go new file mode 100644 index 0000000000..c89b5e0587 --- /dev/null +++ b/backend/internal/service/account_header_override_test.go @@ -0,0 +1,339 @@ +//go:build unit + +package service + +import ( + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func headerOverrideTestAccount(platform, accountType string, credentials map[string]any) *Account { + return &Account{ + Platform: platform, + Type: accountType, + Credentials: credentials, + } +} + +func TestIsHeaderOverrideEligible(t *testing.T) { + tests := []struct { + name string + platform string + accType string + want bool + }{ + {"anthropic apikey", PlatformAnthropic, AccountTypeAPIKey, true}, + {"openai apikey", PlatformOpenAI, AccountTypeAPIKey, true}, + {"anthropic oauth", PlatformAnthropic, AccountTypeOAuth, false}, + {"openai oauth", PlatformOpenAI, AccountTypeOAuth, false}, + {"gemini apikey", PlatformGemini, AccountTypeAPIKey, false}, + {"grok apikey", PlatformGrok, AccountTypeAPIKey, false}, + {"antigravity apikey", PlatformAntigravity, AccountTypeAPIKey, false}, + {"anthropic bedrock", PlatformAnthropic, AccountTypeBedrock, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + acc := headerOverrideTestAccount(tt.platform, tt.accType, nil) + require.Equal(t, tt.want, acc.IsHeaderOverrideEligible()) + }) + } + + var nilAccount *Account + require.False(t, nilAccount.IsHeaderOverrideEligible()) + require.False(t, nilAccount.IsHeaderOverrideEnabled()) + require.Nil(t, nilAccount.GetHeaderOverrides()) +} + +func TestIsHeaderOverrideEnabled(t *testing.T) { + acc := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + }) + require.True(t, acc.IsHeaderOverrideEnabled()) + + // 未配置 / 非 bool / false 均视为未启用 + require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, nil).IsHeaderOverrideEnabled()) + require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: "true", + }).IsHeaderOverrideEnabled()) + require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: false, + }).IsHeaderOverrideEnabled()) + + // 不符合平台/类型条件时即使配置了 true 也不启用 + require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeOAuth, map[string]any{ + credKeyHeaderOverrideEnabled: true, + }).IsHeaderOverrideEnabled()) + require.False(t, headerOverrideTestAccount(PlatformGemini, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + }).IsHeaderOverrideEnabled()) +} + +func TestGetHeaderOverrides(t *testing.T) { + acc := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{ + "User-Agent": "my-agent/1.0", // 大写 key 归一化为小写 + " X-App ": "cli", // 名称去空白 + "x-empty": "", // 空 value（模板占位）跳过 + "authorization": "Bearer leaked", // 禁止覆写的头跳过 + "bad name": "value", // 非法 header 名跳过 + "x-padded": " padded ", // value 去空白 + }, + }) + overrides := acc.GetHeaderOverrides() + require.Equal(t, map[string]string{ + "user-agent": "my-agent/1.0", + "x-app": "cli", + "x-padded": "padded", + }, overrides) + + // 未启用时返回 nil + disabled := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrides: map[string]any{"user-agent": "x"}, + }) + require.Nil(t, disabled.GetHeaderOverrides()) + + // 启用但全部为空 value 时返回 nil + empty := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{"user-agent": ""}, + }) + require.Nil(t, empty.GetHeaderOverrides()) + + // 未经 Normalize 落库的超长数据 / WebSocket 握手头在应用时被防御性跳过 + oversizedValue := strings.Repeat("a", maxHeaderOverrideValueLength+1) + defensive := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{ + "x-big": oversizedValue, + "sec-websocket-key": "forged", + "content-type": "application/json", // 名单扩充前落库的数据也要被拦截 + "x-claude-code-session-id": "pinned-session", + "x-ok": "ok", + }, + }) + require.Equal(t, map[string]string{"x-ok": "ok"}, defensive.GetHeaderOverrides()) +} + +func TestApplyHeaderOverrides(t *testing.T) { + acc := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{ + "user-agent": "override-agent/2.0", + "anthropic-beta": "custom-beta-1", + "x-custom": "custom-value", + }, + }) + + h := http.Header{} + // 模拟转发链路：canonical key 与 wire casing 原样 key 混合存在 + h.Set("User-Agent", "claude-cli/2.1.161 (external, cli)") + h["anthropic-beta"] = []string{"claude-code-20250219,oauth-2025-04-20"} // 非 canonical 原样 key + h.Set("Content-Type", "application/json") + + acc.ApplyHeaderOverrides(h) + + // user-agent 覆盖且只有一个值（已知头恢复 wire casing） + require.Equal(t, []string{"override-agent/2.0"}, h["User-Agent"]) + // anthropic-beta：非 canonical 旧值被清除，写入 wire casing（小写） + require.Equal(t, []string{"custom-beta-1"}, h["anthropic-beta"]) + require.Empty(t, h["Anthropic-Beta"]) + // 新增头（未知头以小写原样键写入，与转发链路 wire casing 约定一致） + require.Equal(t, []string{"custom-value"}, h["x-custom"]) + require.Equal(t, "custom-value", getHeaderRaw(h, "x-custom")) + // 未覆写的头不受影响 + require.Equal(t, "application/json", h.Get("Content-Type")) + + // 覆盖后不存在任何大小写重复 + count := 0 + for k := range h { + if k == "anthropic-beta" || k == "Anthropic-Beta" { + count++ + } + } + require.Equal(t, 1, count) +} + +func TestApplyHeaderOverridesNoOpPaths(t *testing.T) { + baseline := func() http.Header { + h := http.Header{} + h.Set("User-Agent", "orig") + return h + } + + // OAuth 账号：即使配置了覆写也不生效 + oauth := headerOverrideTestAccount(PlatformAnthropic, AccountTypeOAuth, map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{"user-agent": "hacked"}, + }) + h := baseline() + oauth.ApplyHeaderOverrides(h) + require.Equal(t, "orig", h.Get("User-Agent")) + + // 未启用开关 + off := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrides: map[string]any{"user-agent": "hacked"}, + }) + h = baseline() + off.ApplyHeaderOverrides(h) + require.Equal(t, "orig", h.Get("User-Agent")) + + // 禁止覆写的头（authorization / x-api-key / host 等）不会被应用 + blocked := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{ + "Authorization": "Bearer evil", + "X-Api-Key": "evil", + "Host": "evil.example.com", + "Content-Length": "0", + }, + }) + h = http.Header{} + h.Set("Authorization", "Bearer real-key") + blocked.ApplyHeaderOverrides(h) + require.Equal(t, "Bearer real-key", h.Get("Authorization")) + require.Empty(t, h.Get("X-Api-Key")) + require.Empty(t, h.Get("Host")) + + // nil header 不 panic + blocked.ApplyHeaderOverrides(nil) +} + +func TestNormalizeHeaderOverrideCredentials(t *testing.T) { + t.Run("nil credentials no-op", func(t *testing.T) { + require.NoError(t, NormalizeHeaderOverrideCredentials(nil)) + }) + + t.Run("missing keys no-op", func(t *testing.T) { + creds := map[string]any{"api_key": "sk-xxx"} + require.NoError(t, NormalizeHeaderOverrideCredentials(creds)) + _, exists := creds[credKeyHeaderOverrides] + require.False(t, exists) + }) + + t.Run("normalizes names and values", func(t *testing.T) { + creds := map[string]any{ + credKeyHeaderOverrideEnabled: true, + credKeyHeaderOverrides: map[string]any{ + " User-Agent ": " my-agent ", + "X-App": "", + "": "", // 完全空行被丢弃 + }, + } + require.NoError(t, NormalizeHeaderOverrideCredentials(creds)) + require.Equal(t, map[string]any{ + "user-agent": "my-agent", + "x-app": "", + }, creds[credKeyHeaderOverrides]) + }) + + t.Run("accepts map[string]string input", func(t *testing.T) { + creds := map[string]any{ + credKeyHeaderOverrides: map[string]string{"X-App": "cli"}, + } + require.NoError(t, NormalizeHeaderOverrideCredentials(creds)) + require.Equal(t, map[string]any{"x-app": "cli"}, creds[credKeyHeaderOverrides]) + }) + + t.Run("rejects non-bool enabled", func(t *testing.T) { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrideEnabled: "yes", + }) + require.Error(t, err) + }) + + t.Run("rejects non-object overrides", func(t *testing.T) { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: []any{"user-agent"}, + }) + require.Error(t, err) + }) + + t.Run("rejects non-string value", func(t *testing.T) { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{"x-app": 123}, + }) + require.Error(t, err) + }) + + t.Run("rejects invalid header name", func(t *testing.T) { + for _, name := range []string{"bad name", "bad:name", "bad\nname", "值"} { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{name: "v"}, + }) + require.Error(t, err, "name %q should be rejected", name) + } + }) + + t.Run("rejects empty name with value", func(t *testing.T) { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{" ": "v"}, + }) + require.Error(t, err) + }) + + t.Run("rejects blocked headers", func(t *testing.T) { + for _, name := range []string{ + "Authorization", "x-api-key", "Host", "content-length", "Transfer-Encoding", + "connection", "accept-encoding", "Sec-WebSocket-Key", "session_id", + "conversation_id", "x-codex-turn-state", "chatgpt-account-id", + "Content-Type", "Cookie", "x-goog-api-key", + "X-Claude-Code-Session-Id", "x-client-request-id", + } { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{name: "v"}, + }) + require.Error(t, err, "blocked header %q should be rejected", name) + } + }) + + t.Run("allows tab inside value", func(t *testing.T) { + creds := map[string]any{ + credKeyHeaderOverrides: map[string]any{"x-app": "a\tb"}, + } + require.NoError(t, NormalizeHeaderOverrideCredentials(creds)) + require.Equal(t, map[string]any{"x-app": "a\tb"}, creds[credKeyHeaderOverrides]) + }) + + t.Run("rejects invalid value", func(t *testing.T) { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{"x-app": "bad\nvalue"}, + }) + require.Error(t, err) + }) + + t.Run("rejects duplicate names case-insensitively", func(t *testing.T) { + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{ + "User-Agent": "a", + "user-agent": "b", + }, + }) + require.Error(t, err) + }) + + t.Run("rejects too many entries", func(t *testing.T) { + entries := make(map[string]any, maxHeaderOverrideEntries+1) + for i := 0; i <= maxHeaderOverrideEntries; i++ { + entries["x-h-"+string(rune('a'+i%26))+string(rune('a'+(i/26)%26))+string(rune('a'+(i/676)%26))] = "v" + } + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: entries, + }) + require.Error(t, err) + }) + + t.Run("rejects oversized value", func(t *testing.T) { + big := make([]byte, maxHeaderOverrideValueLength+1) + for i := range big { + big[i] = 'a' + } + err := NormalizeHeaderOverrideCredentials(map[string]any{ + credKeyHeaderOverrides: map[string]any{"x-app": string(big)}, + }) + require.Error(t, err) + }) +} diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index dcba614c2c..5956684f98 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -39,6 +39,9 @@ type AccountRepository interface { List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) + // ListAllWithFilters 返回符合过滤条件的全部账号（不分页），用于账号列表页 + // 计算 OpenAI 调度分数的过滤范围池。 + ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) ListActive(ctx context.Context) ([]Account, error) ListOAuthRefreshCandidates(ctx context.Context) ([]Account, error) diff --git a/backend/internal/service/account_service_delete_test.go b/backend/internal/service/account_service_delete_test.go index a304356c09..ee6163239e 100644 --- a/backend/internal/service/account_service_delete_test.go +++ b/backend/internal/service/account_service_delete_test.go @@ -79,6 +79,10 @@ func (s *accountRepoStub) List(ctx context.Context, params pagination.Pagination panic("unexpected List call") } +func (s *accountRepoStub) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) { + return nil, nil +} + func (s *accountRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) { panic("unexpected ListWithFilters call") } diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 80a862b971..7bf02afd51 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -295,6 +295,9 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account setAnthropicAPIKeyAuthHeader(req.Header, account, authToken) } + // 账号级请求头覆写：测试请求与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) + // Get proxy URL proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -600,9 +603,19 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account if isOAuth { req.Host = "chatgpt.com" req.Header.Set("accept", "text/event-stream") + req.Header.Set("OpenAI-Beta", "responses=experimental") + req.Header.Set("Originator", "codex_cli_rs") + if customUA := strings.TrimSpace(credentialAccount.GetOpenAIUserAgent()); customUA != "" { + req.Header.Set("User-Agent", customUA) + } else { + req.Header.Set("User-Agent", codexCLIUserAgent) + } setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount) } + // 账号级请求头覆写：测试请求与真实转发保持一致的最终头 + credentialAccount.ApplyHeaderOverrides(req.Header) + // Get proxy URL proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -756,6 +769,9 @@ func (s *AccountTestService) testOpenAIChatCompletionsConnection( req.Header.Set("Accept", "text/event-stream") req.Header.Set("Authorization", "Bearer "+authToken) + // 账号级请求头覆写：测试请求与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) + proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() @@ -848,6 +864,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account setOpenAIChatGPTAccountHeaders(req.Header, account) } + // 账号级请求头覆写：测试请求与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) + proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() @@ -1599,6 +1618,9 @@ func (s *AccountTestService) testOpenAIImageAPIKey(c *gin.Context, ctx context.C req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+authToken) + // 账号级请求头覆写：测试请求与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) + proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index 7685174f3b..8c67ae987e 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -78,6 +78,12 @@ type AdminService interface { // Account management ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string, sortBy, sortOrder string) ([]Account, int64, error) + // ListAccountsForSchedulerScoreFilter 返回符合过滤条件的全部账号（不分页）， + // 作为账号列表页计算 OpenAI 调度分数的过滤范围池。 + ListAccountsForSchedulerScoreFilter(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) + // ListOpenAISchedulableAccountsForSchedulerScore 返回指定分组（nil 为未分组）内 + // 可调度的 OpenAI 账号，用于按组计算调度分数。 + ListOpenAISchedulableAccountsForSchedulerScore(ctx context.Context, groupID *int64) ([]Account, error) GetAccount(ctx context.Context, id int64) (*Account, error) GetAccountsByIDs(ctx context.Context, ids []int64) ([]*Account, error) CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error) @@ -209,9 +215,12 @@ 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 // 高峰时段倍率配置（PeakRateMultiplier 为 nil 时按 1.0 处理） PeakRateEnabled bool PeakStart string @@ -255,9 +264,12 @@ 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 // 高峰时段倍率配置（nil 表示不修改） PeakRateEnabled *bool PeakStart *string @@ -1851,6 +1863,20 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn } 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 + } peakRateMultiplier := 1.0 if input.PeakRateMultiplier != nil { @@ -1886,6 +1912,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn } allowImageGeneration := input.AllowImageGeneration || defaultAllowImageGenerationForPlatform(platform) + allowBatchImageGeneration := input.AllowBatchImageGeneration && allowImageGeneration && platform == PlatformGemini // 如果指定了复制账号的源分组，先获取账号 ID 列表 var accountIDsToCopy []int64 @@ -1931,8 +1958,11 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn WeeklyLimitUSD: weeklyLimit, MonthlyLimitUSD: monthlyLimit, AllowImageGeneration: allowImageGeneration, + AllowBatchImageGeneration: allowBatchImageGeneration, ImageRateIndependent: input.ImageRateIndependent, ImageRateMultiplier: imageRateMultiplier, + BatchImageDiscountMultiplier: batchImageDiscountMultiplier, + BatchImageHoldMultiplier: batchImageHoldMultiplier, PeakRateEnabled: peakRateEnabled, PeakStart: peakStart, PeakEnd: peakEnd, @@ -2117,6 +2147,12 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd 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 } @@ -2126,6 +2162,18 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd } 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 + } if input.PeakRateEnabled != nil { group.PeakRateEnabled = *input.PeakRateEnabled } @@ -2618,6 +2666,23 @@ func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, 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) } @@ -2668,6 +2733,11 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou } } + // 校验并规范化请求头覆写配置（header 名小写化、格式检查） + if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { + return nil, err + } + account := &Account{ Name: input.Name, Notes: normalizeAccountNotes(input.Notes), @@ -2798,6 +2868,10 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U // 敏感子键采用"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:{}，此时也必须落库。 @@ -3016,6 +3090,11 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp } } + // 校验并规范化请求头覆写配置（批量路径为 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, diff --git a/backend/internal/service/admin_service_bulk_update_test.go b/backend/internal/service/admin_service_bulk_update_test.go index df415295b1..2f44b1741d 100644 --- a/backend/internal/service/admin_service_bulk_update_test.go +++ b/backend/internal/service/admin_service_bulk_update_test.go @@ -88,6 +88,10 @@ func (s *accountRepoStubForBulkUpdate) ListByGroup(_ context.Context, groupID in return nil, nil } +func (s *accountRepoStubForBulkUpdate) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) { + return nil, nil +} + func (s *accountRepoStubForBulkUpdate) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) { s.listCalled = true s.lastListParams = params diff --git a/backend/internal/service/admin_service_group_test.go b/backend/internal/service/admin_service_group_test.go index 0b360c61c7..0c229c389e 100644 --- a/backend/internal/service/admin_service_group_test.go +++ b/backend/internal/service/admin_service_group_test.go @@ -232,6 +232,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{ @@ -326,6 +366,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 +471,58 @@ 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.6 + + 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.6, repo.created.BatchImageHoldMultiplier, 1e-12) +} + +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_InvalidatesAuthCacheOnRPMLimitChange(t *testing.T) { existingGroup := &Group{ ID: 1, diff --git a/backend/internal/service/admin_service_search_test.go b/backend/internal/service/admin_service_search_test.go index 595e99e344..76acd1b5e2 100644 --- a/backend/internal/service/admin_service_search_test.go +++ b/backend/internal/service/admin_service_search_test.go @@ -25,6 +25,10 @@ type accountRepoStubForAdminList struct { listWithFiltersErr error } +func (s *accountRepoStubForAdminList) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) { + return nil, nil +} + func (s *accountRepoStubForAdminList) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) { s.listWithFiltersCalls++ s.listWithFiltersParams = params diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index aa4cab22d7..585170438e 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -153,14 +153,24 @@ type antigravityRetryLoopResult struct { } // resolveAntigravityForwardBaseURL 解析转发用 base URL。 -// 默认使用 daily（ForwardBaseURLs 的首个地址）；当环境变量为 prod 时使用第二个地址。 +// +// 默认使用生产端点 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.ForwardBaseURLs() + baseURLs := antigravity.BaseURLs if len(baseURLs) == 0 { return "" } mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv))) - if mode == "prod" && len(baseURLs) > 1 { + if (mode == "daily" || mode == "sandbox") && len(baseURLs) > 1 { return baseURLs[1] } return baseURLs[0] diff --git a/backend/internal/service/antigravity_token_refresher.go b/backend/internal/service/antigravity_token_refresher.go index 7ce0ccf0fe..f40df7d95f 100644 --- a/backend/internal/service/antigravity_token_refresher.go +++ b/backend/internal/service/antigravity_token_refresher.go @@ -12,6 +12,10 @@ const ( // antigravityRefreshWindow Antigravity token 提前刷新窗口：15分钟 // Google OAuth token 有效期55分钟，提前15分钟刷新 antigravityRefreshWindow = 15 * time.Minute + + antigravityForceTokenRefreshExtraKey = "antigravity_force_token_refresh" + antigravityForceTokenRefreshReasonExtraKey = "antigravity_force_token_refresh_reason" + antigravityForceTokenRefreshAtExtraKey = "antigravity_force_token_refresh_at" ) // AntigravityTokenRefresher 实现 TokenRefresher 接口 @@ -41,6 +45,9 @@ func (r *AntigravityTokenRefresher) NeedsRefresh(account *Account, _ time.Durati if !r.CanRefresh(account) { return false } + if accountNeedsAntigravityForceTokenRefresh(account) { + return true + } expiresAt := account.GetCredentialAsTime("expires_at") if expiresAt == nil { return false @@ -54,6 +61,29 @@ func (r *AntigravityTokenRefresher) NeedsRefresh(account *Account, _ time.Durati return needsRefresh } +func accountNeedsAntigravityForceTokenRefresh(account *Account) bool { + return account != nil && + account.Platform == PlatformAntigravity && + account.Type == AccountTypeOAuth && + account.getExtraBool(antigravityForceTokenRefreshExtraKey) +} + +func antigravityForceTokenRefreshExtra(reason string) map[string]any { + return map[string]any{ + antigravityForceTokenRefreshExtraKey: true, + antigravityForceTokenRefreshReasonExtraKey: reason, + antigravityForceTokenRefreshAtExtraKey: time.Now().UTC().Format(time.RFC3339), + } +} + +func clearAntigravityForceTokenRefreshExtra() map[string]any { + return map[string]any{ + antigravityForceTokenRefreshExtraKey: false, + antigravityForceTokenRefreshReasonExtraKey: "", + antigravityForceTokenRefreshAtExtraKey: "", + } +} + // Refresh 执行 token 刷新 func (r *AntigravityTokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) { tokenInfo, err := r.antigravityOAuthService.RefreshAccountToken(ctx, account) diff --git a/backend/internal/service/api_key.go b/backend/internal/service/api_key.go index ec20b0a9bf..dfc3ec1c5a 100644 --- a/backend/internal/service/api_key.go +++ b/backend/internal/service/api_key.go @@ -44,6 +44,7 @@ type APIKey struct { UpdatedAt time.Time User *User Group *Group + CurrentConcurrency int // Quota fields Quota float64 // Quota limit in USD (0 = unlimited) diff --git a/backend/internal/service/api_key_auth_cache.go b/backend/internal/service/api_key_auth_cache.go index 32c3910c9d..6f927ff3b8 100644 --- a/backend/internal/service/api_key_auth_cache.go +++ b/backend/internal/service/api_key_auth_cache.go @@ -67,6 +67,7 @@ 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"` diff --git a/backend/internal/service/api_key_auth_cache_impl.go b/backend/internal/service/api_key_auth_cache_impl.go index b5aedf271e..f3da3df493 100644 --- a/backend/internal/service/api_key_auth_cache_impl.go +++ b/backend/internal/service/api_key_auth_cache_impl.go @@ -259,6 +259,7 @@ 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, @@ -336,6 +337,7 @@ 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, diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index de9b908dc7..8903be65ee 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -203,6 +203,7 @@ type APIKeyService struct { userGroupRateRepo UserGroupRateRepository cache APIKeyCache rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache + concurrencyService *ConcurrencyService cfg *config.Config authCacheL1 *ristretto.Cache authCfg apiKeyAuthCacheConfig @@ -240,6 +241,10 @@ func (s *APIKeyService) SetRateLimitCacheInvalidator(inv RateLimitCacheInvalidat s.rateLimitCacheInvalid = inv } +func (s *APIKeyService) SetConcurrencyService(concurrencyService *ConcurrencyService) { + s.concurrencyService = concurrencyService +} + func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) { if apiKey == nil { return @@ -436,9 +441,40 @@ func (s *APIKeyService) List(ctx context.Context, userID int64, params paginatio if err != nil { return nil, nil, fmt.Errorf("list api keys: %w", err) } + s.fillCurrentConcurrency(ctx, keys) return keys, pagination, nil } +func (s *APIKeyService) fillCurrentConcurrency(ctx context.Context, keys []APIKey) { + if s == nil || s.concurrencyService == nil || len(keys) == 0 { + return + } + ids := make([]int64, 0, len(keys)) + for i := range keys { + if keys[i].ID > 0 { + ids = append(ids, keys[i].ID) + } + } + counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, ids) + if err != nil { + return + } + for i := range keys { + keys[i].CurrentConcurrency = counts[keys[i].ID] + } +} + +func (s *APIKeyService) currentConcurrencyForAPIKey(ctx context.Context, apiKeyID int64) int { + if s == nil || s.concurrencyService == nil || apiKeyID <= 0 { + return 0 + } + counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, []int64{apiKeyID}) + if err != nil { + return 0 + } + return counts[apiKeyID] +} + func (s *APIKeyService) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) { if len(apiKeyIDs) == 0 { return []int64{}, nil @@ -458,6 +494,9 @@ func (s *APIKeyService) GetByID(ctx context.Context, id int64) (*APIKey, error) return nil, fmt.Errorf("get api key: %w", err) } s.compileAPIKeyIPRules(apiKey) + if apiKey != nil { + apiKey.CurrentConcurrency = s.currentConcurrencyForAPIKey(ctx, apiKey.ID) + } return apiKey, nil } diff --git a/backend/internal/service/api_key_service_delete_test.go b/backend/internal/service/api_key_service_delete_test.go index 8664c03bd7..25ad1edb15 100644 --- a/backend/internal/service/api_key_service_delete_test.go +++ b/backend/internal/service/api_key_service_delete_test.go @@ -300,6 +300,40 @@ func TestApiKeyService_Delete_NotFound(t *testing.T) { require.Empty(t, cache.deleteAuthKeys) } +func TestAPIKeyService_List_FillsCurrentConcurrency(t *testing.T) { + repo := &apiKeyRepoStub{ + allowListByUserID: true, + listByUserIDKeys: []APIKey{ + {ID: 10, UserID: 7, Key: "sk-10", Name: "key-10"}, + {ID: 11, UserID: 7, Key: "sk-11", Name: "key-11"}, + }, + } + concurrency := NewConcurrencyService(&stubConcurrencyCacheForTest{ + apiKeyConcurrency: map[int64]int{10: 2, 11: 0}, + }) + svc := &APIKeyService{apiKeyRepo: repo, concurrencyService: concurrency} + + keys, _, err := svc.List(context.Background(), 7, pagination.PaginationParams{Page: 1, PageSize: 20}, APIKeyListFilters{}) + require.NoError(t, err) + require.Len(t, keys, 2) + require.Equal(t, 2, keys[0].CurrentConcurrency) + require.Equal(t, 0, keys[1].CurrentConcurrency) +} + +func TestAPIKeyService_GetByID_FillsCurrentConcurrency(t *testing.T) { + repo := &apiKeyRepoStub{ + apiKey: &APIKey{ID: 10, UserID: 7, Key: "sk-10", Name: "key-10"}, + } + concurrency := NewConcurrencyService(&stubConcurrencyCacheForTest{ + apiKeyConcurrency: map[int64]int{10: 4}, + }) + svc := &APIKeyService{apiKeyRepo: repo, concurrencyService: concurrency} + + key, err := svc.GetByID(context.Background(), 10) + require.NoError(t, err) + require.Equal(t, 4, key.CurrentConcurrency) +} + // TestApiKeyService_Delete_DeleteFails 测试删除操作失败时的错误处理。 // 预期行为： // - GetKeyAndOwnerID 返回正确的所有者 ID diff --git a/backend/internal/service/batch_image.go b/backend/internal/service/batch_image.go new file mode 100644 index 0000000000..4fad11dd5e --- /dev/null +++ b/backend/internal/service/batch_image.go @@ -0,0 +1,406 @@ +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") + + 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 + 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..af61117db0 --- /dev/null +++ b/backend/internal/service/batch_image_billing_hold.go @@ -0,0 +1,104 @@ +package service + +import ( + "context" + "errors" + "strings" +) + +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 { + 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..d89f508f47 --- /dev/null +++ b/backend/internal/service/batch_image_billing_recovery.go @@ -0,0 +1,62 @@ +package service + +import ( + "context" + "errors" + "time" +) + +const ( + defaultBatchImageBillingRecoveryStaleAfter = 10 * time.Minute + defaultBatchImageBillingRecoveryLimit = 100 +) + +type BatchImageBillingRecoveryService struct { + Repo BatchImageRepository + Billing UsageBillingRepository + AuthCache APIKeyAuthCacheInvalidator + 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 + } + jobs, err := s.Repo.ListStaleUnsubmittedBatchImageJobs(ctx, time.Now().Add(-staleAfter), limit) + if err != nil { + return 0, err + } + released := 0 + for _, job := range jobs { + if job == nil { + continue + } + msg := "batch image submission did not reach provider before recovery cutoff" + if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{ + EventType: "billing_hold_recovery_failed_unsubmitted", + EventPayload: map[string]any{"batch_id": job.BatchID}, + ErrorCode: batchImageStringPtr("SUBMIT_STALE_BEFORE_PROVIDER"), + ErrorMessage: batchImageStringPtr(msg), + }); err != nil && !errors.Is(err, ErrBatchImageInvalidTransition) { + return released, err + } + job.Status = BatchImageJobStatusFailed + if err := releaseBatchImageBalanceHold(ctx, s.Billing, job, batchImageDerefString(job.RequestHash)); err != nil { + return released, err + } + if s.AuthCache != nil && job.UserID > 0 { + s.AuthCache.InvalidateAuthCacheByUserID(ctx, job.UserID) + } + released++ + } + return released, nil +} 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..2ab2783f83 --- /dev/null +++ b/backend/internal/service/batch_image_billing_recovery_test.go @@ -0,0 +1,52 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +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) +} diff --git a/backend/internal/service/batch_image_cleanup.go b/backend/internal/service/batch_image_cleanup.go new file mode 100644 index 0000000000..7b2b527080 --- /dev/null +++ b/backend/internal/service/batch_image_cleanup.go @@ -0,0 +1,294 @@ +package service + +import ( + "context" + "errors" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +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, + } +} + +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.Repo.AppendBatchImageEvent(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.Repo.AppendBatchImageEvent(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.Repo.AppendBatchImageEvent(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()) + _ = s.Repo.RecordBatchImageCleanupFailure(ctx, job.BatchID, code, msg) + event := string(target) + "_cleanup_failed" + _ = s.Repo.AppendBatchImageEvent(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..b08b99c75f --- /dev/null +++ b/backend/internal/service/batch_image_download.go @@ -0,0 +1,659 @@ +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 maxItems <= 0 { + maxItems = s.maxZipItems() + } + 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..24496f6bfa --- /dev/null +++ b/backend/internal/service/batch_image_processor.go @@ -0,0 +1,596 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "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 + } + 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 + } + 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) + 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 + } + 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 + } + if result.TotalCount == 0 { + return nil, ErrBatchImageIndexNoResultLines + } + if err := i.Repo.ReplaceBatchImageItemsForJob(ctx, job.BatchID, items, BatchImageCounts{ + SuccessCount: result.SuccessCount, + FailCount: result.FailCount, + }); err != nil { + return nil, err + } + return result, nil +} + +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..8a37126161 --- /dev/null +++ b/backend/internal/service/batch_image_processor_test.go @@ -0,0 +1,847 @@ +//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) + + 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, 1, result.TotalCount) + require.Len(t, repo.items[job.BatchID], 1) + require.Equal(t, "ok2", repo.items[job.BatchID][0].CustomID) +} + +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) 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 { + 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..ccf91de714 --- /dev/null +++ b/backend/internal/service/batch_image_provider_gemini.go @@ -0,0 +1,680 @@ +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" +) + +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 = http.DefaultClient + } + return &GeminiBatchHTTPClient{baseURL: baseURL, client: 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" + } + 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 +} + +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..9289c1484c --- /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 = http.DefaultClient + } + 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 = http.DefaultClient + } + 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..0b792aacb2 --- /dev/null +++ b/backend/internal/service/batch_image_public.go @@ -0,0 +1,1306 @@ +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" +) + +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 + } + 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, + }) + } + + providerJob, err := provider.Submit(ctx, job, account, input) + 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 { + 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 +} + +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) { + _ = s.Repo.AppendBatchImageEvent(ctx, batchID, "billing_retry_enqueue_failed", map[string]any{ + "batch_id": batchID, + "error": sanitizeBatchImagePublicMessage(err.Error()), + }) + } +} + +func (s *BatchImagePublicService) hidePreUpstreamSubmitFailure(ctx context.Context, owner BatchImageOwner, job *BatchImageJob) { + if s == nil || s.Repo == nil || job == nil || job.ProviderJobName != nil { + return + } + _ = s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now()) +} + +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 + } + _ = s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "job_cancel_requested", map[string]any{"batch_id": job.BatchID}) + 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 + } + 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..d59784269f --- /dev/null +++ b/backend/internal/service/batch_image_public_test.go @@ -0,0 +1,961 @@ +//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) + require.InDelta(t, 0.6, job.HoldMultiplier, 1e-12) + require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12) + require.InDelta(t, 0.09375, job.HoldUnitPrice, 1e-12) + require.InDelta(t, 0.1875, *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..cbd1ca7ae6 --- /dev/null +++ b/backend/internal/service/batch_image_settlement.go @@ -0,0 +1,332 @@ +package service + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +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.SuccessCount < 0 || job.FailCount < 0 || job.ItemCount < 0 || job.SuccessCount+job.FailCount > job.ItemCount { + return nil, ErrBatchImageSettlementInvalidCounts + } + if strings.TrimSpace(batchImageDerefString(job.ManifestHash)) != "" && batchImageDerefString(job.ManifestHash) != manifestHash { + return nil, ErrBatchImageSettlementManifestConflict + } + if job.APIKeyID == nil || *job.APIKeyID <= 0 { + return nil, ErrBatchImageSettlementMissingAPIKeyID + } + if job.AccountID == nil || *job.AccountID <= 0 { + return nil, ErrBatchImageSettlementMissingAccountID + } + if isBatchImageSettlementRetryExhausted(job) { + return nil, s.failExhaustedSettlement(ctx, job, manifestHash, "settlement billing retry limit reached") + } + + unitPrice, err := s.settlementUnitPrice(ctx, job) + if err != nil { + return nil, err + } + if unitPrice < 0 { + return nil, ErrBatchImageSettlementPricingMissing + } + 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) + _, _ = s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_COST_EXCEEDS_HOLD", msg) + return nil, ErrBatchImageSettlementCostExceedsHold + } + + if err := captureBatchImageBalanceHold(ctx, s.BillingRepo, job, actualCost, manifestHash); err != nil { + msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength) + retryCount, recordErr := s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_BILLING_FAILED", msg) + if recordErr == nil && retryCount >= batchImageSettlementMaxRetries { + job.RetryCount = retryCount + return nil, s.failExhaustedSettlement(ctx, job, manifestHash, msg) + } + 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 +} + +func isBatchImageSettlementRetryExhausted(job *BatchImageJob) bool { + return job != nil && + job.Status == BatchImageJobStatusSettling && + job.RetryCount >= batchImageSettlementMaxRetries && + batchImageDerefString(job.LastErrorCode) == "SETTLEMENT_BILLING_FAILED" +} + +func (s *BatchImageSettlementService) failExhaustedSettlement(ctx context.Context, job *BatchImageJob, manifestHash, message string) error { + if s == nil || s.Repo == nil { + return ErrBatchImageSettlementBillingFailed + } + if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, manifestHash); err != nil { + msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength) + _, _ = s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_RELEASE_FAILED", msg) + 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..8837a85358 --- /dev/null +++ b/backend/internal/service/batch_image_settlement_test.go @@ -0,0 +1,440 @@ +//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 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..5027350689 --- /dev/null +++ b/backend/internal/service/batch_image_worker.go @@ -0,0 +1,230 @@ +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 { + return nil + } + defer func() { + _ = lock.Release(ctx) + }() + + result, err := w.processor.Process(ctx, reserved.BatchID) + 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) +} + +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..de3b4cb42a --- /dev/null +++ b/backend/internal/service/batch_image_worker_runtime.go @@ -0,0 +1,142 @@ +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, + 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..934cd9fd1c --- /dev/null +++ b/backend/internal/service/batch_image_worker_test.go @@ -0,0 +1,154 @@ +//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_SkipsWhenJobLockNotAcquired(t *testing.T) { + queue := newFakeBatchImageQueue("imgbatch_worker_locked") + queue.lockAcquired = false + processor := &fakeBatchImageProcessor{} + worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{}) + + require.NoError(t, worker.RunOnce(context.Background())) + require.Empty(t, processor.processed) + require.Empty(t, queue.requeued) + 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 a781936598..dc54a1b1f3 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -280,6 +280,11 @@ func (s *BillingService) initFallbackPricing() { s.fallbackPrices["gpt-5.5"] = s.fallbackPrices["gpt-5.4"] s.fallbackPrices["gpt-5.5-pro"] = s.fallbackPrices["gpt-5.4"] + // GPT-5.6（sol / terra / luna）暂无独立定价，回退到 GPT-5.4。 + s.fallbackPrices["gpt-5.6-sol"] = s.fallbackPrices["gpt-5.4"] + s.fallbackPrices["gpt-5.6-terra"] = s.fallbackPrices["gpt-5.4"] + s.fallbackPrices["gpt-5.6-luna"] = s.fallbackPrices["gpt-5.4"] + s.fallbackPrices["gpt-5.4-mini"] = &ModelPricing{ InputPricePerToken: 7.5e-7, OutputPricePerToken: 4.5e-6, @@ -667,6 +672,12 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing { // OpenAI（GPT-5 / Codex 族）：仅匹配已知型号，避免未知 OpenAI 型号误计价。 if normalized := normalizeKnownOpenAICodexModel(modelLower); normalized != "" { switch normalized { + case "gpt-5.6-sol": + return s.fallbackPrices["gpt-5.6-sol"] + case "gpt-5.6-terra": + return s.fallbackPrices["gpt-5.6-terra"] + case "gpt-5.6-luna": + return s.fallbackPrices["gpt-5.6-luna"] case "gpt-5.5-pro": return s.fallbackPrices["gpt-5.5-pro"] case "gpt-5.5": @@ -1060,7 +1071,8 @@ func isOpenAIGPT54Model(model string) bool { // normalizeCodexModel 的默认兜底把非 OpenAI 模型（claude-*、gemini-*、gpt-4o） // 误识别为 gpt-5.4。 normalized := normalizeKnownOpenAICodexModel(model) - return normalized == "gpt-5.4" || normalized == "gpt-5.5" || normalized == "gpt-5.5-pro" + return normalized == "gpt-5.4" || normalized == "gpt-5.5" || normalized == "gpt-5.5-pro" || + normalized == "gpt-5.6-sol" || normalized == "gpt-5.6-terra" || normalized == "gpt-5.6-luna" } // CalculateCostWithConfig 使用配置中的默认倍率计算费用 diff --git a/backend/internal/service/concurrency_service.go b/backend/internal/service/concurrency_service.go index 50e0077b94..f2f2aade89 100644 --- a/backend/internal/service/concurrency_service.go +++ b/backend/internal/service/concurrency_service.go @@ -37,7 +37,7 @@ type ConcurrencyCache interface { ReleaseUserSlot(ctx context.Context, userID int64, requestID string) error GetUserConcurrency(ctx context.Context, userID int64) (int, error) - // 等待队列计数（只在首次创建时设置 TTL） + // 等待队列计数（每次入队都会刷新 TTL，避免长时间排队时计数提前过期） IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) DecrementWaitCount(ctx context.Context, userID int64) error @@ -53,6 +53,12 @@ type ConcurrencyCache interface { CleanupStaleProcessSlots(ctx context.Context, activeRequestPrefix string) error } +type APIKeyConcurrencyCache interface { + TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) +} + var ( requestIDPrefix = initRequestIDPrefix() requestIDCounter atomic.Uint64 @@ -90,6 +96,8 @@ const ( defaultAccountLoadBatchCacheTTL = 200 * time.Millisecond accountLoadBatchFetchTimeout = 3 * time.Second maxAccountLoadBatchCacheEntries = 256 + apiKeyConcurrencyFetchTimeout = 3 * time.Second + apiKeySlotTrackTimeout = 2 * time.Second ) // ConcurrencyService 管理账号和用户的并发限制。 @@ -238,6 +246,77 @@ func (s *ConcurrencyService) AcquireUserSlot(ctx context.Context, userID int64, }, nil } +// TrackAPIKeySlot records one active request slot for an API key without +// applying key-level concurrency limits. It is fail-open: Redis errors are +// logged and return a no-op release function. +func (s *ConcurrencyService) TrackAPIKeySlot(ctx context.Context, apiKeyID int64) func() { + if s == nil || s.cache == nil || apiKeyID <= 0 { + return func() {} + } + cache, ok := s.cache.(APIKeyConcurrencyCache) + if !ok { + return func() {} + } + + requestID := generateRequestID() + baseCtx := context.Background() + if ctx != nil { + baseCtx = context.WithoutCancel(ctx) + } + trackCtx, cancel := context.WithTimeout(baseCtx, apiKeySlotTrackTimeout) + err := cache.TrackAPIKeySlot(trackCtx, apiKeyID, requestID) + cancel() + if err != nil { + logger.LegacyPrintf("service.concurrency", "Warning: failed to track api key slot for %d (req=%s): %v", apiKeyID, requestID, err) + return func() {} + } + + return func() { + bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := cache.ReleaseAPIKeySlot(bgCtx, apiKeyID, requestID); err != nil { + logger.LegacyPrintf("service.concurrency", "Warning: failed to release api key slot for %d (req=%s): %v", apiKeyID, requestID, err) + } + } +} + +// GetAPIKeyConcurrencyBatch gets real-time active request counts for API keys. +// Stats are best-effort: missing Redis support or Redis errors return zeroes. +func (s *ConcurrencyService) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) { + result := zeroAPIKeyConcurrencyMap(apiKeyIDs) + if len(apiKeyIDs) == 0 { + return result, nil + } + if s == nil || s.cache == nil { + return result, nil + } + cache, ok := s.cache.(APIKeyConcurrencyCache) + if !ok { + return result, nil + } + + redisCtx, cancel := context.WithTimeout(context.Background(), apiKeyConcurrencyFetchTimeout) + defer cancel() + + counts, err := cache.GetAPIKeyConcurrencyBatch(redisCtx, apiKeyIDs) + if err != nil { + logger.LegacyPrintf("service.concurrency", "Warning: get api key concurrency batch failed: %v", err) + return result, nil + } + for _, apiKeyID := range apiKeyIDs { + result[apiKeyID] = counts[apiKeyID] + } + return result, nil +} + +func zeroAPIKeyConcurrencyMap(apiKeyIDs []int64) map[int64]int { + result := make(map[int64]int, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + result[apiKeyID] = 0 + } + return result +} + // ============================================ // Wait Queue Count Methods // ============================================ diff --git a/backend/internal/service/concurrency_service_test.go b/backend/internal/service/concurrency_service_test.go index bacad0245e..3f358bbe6a 100644 --- a/backend/internal/service/concurrency_service_test.go +++ b/backend/internal/service/concurrency_service_test.go @@ -16,25 +16,33 @@ import ( // stubConcurrencyCacheForTest 用于并发服务单元测试的缓存桩 type stubConcurrencyCacheForTest struct { - acquireResult bool - acquireErr error - releaseErr error - concurrency int - concurrencyErr error - waitAllowed bool - waitErr error - waitCount int - waitCountErr error - loadBatch map[int64]*AccountLoadInfo - loadBatchErr error - usersLoadBatch map[int64]*UserLoadInfo - usersLoadErr error - cleanupErr error + acquireResult bool + acquireErr error + releaseErr error + concurrency int + concurrencyErr error + waitAllowed bool + waitErr error + waitCount int + waitCountErr error + loadBatch map[int64]*AccountLoadInfo + loadBatchErr error + usersLoadBatch map[int64]*UserLoadInfo + usersLoadErr error + cleanupErr error + apiKeyTrackErr error + apiKeyReleaseErr error + apiKeyConcurrency map[int64]int + apiKeyConcurrencyErr error // 记录调用 - releasedAccountIDs []int64 - releasedRequestIDs []string - loadBatchCalls atomic.Int64 + releasedAccountIDs []int64 + releasedRequestIDs []string + loadBatchCalls atomic.Int64 + trackedAPIKeyIDs []int64 + trackedAPIKeyRequestIDs []string + releasedAPIKeyIDs []int64 + releasedAPIKeyRequestIDs []string } var _ ConcurrencyCache = (*stubConcurrencyCacheForTest)(nil) @@ -78,6 +86,26 @@ func (c *stubConcurrencyCacheForTest) ReleaseUserSlot(_ context.Context, _ int64 func (c *stubConcurrencyCacheForTest) GetUserConcurrency(_ context.Context, _ int64) (int, error) { return c.concurrency, c.concurrencyErr } +func (c *stubConcurrencyCacheForTest) TrackAPIKeySlot(_ context.Context, apiKeyID int64, requestID string) error { + c.trackedAPIKeyIDs = append(c.trackedAPIKeyIDs, apiKeyID) + c.trackedAPIKeyRequestIDs = append(c.trackedAPIKeyRequestIDs, requestID) + return c.apiKeyTrackErr +} +func (c *stubConcurrencyCacheForTest) ReleaseAPIKeySlot(_ context.Context, apiKeyID int64, requestID string) error { + c.releasedAPIKeyIDs = append(c.releasedAPIKeyIDs, apiKeyID) + c.releasedAPIKeyRequestIDs = append(c.releasedAPIKeyRequestIDs, requestID) + return c.apiKeyReleaseErr +} +func (c *stubConcurrencyCacheForTest) GetAPIKeyConcurrencyBatch(_ context.Context, apiKeyIDs []int64) (map[int64]int, error) { + if c.apiKeyConcurrencyErr != nil { + return nil, c.apiKeyConcurrencyErr + } + result := make(map[int64]int, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + result[apiKeyID] = c.apiKeyConcurrency[apiKeyID] + } + return result, nil +} func (c *stubConcurrencyCacheForTest) IncrementWaitCount(_ context.Context, _ int64, _ int) (bool, error) { return c.waitAllowed, c.waitErr } @@ -201,6 +229,62 @@ func TestAcquireUserSlot_UnlimitedConcurrency(t *testing.T) { require.True(t, result.Acquired) } +func TestTrackAPIKeySlot_ReleaseDecrements(t *testing.T) { + cache := &stubConcurrencyCacheForTest{} + svc := NewConcurrencyService(cache) + + release := svc.TrackAPIKeySlot(context.Background(), 88) + require.NotNil(t, release) + require.Equal(t, []int64{88}, cache.trackedAPIKeyIDs) + require.Len(t, cache.trackedAPIKeyRequestIDs, 1) + require.NotEmpty(t, cache.trackedAPIKeyRequestIDs[0]) + + release() + + require.Equal(t, []int64{88}, cache.releasedAPIKeyIDs) + require.Equal(t, cache.trackedAPIKeyRequestIDs, cache.releasedAPIKeyRequestIDs) +} + +func TestTrackAPIKeySlot_FailOpen(t *testing.T) { + cache := &stubConcurrencyCacheForTest{apiKeyTrackErr: errors.New("redis down")} + svc := NewConcurrencyService(cache) + + release := svc.TrackAPIKeySlot(context.Background(), 88) + require.NotNil(t, release) + require.Equal(t, []int64{88}, cache.trackedAPIKeyIDs) + + require.NotPanics(t, release) + require.Empty(t, cache.releasedAPIKeyIDs) +} + +func TestGetAPIKeyConcurrencyBatch_Fallbacks(t *testing.T) { + t.Run("nil cache returns zeroes", func(t *testing.T) { + svc := &ConcurrencyService{cache: nil} + + counts, err := svc.GetAPIKeyConcurrencyBatch(context.Background(), []int64{1, 2}) + require.NoError(t, err) + require.Equal(t, map[int64]int{1: 0, 2: 0}, counts) + }) + + t.Run("redis error returns zeroes", func(t *testing.T) { + cache := &stubConcurrencyCacheForTest{apiKeyConcurrencyErr: errors.New("redis down")} + svc := NewConcurrencyService(cache) + + counts, err := svc.GetAPIKeyConcurrencyBatch(context.Background(), []int64{1, 2}) + require.NoError(t, err) + require.Equal(t, map[int64]int{1: 0, 2: 0}, counts) + }) + + t.Run("success returns counts", func(t *testing.T) { + cache := &stubConcurrencyCacheForTest{apiKeyConcurrency: map[int64]int{1: 3, 2: 0}} + svc := NewConcurrencyService(cache) + + counts, err := svc.GetAPIKeyConcurrencyBatch(context.Background(), []int64{1, 2}) + require.NoError(t, err) + require.Equal(t, map[int64]int{1: 3, 2: 0}, counts) + }) +} + func TestGenerateRequestID_UsesStablePrefixAndMonotonicCounter(t *testing.T) { id1 := generateRequestID() id2 := generateRequestID() diff --git a/backend/internal/service/domain_constants.go b/backend/internal/service/domain_constants.go index ab853c280a..bc1db19105 100644 --- a/backend/internal/service/domain_constants.go +++ b/backend/internal/service/domain_constants.go @@ -431,6 +431,20 @@ const ( // SettingKeyAllowUngroupedKeyScheduling 允许未分组 API Key 调度（默认 false：未分组 Key 返回 403） SettingKeyAllowUngroupedKeyScheduling = "allow_ungrouped_key_scheduling" + // SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled OpenAI 高级调度下是否启用粘性加权。 + SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled = "openai_advanced_scheduler_sticky_weighted_enabled" + // SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled OpenAI 高级调度下是否优先使用订阅账号池。 + SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled = "openai_advanced_scheduler_subscription_priority_enabled" + SettingKeyOpenAIAdvancedSchedulerLBTopK = "openai_advanced_scheduler_lb_top_k" + SettingKeyOpenAIAdvancedSchedulerWeightPriority = "openai_advanced_scheduler_weight_priority" + SettingKeyOpenAIAdvancedSchedulerWeightLoad = "openai_advanced_scheduler_weight_load" + SettingKeyOpenAIAdvancedSchedulerWeightQueue = "openai_advanced_scheduler_weight_queue" + SettingKeyOpenAIAdvancedSchedulerWeightErrorRate = "openai_advanced_scheduler_weight_error_rate" + SettingKeyOpenAIAdvancedSchedulerWeightTTFT = "openai_advanced_scheduler_weight_ttft" + SettingKeyOpenAIAdvancedSchedulerWeightReset = "openai_advanced_scheduler_weight_reset" + SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom = "openai_advanced_scheduler_weight_quota_headroom" + SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse = "openai_advanced_scheduler_weight_previous_response" + SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky = "openai_advanced_scheduler_weight_session_sticky" // SettingKeyBackendModeEnabled Backend 模式：禁用用户注册和自助服务，仅管理员可登录 SettingKeyBackendModeEnabled = "backend_mode_enabled" diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index f843ba3e45..35a60c8124 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -95,6 +95,9 @@ func (m *mockAccountRepoForPlatform) List(ctx context.Context, params pagination func (m *mockAccountRepoForPlatform) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) { return nil, nil, nil } +func (m *mockAccountRepoForPlatform) ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) { + return nil, nil +} func (m *mockAccountRepoForPlatform) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) { return nil, nil } diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index a90714ca1d..1665b1fe47 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -3,6 +3,7 @@ package service import ( "bytes" "encoding/json" + "errors" "fmt" "math" "regexp" @@ -168,7 +169,7 @@ func parseGatewayRequestCurrentBody(parsed *ParsedRequest, protocol string) erro bodyBytes := parsed.Body.Bytes() if !gjson.ValidBytes(bodyBytes) { - return fmt.Errorf("invalid json") + return DescribeInvalidJSON(bodyBytes) } // 只在当前函数内零拷贝读取 JSON 字段；ReplaceBody 后必须重新进入本函数刷新派生状态。 @@ -216,6 +217,26 @@ func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error { return parseGatewayRequestCurrentBody(parsed, protocol) } +// DescribeInvalidJSON returns a diagnostic error for a request body that +// failed JSON validation. It re-parses with encoding/json (failure path only) +// to pinpoint the first offending byte, so operators can distinguish genuinely +// invalid JSON from a truncated / partially consumed body. The error carries +// only length/offset/character information — never body content — so callers +// may safely wrap or log it. +func DescribeInvalidJSON(body []byte) error { + var raw json.RawMessage + if err := json.Unmarshal(body, &raw); err != nil { + var syntaxErr *json.SyntaxError + if errors.As(err, &syntaxErr) { + return fmt.Errorf("invalid json (len=%d, offset=%d): %s", len(body), syntaxErr.Offset, syntaxErr.Error()) + } + return fmt.Errorf("invalid json (len=%d): %w", len(body), err) + } + // gjson rejected the body but encoding/json accepted it (divergent edge + // cases, e.g. certain malformed UTF-8 sequences); report the basics. + return fmt.Errorf("invalid json (len=%d)", len(body)) +} + // ParsedRequest 保存网关请求的预解析结果 // // 性能优化说明： diff --git a/backend/internal/service/gateway_request_invalid_json_test.go b/backend/internal/service/gateway_request_invalid_json_test.go new file mode 100644 index 0000000000..cc69a41e72 --- /dev/null +++ b/backend/internal/service/gateway_request_invalid_json_test.go @@ -0,0 +1,52 @@ +//go:build unit + +package service + +import ( + "fmt" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/domain" + "github.com/stretchr/testify/require" +) + +func TestDescribeInvalidJSON_TruncatedBody(t *testing.T) { + // Simulates a body cut off mid-stream (e.g. partially consumed by middleware). + body := []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi`) + + err := DescribeInvalidJSON(body) + + require.Error(t, err) + require.Contains(t, err.Error(), fmt.Sprintf("len=%d", len(body))) + require.Contains(t, err.Error(), "unexpected end of JSON input") +} + +func TestDescribeInvalidJSON_InvalidCharacterWithOffset(t *testing.T) { + body := []byte(`{"model": bad}`) + + err := DescribeInvalidJSON(body) + + require.Error(t, err) + require.Contains(t, err.Error(), "offset=11") + require.Contains(t, err.Error(), "invalid character") +} + +func TestDescribeInvalidJSON_DoesNotLeakBodyContent(t *testing.T) { + secret := "sk-super-secret-value" + body := []byte(`{"api_key":"` + secret + `","broken":`) + + err := DescribeInvalidJSON(body) + + require.Error(t, err) + require.NotContains(t, err.Error(), secret) +} + +func TestParseGatewayRequest_InvalidJSONErrorIsDiagnostic(t *testing.T) { + body := []byte(`{"model":"claude-sonnet-4-6","messages":[`) + + _, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) + + require.Error(t, err) + require.True(t, strings.HasPrefix(err.Error(), "invalid json (len="), "error should carry diagnostics, got: %s", err.Error()) +} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 54035345d9..2ca554edb3 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -5105,6 +5105,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { return nil, err } + // Pre-filter: strip web-search history blocks the upstream cannot accept + // (emulation-synthesized server_tool_use / web_search_tool_result always; + // genuine ones additionally for passback-required upstreams). See + // FilterWebSearchHistoryBlocks. reqModel 此时已是映射后的模型 ID。 + if err := replaceBody(FilterWebSearchHistoryBlocks(body, reqModel)); err != nil { + return nil, err + } // Pre-filter: remove thinking blocks with missing/invalid signatures before forwarding. // Clients (e.g. Claude Code) sometimes send multi-turn conversations where a historical // assistant message contains a thinking block that is missing the required "signature" field, @@ -5688,6 +5695,11 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( } // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. input.Body = StripEmptyTextBlocks(input.Body) + // Pre-filter: strip web-search history blocks the upstream cannot accept + // (emulation-synthesized ones always; genuine ones additionally for + // passback-required third-party upstreams such as GLM/Kimi/DeepSeek, + // which reject server_tool_use with 400). input.RequestModel 已是映射后的模型 ID。 + input.Body = FilterWebSearchHistoryBlocks(input.Body, input.RequestModel) if input.Parsed != nil { // 透传分支也会改写实际 wire body，成功 usage hash 依赖这里同步当前 body。 if err := input.Parsed.ReplaceBody(input.Body); err != nil { @@ -5920,6 +5932,10 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( if c != nil && c.Request != nil { clientBeta = getHeaderRaw(c.Request.Header, "anthropic-beta") } + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + clientBeta = beta + } if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { body = sanitized } @@ -5956,6 +5972,9 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") } + // 账号级请求头覆写（最终生效，覆盖上面所有来源的同名头） + account.ApplyHeaderOverrides(req.Header) + return req, body, nil } @@ -6886,6 +6905,12 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex tokenType, mimicClaudeCode, modelID, clientHeaders, body, effectiveDropSet, ) + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值（由下方 ApplyHeaderOverrides 写入）： + // body 能力净化必须以覆写值为准，否则 header/body 不对称会被上游 400。 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + finalBetaHeader, finalBetaShouldSet = beta, true + } + // 能力维度 body sanitize：与最终 anthropic-beta header 对称 if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBetaHeader); changed { body = sanitized @@ -6959,6 +6984,10 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex } } + // 账号级请求头覆写（仅 anthropic/openai api_key 账号启用时生效；OAuth 路径 no-op）。 + // 放在所有 header 逻辑之后，确保配置值对同名头拥有最终决定权。 + account.ApplyHeaderOverrides(req.Header) + // === DEBUG: 打印上游转发请求（headers + body 摘要），与 CLIENT_ORIGINAL 对比 === s.debugLogGatewaySnapshot("UPSTREAM_FORWARD", req.Header, body, map[string]string{ "url": req.URL.String(), @@ -10410,6 +10439,10 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( if c != nil && c.Request != nil { clientBeta = getHeaderRaw(c.Request.Header, "anthropic-beta") } + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + clientBeta = beta + } if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { body = sanitized } @@ -10445,6 +10478,9 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( req.Header.Set("anthropic-version", "2023-06-01") } + // 账号级请求头覆写（最终生效，覆盖上面所有来源的同名头） + account.ApplyHeaderOverrides(req.Header) + return req, nil } @@ -10512,6 +10548,11 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con tokenType, mimicClaudeCode, modelID, clientHeaders, body, ctEffectiveDropSet, ) + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + finalBetaHeader, finalBetaShouldSet = beta, true + } + // 能力维度 body sanitize：与最终 anthropic-beta header 对称 if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBetaHeader); changed { body = sanitized @@ -10578,6 +10619,9 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con } } + // 账号级请求头覆写（仅 anthropic/openai api_key 账号启用时生效；OAuth 路径 no-op） + account.ApplyHeaderOverrides(req.Header) + if c != nil && tokenType == "oauth" { c.Set(claudeMimicDebugInfoKey, buildClaudeMimicDebugLine(req, body, account, tokenType, mimicClaudeCode)) } diff --git a/backend/internal/service/gateway_websearch_block_filter.go b/backend/internal/service/gateway_websearch_block_filter.go new file mode 100644 index 0000000000..a0706c5c2d --- /dev/null +++ b/backend/internal/service/gateway_websearch_block_filter.go @@ -0,0 +1,138 @@ +package service + +import ( + "bytes" + "encoding/json" + "strings" + "unsafe" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const ( + blockTypeServerToolUse = "server_tool_use" + blockTypeWebSearchToolResult = "web_search_tool_result" +) + +// Fast-path byte patterns: both block types only ever appear as quoted JSON +// string values, so a raw substring check is a safe pre-filter regardless of +// key/value spacing. +var ( + patternServerToolUse = []byte(`"server_tool_use"`) + patternWebSearchToolResult = []byte(`"web_search_tool_result"`) +) + +// FilterWebSearchHistoryBlocks removes web-search content blocks from +// historical messages when the upstream cannot accept them: +// +// 1. Emulation-synthesized blocks — server_tool_use / web_search_tool_result +// whose tool-use ID carries webSearchToolUseIDPrefix — are fabricated +// locally by the web-search emulation (gateway_websearch_emulation.go). +// No upstream ever issued them, so clients replaying the conversation +// (e.g. Claude Code) poison every follow-up request. They are stripped +// for all upstreams. +// 2. For passback-required upstreams (DeepSeek/Kimi/GLM …, see +// ResolveThinkingProtocol) all server_tool_use / web_search_tool_result +// blocks are stripped: these upstreams only accept +// text/thinking/image/tool_use/tool_result and reject anything else with +// 400 "invalid value: `server_tool_use`". anthropic-strict and unknown +// upstreams keep genuine blocks untouched. +// +// The emulated assistant turn always carries a trailing text summary, so the +// search context survives the strip. A message whose content would become +// empty gets a placeholder text block (mirroring FilterThinkingBlocksForRetry). +// Returns the original body unchanged when nothing needs stripping. +func FilterWebSearchHistoryBlocks(body []byte, mappedModel string) []byte { + if !bytes.Contains(body, patternServerToolUse) && !bytes.Contains(body, patternWebSearchToolResult) { + return body + } + + stripAll := ResolveThinkingProtocol(mappedModel) == ThinkingProtocolPassbackRequired + + jsonStr := *(*string)(unsafe.Pointer(&body)) + msgsRes := gjson.Get(jsonStr, "messages") + if !msgsRes.Exists() || !msgsRes.IsArray() { + return body + } + + var messages []any + if err := json.Unmarshal(sliceRawFromBody(body, msgsRes), &messages); err != nil { + return body + } + + modified := false + for _, msg := range messages { + msgMap, ok := msg.(map[string]any) + if !ok { + continue + } + content, ok := msgMap["content"].([]any) + if !ok { + continue + } + + // 延迟分配：只有命中需剥离的块才构建新 slice。 + var newContent []any + for i, block := range content { + blockMap, isMap := block.(map[string]any) + if isMap && shouldStripWebSearchBlock(blockMap, stripAll) { + if newContent == nil { + newContent = make([]any, 0, len(content)) + newContent = append(newContent, content[:i]...) + } + continue + } + if newContent != nil { + newContent = append(newContent, block) + } + } + if newContent == nil { + continue + } + modified = true + if len(newContent) == 0 { + role, _ := msgMap["role"].(string) + placeholder := "(content removed)" + if role == "assistant" { + placeholder = "(assistant content removed)" + } + newContent = []any{map[string]any{"type": "text", "text": placeholder}} + } + msgMap["content"] = newContent + } + + if !modified { + return body + } + + msgsBytes, err := json.Marshal(messages) + if err != nil { + return body + } + out, err := sjson.SetRawBytes(body, "messages", msgsBytes) + if err != nil { + return body + } + return out +} + +func shouldStripWebSearchBlock(block map[string]any, stripAll bool) bool { + blockType, _ := block["type"].(string) + switch blockType { + case blockTypeServerToolUse: + if stripAll { + return true + } + id, _ := block["id"].(string) + return strings.HasPrefix(id, webSearchToolUseIDPrefix) + case blockTypeWebSearchToolResult: + if stripAll { + return true + } + id, _ := block["tool_use_id"].(string) + return strings.HasPrefix(id, webSearchToolUseIDPrefix) + default: + return false + } +} diff --git a/backend/internal/service/gateway_websearch_block_filter_test.go b/backend/internal/service/gateway_websearch_block_filter_test.go new file mode 100644 index 0000000000..cebda4d37b --- /dev/null +++ b/backend/internal/service/gateway_websearch_block_filter_test.go @@ -0,0 +1,140 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +// emulatedWebSearchBody is a follow-up /v1/messages request whose history +// contains an assistant turn synthesized by the web-search emulation +// (server_tool_use + web_search_tool_result with the local srvtoolu_ws_ ID +// prefix, followed by the text summary). +const emulatedWebSearchBody = `{"model":"claude-sonnet-4-6","max_tokens":1024,"messages":[` + + `{"role":"user","content":[{"type":"text","text":"search the weather"}]},` + + `{"role":"assistant","content":[` + + `{"type":"server_tool_use","id":"srvtoolu_ws_0123456789abcdef","name":"web_search","input":{"query":"weather"}},` + + `{"type":"web_search_tool_result","tool_use_id":"srvtoolu_ws_0123456789abcdef","content":[{"type":"web_search_result","url":"https://example.com","title":"Weather"}]},` + + `{"type":"text","text":"Here are the search results for \"weather\":"}]},` + + `{"role":"user","content":[{"type":"text","text":"thanks, continue"}]}]}` + +// genuineWebSearchBody carries real Anthropic web-search blocks (upstream IDs +// do NOT have the local srvtoolu_ws_ prefix). +const genuineWebSearchBody = `{"model":"claude-sonnet-4-6","max_tokens":1024,"messages":[` + + `{"role":"user","content":[{"type":"text","text":"search"}]},` + + `{"role":"assistant","content":[` + + `{"type":"server_tool_use","id":"srvtoolu_01ABCDEF","name":"web_search","input":{"query":"weather"}},` + + `{"type":"web_search_tool_result","tool_use_id":"srvtoolu_01ABCDEF","content":[{"type":"web_search_result","url":"https://example.com","title":"Weather"}]},` + + `{"type":"text","text":"summary with citations"}]}]}` + +func collectContentTypes(t *testing.T, body []byte) []string { + t.Helper() + var types []string + for _, msg := range gjson.GetBytes(body, "messages").Array() { + for _, block := range msg.Get("content").Array() { + types = append(types, block.Get("type").String()) + } + } + return types +} + +func TestFilterWebSearchHistoryBlocks_StripsEmulatedBlocksForAnthropicStrict(t *testing.T) { + out := FilterWebSearchHistoryBlocks([]byte(emulatedWebSearchBody), "claude-sonnet-4-6") + + require.Equal(t, []string{"text", "text", "text"}, collectContentTypes(t, out)) + // The emulated text summary must survive so the search context is preserved. + require.Contains(t, string(out), "Here are the search results") + require.NotContains(t, string(out), "srvtoolu_ws_") + require.True(t, gjson.ValidBytes(out)) +} + +func TestFilterWebSearchHistoryBlocks_KeepsGenuineBlocksForAnthropicStrict(t *testing.T) { + body := []byte(genuineWebSearchBody) + out := FilterWebSearchHistoryBlocks(body, "claude-sonnet-4-6") + + require.Equal(t, string(body), string(out)) +} + +func TestFilterWebSearchHistoryBlocks_StripsAllBlocksForPassbackRequired(t *testing.T) { + // GLM only accepts text/thinking/image/tool_use/tool_result and rejects + // server_tool_use with 400, so genuine blocks must be stripped as well. + out := FilterWebSearchHistoryBlocks([]byte(genuineWebSearchBody), "glm-4.7") + + require.Equal(t, []string{"text", "text"}, collectContentTypes(t, out)) + require.NotContains(t, string(out), "server_tool_use") + require.NotContains(t, string(out), "web_search_tool_result") + require.Contains(t, string(out), "summary with citations") +} + +func TestFilterWebSearchHistoryBlocks_StripsEmulatedBlocksForUnknownModel(t *testing.T) { + out := FilterWebSearchHistoryBlocks([]byte(emulatedWebSearchBody), "totally-unknown-model") + + require.Equal(t, []string{"text", "text", "text"}, collectContentTypes(t, out)) + require.NotContains(t, string(out), "srvtoolu_ws_") +} + +func TestFilterWebSearchHistoryBlocks_KeepsGenuineBlocksForUnknownModel(t *testing.T) { + body := []byte(genuineWebSearchBody) + out := FilterWebSearchHistoryBlocks(body, "totally-unknown-model") + + require.Equal(t, string(body), string(out)) +} + +func TestFilterWebSearchHistoryBlocks_NoWebSearchBlocksFastPath(t *testing.T) { + body := []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) + out := FilterWebSearchHistoryBlocks(body, "claude-sonnet-4-6") + + require.Equal(t, string(body), string(out)) +} + +func TestFilterWebSearchHistoryBlocks_EmptiedMessageGetsPlaceholder(t *testing.T) { + body := []byte(`{"model":"glm-4.7","messages":[` + + `{"role":"user","content":[{"type":"text","text":"search"}]},` + + `{"role":"assistant","content":[` + + `{"type":"server_tool_use","id":"srvtoolu_01X","name":"web_search","input":{"query":"q"}},` + + `{"type":"web_search_tool_result","tool_use_id":"srvtoolu_01X","content":[]}]}]}`) + + out := FilterWebSearchHistoryBlocks(body, "glm-4.7") + + msgs := gjson.GetBytes(out, "messages").Array() + require.Len(t, msgs, 2) + assistant := msgs[1] + require.Equal(t, "assistant", assistant.Get("role").String()) + content := assistant.Get("content").Array() + require.Len(t, content, 1) + require.Equal(t, "text", content[0].Get("type").String()) + require.Equal(t, "(assistant content removed)", content[0].Get("text").String()) +} + +func TestFilterWebSearchHistoryBlocks_StringContentUntouched(t *testing.T) { + // A string mentioning the pattern inside a text value must not trigger a rewrite. + body := []byte(`{"model":"claude-sonnet-4-6","messages":[` + + `{"role":"user","content":"please explain \"server_tool_use\" blocks"}]}`) + + out := FilterWebSearchHistoryBlocks(body, "claude-sonnet-4-6") + + require.Equal(t, string(body), string(out)) +} + +func TestFilterWebSearchHistoryBlocks_InvalidMessagesUnchanged(t *testing.T) { + body := []byte(`{"model":"claude-sonnet-4-6","messages":"server_tool_use"}`) + out := FilterWebSearchHistoryBlocks(body, "claude-sonnet-4-6") + + require.Equal(t, string(body), string(out)) +} + +func TestFilterWebSearchHistoryBlocks_PreservesOtherToolBlocks(t *testing.T) { + body := []byte(`{"model":"glm-4.7","messages":[` + + `{"role":"assistant","content":[` + + `{"type":"tool_use","id":"toolu_01A","name":"get_weather","input":{}},` + + `{"type":"server_tool_use","id":"srvtoolu_ws_abc","name":"web_search","input":{"query":"q"}},` + + `{"type":"text","text":"result"}]},` + + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_01A","content":"sunny"}]}]}`) + + out := FilterWebSearchHistoryBlocks(body, "glm-4.7") + + require.Equal(t, []string{"tool_use", "text", "tool_result"}, collectContentTypes(t, out)) +} diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index c021e88edf..7d5ed0ec9e 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -82,6 +82,9 @@ func (m *mockAccountRepoForGemini) List(ctx context.Context, params pagination.P func (m *mockAccountRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) { return nil, nil, nil } +func (m *mockAccountRepoForGemini) ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error) { + return nil, nil +} func (m *mockAccountRepoForGemini) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) { return nil, nil } diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 8942404eaa..f8b7ff3e83 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -483,6 +483,7 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody) case GrokMediaEndpointVideosGenerations: meta.ResponseID = extractGrokMediaVideoRequestID(responseBody) + // Video generation is one billable media unit; the legacy usage schema stores it in ImageCount. meta.ImageCount = 1 meta.ImageSize = requestInfo.SizeTier meta.ImageInputSize = requestInfo.Size diff --git a/backend/internal/service/group.go b/backend/internal/service/group.go index 6d0a11f766..e3a1697b57 100644 --- a/backend/internal/service/group.go +++ b/backend/internal/service/group.go @@ -36,12 +36,15 @@ type Group struct { DefaultValidityDays int // 图片生成计费配置（antigravity 和 gemini 平台使用） - AllowImageGeneration bool - ImageRateIndependent bool - ImageRateMultiplier float64 - ImagePrice1K *float64 - ImagePrice2K *float64 - ImagePrice4K *float64 + AllowImageGeneration bool + AllowBatchImageGeneration bool + ImageRateIndependent bool + ImageRateMultiplier float64 + ImagePrice1K *float64 + ImagePrice2K *float64 + ImagePrice4K *float64 + BatchImageDiscountMultiplier float64 + BatchImageHoldMultiplier float64 // Claude Code 客户端限制 ClaudeCodeOnly bool diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index ab298d4521..a0a4fbff3e 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -36,8 +36,20 @@ const ( ) type cachedOpenAIAdvancedSchedulerSetting struct { - enabled bool - expiresAt int64 + enabled bool + stickyWeightedEnabled bool + subscriptionPriorityEnabled bool + lbTopKOverride int + weightOverrides map[string]float64 + expiresAt int64 +} + +type openAIAdvancedSchedulerRuntimeSettings struct { + enabled bool + stickyWeightedEnabled bool + subscriptionPriorityEnabled bool + lbTopKOverride int + weightOverrides map[string]float64 } var openAIAdvancedSchedulerSettingCache atomic.Value // *cachedOpenAIAdvancedSchedulerSetting @@ -48,8 +60,12 @@ type OpenAIAccountScheduleRequest struct { Platform string SessionHash string StickyAccountID int64 + StickyPreviousAccountID int64 + StickyWeighted bool + SubscriptionPriority bool PreserveStickyBinding bool PreviousResponseID string + PreviousResponseCanMove bool RequestedModel string RequiredTransport OpenAIUpstreamTransport RequiredCapability OpenAIEndpointCapability @@ -111,6 +127,17 @@ type openAIAccountLoadPlan struct { loadSkew float64 } +type openAIAccountLoadSelectionAttempt struct { + result *AccountSelectionResult + selectionOrder []openAIAccountCandidateScore + candidateCount int + topK int + loadSkew float64 + compactBlocked bool + noCompactCandidates bool + err error +} + func (m *openAIAccountSchedulerMetrics) recordSelect(decision OpenAIAccountScheduleDecision) { if m == nil { return @@ -277,7 +304,8 @@ func (s *defaultOpenAIAccountScheduler) Select( }() previousResponseID := strings.TrimSpace(req.PreviousResponseID) - if previousResponseID != "" && normalizeOpenAICompatiblePlatform(req.Platform) == PlatformOpenAI { + if previousResponseID != "" && normalizeOpenAICompatiblePlatform(req.Platform) == PlatformOpenAI && + (!req.StickyWeighted || !req.PreviousResponseCanMove) { selection, err := s.service.selectAccountByPreviousResponseIDForCapability( ctx, req.GroupID, @@ -310,19 +338,21 @@ func (s *defaultOpenAIAccountScheduler) Select( } } - selection, escapedSticky, err := s.selectBySessionHash(ctx, req) - if err != nil { - return nil, decision, err - } - if selection != nil && selection.Account != nil { - decision.Layer = openAIAccountScheduleLayerSessionSticky - decision.StickySessionHit = true - decision.SelectedAccountID = selection.Account.ID - decision.SelectedAccountType = selection.Account.Type - return selection, decision, nil - } - if escapedSticky { - req.PreserveStickyBinding = true + if !req.StickyWeighted { + selection, escapedSticky, err := s.selectBySessionHash(ctx, req) + if err != nil { + return nil, decision, err + } + if selection != nil && selection.Account != nil { + decision.Layer = openAIAccountScheduleLayerSessionSticky + decision.StickySessionHit = true + decision.SelectedAccountID = selection.Account.ID + decision.SelectedAccountType = selection.Account.Type + return selection, decision, nil + } + if escapedSticky { + req.PreserveStickyBinding = true + } } selection, candidateCount, topK, loadSkew, err := s.selectByLoadBalance(ctx, req) @@ -336,6 +366,14 @@ func (s *defaultOpenAIAccountScheduler) Select( if selection != nil && selection.Account != nil { decision.SelectedAccountID = selection.Account.ID decision.SelectedAccountType = selection.Account.Type + if req.StickyWeighted { + if req.StickyPreviousAccountID > 0 && selection.Account.ID == req.StickyPreviousAccountID { + decision.StickyPreviousHit = true + } + if req.StickyAccountID > 0 && selection.Account.ID == req.StickyAccountID { + decision.StickySessionHit = true + } + } } return selection, decision, nil } @@ -453,6 +491,13 @@ func openAIStickyAccountMatchesGroup(account *Account, groupID *int64) bool { return false } +func openAIAccountSchedulingPriority(account *Account) int { + if account == nil { + return 0 + } + return account.Priority +} + func (s *defaultOpenAIAccountScheduler) shouldEscapeStickyAccount(accountID int64, cfg openAIStickyEscapeConfig) (reason string, errorRate float64, ttft float64, shouldEscape bool) { if !cfg.enabled || s == nil || s.stats == nil || accountID <= 0 { return "", 0, 0, false @@ -471,6 +516,7 @@ type openAIAccountCandidateScore struct { account *Account loadInfo *AccountLoadInfo score float64 + priority int errorRate float64 ttft float64 hasTTFT bool @@ -669,6 +715,7 @@ func buildOpenAIWeightedSelectionOrder( } func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan( + ctx context.Context, req OpenAIAccountScheduleRequest, filtered []*Account, loadMap map[int64]*AccountLoadInfo, @@ -716,18 +763,20 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan( return plan } - minPriority, maxPriority := candidates[0].account.Priority, candidates[0].account.Priority + minPriority, maxPriority := openAIAccountSchedulingPriority(candidates[0].account), openAIAccountSchedulingPriority(candidates[0].account) maxWaiting := 1 loadRateSum := 0.0 loadRateSumSquares := 0.0 minTTFT, maxTTFT := 0.0, 0.0 hasTTFTSample := false - for _, candidate := range candidates { - if candidate.account.Priority < minPriority { - minPriority = candidate.account.Priority + for i := range candidates { + candidate := &candidates[i] + candidate.priority = openAIAccountSchedulingPriority(candidate.account) + if candidate.priority < minPriority { + minPriority = candidate.priority } - if candidate.account.Priority > maxPriority { - maxPriority = candidate.account.Priority + if candidate.priority > maxPriority { + maxPriority = candidate.priority } if candidate.loadInfo.WaitingCount > maxWaiting { maxWaiting = candidate.loadInfo.WaitingCount @@ -751,7 +800,7 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan( } plan.loadSkew = calcLoadSkewByMoments(loadRateSum, loadRateSumSquares, len(candidates)) - weights := s.service.openAIWSSchedulerWeights() + weights := s.service.openAIWSSchedulerWeightsForRequest(ctx) // Reset 因子（use-it-or-lose-it）：在拥有「未来会话窗口结束时间」的账号中， // 剩余时间越短 → 因子越接近 1（越早重置越优先用尽）。无活跃窗口的账号因子为 0。 @@ -785,7 +834,7 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan( item := &candidates[i] priorityFactor := 1.0 if maxPriority > minPriority { - priorityFactor = 1 - float64(item.account.Priority-minPriority)/float64(maxPriority-minPriority) + priorityFactor = 1 - float64(item.priority-minPriority)/float64(maxPriority-minPriority) } loadFactor := 1 - clamp01(float64(item.loadInfo.LoadRate)/100.0) queueFactor := 1 - clamp01(float64(item.loadInfo.WaitingCount)/float64(maxWaiting)) @@ -817,10 +866,18 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan( weights.TTFT*ttftFactor + weights.Reset*resetFactor + weights.QuotaHeadroom*quotaHeadroomFactor + if req.StickyWeighted { + if req.PreviousResponseCanMove && req.StickyPreviousAccountID > 0 && item.account.ID == req.StickyPreviousAccountID { + item.score += weights.Previous + } + if req.StickyAccountID > 0 && item.account.ID == req.StickyAccountID { + item.score += weights.SessionSticky + } + } } plan.candidates = candidates - plan.topK = s.service.openAIWSLBTopK() + plan.topK = s.service.openAIWSLBTopKForRequest(ctx) if plan.topK > len(candidates) { plan.topK = len(candidates) } @@ -845,6 +902,20 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAISelectionOrder( groupTopK = len(pool) } ranked := selectTopKOpenAICandidates(pool, groupTopK) + if req.StickyWeighted { + for _, stickyID := range []int64{req.StickyPreviousAccountID, req.StickyAccountID} { + if stickyID <= 0 { + continue + } + for i, candidate := range ranked { + if candidate.account != nil && candidate.account.ID == stickyID { + ordered := append([]openAIAccountCandidateScore{candidate}, ranked[:i]...) + ordered = append(ordered, ranked[i+1:]...) + return ordered + } + } + } + } return buildOpenAIWeightedSelectionOrder(ranked, req) } @@ -939,6 +1010,74 @@ func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAISelectionOrder( return nil, compactBlocked, nil } +func (s *defaultOpenAIAccountScheduler) tryFallbackToWeightedSticky( + ctx context.Context, + req OpenAIAccountScheduleRequest, +) (*AccountSelectionResult, error) { + if !req.StickyWeighted { + return nil, nil + } + for _, accountID := range []int64{req.StickyPreviousAccountID, req.StickyAccountID} { + if accountID <= 0 { + continue + } + if req.ExcludedIDs != nil { + if _, excluded := req.ExcludedIDs[accountID]; excluded { + continue + } + } + account, err := s.service.getSchedulableAccount(ctx, accountID) + if err != nil || account == nil { + continue + } + if !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) { + continue + } + account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability) + if account == nil || !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) { + continue + } + // 粘性绑定只证明绑定时账号在分组内；账号被移出分组后绑定仍会在 TTL 内存活， + // 必须与 selectBySessionHash 一样重验分组归属，否则会把分组流量泄漏到组外账号。 + if !openAIStickyAccountMatchesGroup(account, req.GroupID) { + if accountID == req.StickyAccountID && strings.TrimSpace(req.SessionHash) != "" { + _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, req.SessionHash) + } + continue + } + if req.RequireCompact && openAICompactSupportTier(account) == 0 { + continue + } + result, acquireErr := s.service.tryAcquireAccountSlot(ctx, account.ID, account.Concurrency) + if acquireErr != nil { + return nil, acquireErr + } + if result != nil && result.Acquired { + if req.SessionHash != "" && !req.PreserveStickyBinding { + _ = s.service.BindStickySession(ctx, req.GroupID, req.SessionHash, account.ID) + } + return &AccountSelectionResult{ + Account: account, + Acquired: true, + ReleaseFunc: result.ReleaseFunc, + }, nil + } + if s.service.concurrencyService != nil { + cfg := s.service.schedulingConfig() + return &AccountSelectionResult{ + Account: account, + WaitPlan: &AccountWaitPlan{ + AccountID: account.ID, + MaxConcurrency: account.Concurrency, + Timeout: cfg.StickySessionWaitTimeout, + MaxWaiting: cfg.StickySessionMaxWaiting, + }, + }, nil + } + } + return nil, nil +} + func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( ctx context.Context, req OpenAIAccountScheduleRequest, @@ -1002,52 +1141,175 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( } } - plan := s.buildOpenAIAccountLoadPlan(req, filtered, loadMap) - candidateCount := plan.candidateCount - topK := plan.topK - loadSkew := plan.loadSkew - selectionOrder := plan.selectionOrder - if req.RequireCompact && len(plan.candidates) == 0 && len(plan.staleSnapshotCompactRetry) == 0 { - return nil, 0, 0, 0, ErrNoAvailableCompactAccounts - } - if req.RequireCompact && len(selectionOrder) == 0 && s.service.schedulerSnapshot == nil { - return nil, candidateCount, topK, loadSkew, ErrNoAvailableCompactAccounts - } - if len(selectionOrder) == 0 { - return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, req.RequireCompact && len(plan.allCandidates) > 0) + if req.SubscriptionPriority { + subscriptionAccounts, regularAccounts := partitionOpenAIChatGPTSubscriptionAccounts(filtered) + if len(subscriptionAccounts) > 0 { + attempt := s.trySelectByLoadBalancePool(ctx, req, subscriptionAccounts, loadMap) + if attempt.err != nil && (!attempt.noCompactCandidates || len(regularAccounts) <= 0) { + return nil, attempt.candidateCount, attempt.topK, attempt.loadSkew, attempt.err + } + if attempt.result != nil { + return attempt.result, attempt.candidateCount, attempt.topK, attempt.loadSkew, nil + } + if len(regularAccounts) > 0 { + regularAttempt := s.trySelectByLoadBalancePool(ctx, req, regularAccounts, loadMap) + if regularAttempt.err != nil && !regularAttempt.noCompactCandidates { + return nil, regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew, regularAttempt.err + } + if regularAttempt.result != nil { + return regularAttempt.result, regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew, nil + } + var result *AccountSelectionResult + candidateCount, topK, loadSkew := regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew + fallbackErr := regularAttempt.err + if regularAttempt.err == nil { + result, candidateCount, topK, loadSkew, fallbackErr = s.finishLoadBalanceSelectionFallback(ctx, req, regularAttempt) + if fallbackErr == nil && result != nil { + return result, candidateCount, topK, loadSkew, nil + } + } + // 常规池既无法获取也无法排队（含仅剩不支持 compact 的候选）时， + // 回退到订阅池的等待计划：busy-but-waitable 的订阅账号不应因常规池存在 + // 而被丢弃，否则开启订阅优先反而让本可排队成功的请求硬失败。 + subResult, subCandidateCount, subTopK, subLoadSkew, subErr := s.finishLoadBalanceSelectionFallback(ctx, req, attempt) + if subErr == nil && subResult != nil { + return subResult, subCandidateCount, subTopK, subLoadSkew, nil + } + return result, candidateCount, topK, loadSkew, fallbackErr + } + return s.finishLoadBalanceSelectionFallback(ctx, req, attempt) + } } - result, compactBlocked, acquireErr := s.tryAcquireOpenAISelectionOrder(ctx, req, selectionOrder) + attempt := s.trySelectByLoadBalancePool(ctx, req, filtered, loadMap) + if attempt.err != nil { + return nil, attempt.candidateCount, attempt.topK, attempt.loadSkew, attempt.err + } + if attempt.result != nil { + return attempt.result, attempt.candidateCount, attempt.topK, attempt.loadSkew, nil + } + return s.finishLoadBalanceSelectionFallback(ctx, req, attempt) +} + +func partitionOpenAIChatGPTSubscriptionAccounts(accounts []*Account) ([]*Account, []*Account) { + subscriptionAccounts := make([]*Account, 0, len(accounts)) + regularAccounts := make([]*Account, 0, len(accounts)) + for _, account := range accounts { + if account != nil && account.IsOpenAIChatGPTSubscription() { + subscriptionAccounts = append(subscriptionAccounts, account) + continue + } + regularAccounts = append(regularAccounts, account) + } + return subscriptionAccounts, regularAccounts +} + +func (s *defaultOpenAIAccountScheduler) trySelectByLoadBalancePool( + ctx context.Context, + req OpenAIAccountScheduleRequest, + filtered []*Account, + loadMap map[int64]*AccountLoadInfo, +) openAIAccountLoadSelectionAttempt { + plan := s.buildOpenAIAccountLoadPlan(ctx, req, filtered, loadMap) + attempt := openAIAccountLoadSelectionAttempt{ + selectionOrder: plan.selectionOrder, + candidateCount: plan.candidateCount, + topK: plan.topK, + loadSkew: plan.loadSkew, + } + if req.RequireCompact && len(plan.candidates) == 0 && len(plan.staleSnapshotCompactRetry) == 0 { + attempt.noCompactCandidates = true + attempt.err = ErrNoAvailableCompactAccounts + return attempt + } + if req.RequireCompact && len(attempt.selectionOrder) == 0 && s.service.schedulerSnapshot == nil { + attempt.noCompactCandidates = true + attempt.err = ErrNoAvailableCompactAccounts + return attempt + } + if len(attempt.selectionOrder) == 0 { + attempt.compactBlocked = req.RequireCompact && len(plan.allCandidates) > 0 + return attempt + } + + result, compactBlocked, acquireErr := s.tryAcquireOpenAISelectionOrder(ctx, req, attempt.selectionOrder) + attempt.compactBlocked = compactBlocked if acquireErr != nil { - return nil, candidateCount, topK, loadSkew, acquireErr + attempt.err = acquireErr + return attempt } if result != nil { - return result, candidateCount, topK, loadSkew, nil + attempt.result = result + return attempt } if s.service.concurrencyService != nil { + loadReq := buildOpenAIAccountLoadRequest(filtered) if freshLoadMap, loadErr := s.service.concurrencyService.GetAccountsLoadBatchFresh(ctx, loadReq); loadErr == nil { - freshPlan := s.buildOpenAIAccountLoadPlan(req, filtered, freshLoadMap) + freshPlan := s.buildOpenAIAccountLoadPlan(ctx, req, filtered, freshLoadMap) if len(freshPlan.selectionOrder) > 0 { freshResult, freshCompactBlocked, freshAcquireErr := s.tryAcquireOpenAISelectionOrder(ctx, req, freshPlan.selectionOrder) if freshAcquireErr != nil { - return nil, candidateCount, topK, loadSkew, freshAcquireErr + attempt.err = freshAcquireErr + return attempt } if freshResult != nil { - return freshResult, freshPlan.candidateCount, freshPlan.topK, freshPlan.loadSkew, nil + attempt.result = freshResult + attempt.selectionOrder = freshPlan.selectionOrder + attempt.candidateCount = freshPlan.candidateCount + attempt.topK = freshPlan.topK + attempt.loadSkew = freshPlan.loadSkew + return attempt } - compactBlocked = compactBlocked || freshCompactBlocked - selectionOrder = freshPlan.selectionOrder - candidateCount = freshPlan.candidateCount - topK = freshPlan.topK - loadSkew = freshPlan.loadSkew + attempt.compactBlocked = attempt.compactBlocked || freshCompactBlocked + attempt.selectionOrder = freshPlan.selectionOrder + attempt.candidateCount = freshPlan.candidateCount + attempt.topK = freshPlan.topK + attempt.loadSkew = freshPlan.loadSkew } } } + return attempt +} + +func buildOpenAIAccountLoadRequest(accounts []*Account) []AccountWithConcurrency { + loadReq := make([]AccountWithConcurrency, 0, len(accounts)) + for _, account := range accounts { + if account == nil { + continue + } + loadReq = append(loadReq, AccountWithConcurrency{ + ID: account.ID, + MaxConcurrency: account.EffectiveLoadFactor(), + }) + } + return loadReq +} + +func (s *defaultOpenAIAccountScheduler) finishLoadBalanceSelectionFallback( + ctx context.Context, + req OpenAIAccountScheduleRequest, + attempt openAIAccountLoadSelectionAttempt, +) (*AccountSelectionResult, int, int, float64, error) { + candidateCount := attempt.candidateCount + topK := attempt.topK + loadSkew := attempt.loadSkew + + if len(attempt.selectionOrder) == 0 { + return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, attempt.compactBlocked) + } + + if stickyFallback, stickyErr := s.tryFallbackToWeightedSticky(ctx, req); stickyErr != nil { + return nil, candidateCount, topK, loadSkew, stickyErr + } else if stickyFallback != nil { + return stickyFallback, candidateCount, topK, loadSkew, nil + } + cfg := s.service.schedulingConfig() + compactBlocked := attempt.compactBlocked // WaitPlan.MaxConcurrency 使用 Concurrency（非 EffectiveLoadFactor），因为 WaitPlan 控制的是 Redis 实际并发槽位等待。 - for _, candidate := range selectionOrder { + for _, candidate := range attempt.selectionOrder { fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.Platform, req.RequestedModel, false, req.RequiredCapability) if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) { continue @@ -1184,40 +1446,169 @@ func (s *OpenAIGatewayService) openAIAdvancedSchedulerSettingRepo() SettingRepos return s.rateLimitService.settingService.settingRepo } -func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerEnabled(ctx context.Context) bool { +func (s *OpenAIGatewayService) openAIAdvancedSchedulerRuntimeSettings(ctx context.Context) openAIAdvancedSchedulerRuntimeSettings { if cached, ok := openAIAdvancedSchedulerSettingCache.Load().(*cachedOpenAIAdvancedSchedulerSetting); ok && cached != nil { if time.Now().UnixNano() < cached.expiresAt { - return cached.enabled + return openAIAdvancedSchedulerRuntimeSettings{ + enabled: cached.enabled, + stickyWeightedEnabled: cached.stickyWeightedEnabled, + subscriptionPriorityEnabled: cached.subscriptionPriorityEnabled, + lbTopKOverride: cached.lbTopKOverride, + weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(cached.weightOverrides), + } } } result, _, _ := openAIAdvancedSchedulerSettingSF.Do(openAIAdvancedSchedulerSettingKey, func() (any, error) { if cached, ok := openAIAdvancedSchedulerSettingCache.Load().(*cachedOpenAIAdvancedSchedulerSetting); ok && cached != nil { if time.Now().UnixNano() < cached.expiresAt { - return cached.enabled, nil + return openAIAdvancedSchedulerRuntimeSettings{ + enabled: cached.enabled, + stickyWeightedEnabled: cached.stickyWeightedEnabled, + subscriptionPriorityEnabled: cached.subscriptionPriorityEnabled, + lbTopKOverride: cached.lbTopKOverride, + weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(cached.weightOverrides), + }, nil } } enabled := false + stickyWeightedEnabled := false + subscriptionPriorityEnabled := false + lbTopKOverride := 0 + weightOverrides := map[string]float64{} if repo := s.openAIAdvancedSchedulerSettingRepo(); repo != nil { dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIAdvancedSchedulerSettingDBTimeout) defer cancel() - value, err := repo.GetValue(dbCtx, openAIAdvancedSchedulerSettingKey) - if err == nil { - enabled = strings.EqualFold(strings.TrimSpace(value), "true") + if values, err := repo.GetMultiple(dbCtx, openAIAdvancedSchedulerRuntimeSettingKeys()); err == nil { + enabled = strings.EqualFold(strings.TrimSpace(values[openAIAdvancedSchedulerSettingKey]), "true") + stickyWeightedEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]), "true") + subscriptionPriorityEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]), "true") + lbTopKOverride = parsePositiveIntOverride(values[SettingKeyOpenAIAdvancedSchedulerLBTopK]) + weightOverrides = parseOpenAIAdvancedSchedulerWeightOverrides(values) + } else { + // 批量读取失败时逐键降级，覆盖全部键（含 TopK/权重），避免只加载布尔开关 + // 而静默丢弃管理员配置的覆盖值；降级状态会被缓存一个 TTL，必须留痕。 + slog.Warn("openai_advanced_scheduler_settings_batch_load_failed", "error", err) + fallbackValues := make(map[string]string) + for _, key := range openAIAdvancedSchedulerRuntimeSettingKeys() { + if value, valueErr := repo.GetValue(dbCtx, key); valueErr == nil { + fallbackValues[key] = value + } + } + enabled = strings.EqualFold(strings.TrimSpace(fallbackValues[openAIAdvancedSchedulerSettingKey]), "true") + stickyWeightedEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]), "true") + subscriptionPriorityEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]), "true") + lbTopKOverride = parsePositiveIntOverride(fallbackValues[SettingKeyOpenAIAdvancedSchedulerLBTopK]) + weightOverrides = parseOpenAIAdvancedSchedulerWeightOverrides(fallbackValues) } } openAIAdvancedSchedulerSettingCache.Store(&cachedOpenAIAdvancedSchedulerSetting{ - enabled: enabled, - expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(), + enabled: enabled, + stickyWeightedEnabled: stickyWeightedEnabled, + subscriptionPriorityEnabled: subscriptionPriorityEnabled, + lbTopKOverride: lbTopKOverride, + weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(weightOverrides), + expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(), }) - return enabled, nil + return openAIAdvancedSchedulerRuntimeSettings{ + enabled: enabled, + stickyWeightedEnabled: stickyWeightedEnabled, + subscriptionPriorityEnabled: subscriptionPriorityEnabled, + lbTopKOverride: lbTopKOverride, + weightOverrides: weightOverrides, + }, nil }) - enabled, _ := result.(bool) - return enabled + settings, _ := result.(openAIAdvancedSchedulerRuntimeSettings) + return settings +} + +func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerEnabled(ctx context.Context) bool { + return s.openAIAdvancedSchedulerRuntimeSettings(ctx).enabled +} + +func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx context.Context) bool { + settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx) + return settings.enabled && settings.stickyWeightedEnabled +} + +func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerSubscriptionPriorityEnabled(ctx context.Context) bool { + settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx) + return settings.enabled && settings.subscriptionPriorityEnabled +} + +func openAIAdvancedSchedulerRuntimeSettingKeys() []string { + keys := []string{ + openAIAdvancedSchedulerSettingKey, + SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled, + SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled, + SettingKeyOpenAIAdvancedSchedulerLBTopK, + } + for _, spec := range openAIAdvancedSchedulerWeightOverrideSpecs() { + keys = append(keys, spec.key) + } + return keys +} + +type openAIAdvancedSchedulerWeightOverrideSpec struct { + key string + name string +} + +func openAIAdvancedSchedulerWeightOverrideSpecs() []openAIAdvancedSchedulerWeightOverrideSpec { + return []openAIAdvancedSchedulerWeightOverrideSpec{ + {key: SettingKeyOpenAIAdvancedSchedulerWeightPriority, name: "priority"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightLoad, name: "load"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightQueue, name: "queue"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightErrorRate, name: "error_rate"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightTTFT, name: "ttft"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightReset, name: "reset"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom, name: "quota_headroom"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse, name: "previous_response"}, + {key: SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky, name: "session_sticky"}, + } +} + +func parsePositiveIntOverride(raw string) int { + raw = strings.TrimSpace(raw) + if raw == "" { + return 0 + } + value, err := strconv.Atoi(raw) + if err != nil || value <= 0 { + return 0 + } + return value +} + +func parseOpenAIAdvancedSchedulerWeightOverrides(values map[string]string) map[string]float64 { + overrides := map[string]float64{} + for _, spec := range openAIAdvancedSchedulerWeightOverrideSpecs() { + raw := strings.TrimSpace(values[spec.key]) + if raw == "" { + continue + } + value, err := strconv.ParseFloat(raw, 64) + if err != nil || value < 0 || math.IsNaN(value) || math.IsInf(value, 0) { + continue + } + overrides[spec.name] = value + } + return overrides +} + +func cloneOpenAIAdvancedSchedulerWeightOverrides(in map[string]float64) map[string]float64 { + if len(in) == 0 { + return nil + } + out := make(map[string]float64, len(in)) + for key, value := range in { + out[key] = value + } + return out } func (s *OpenAIGatewayService) getOpenAIAccountScheduler(ctx context.Context) OpenAIAccountScheduler { @@ -1253,9 +1644,12 @@ func (s *OpenAIGatewayService) SelectAccountWithScheduler( requiredTransport OpenAIUpstreamTransport, requireCompact bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { - return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, "", "", requireCompact, PlatformOpenAI) + return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, "", "", requireCompact, PlatformOpenAI, false) } +// SelectAccountWithSchedulerForCapability 按能力要求调度账号。 +// previousResponseCanMove 表示首包 input 可自行重建工具续链，previous_response_id 允许跨账号迁移 +// （粘性加权模式下改为加权偏好而非硬粘连）。 func (s *OpenAIGatewayService) SelectAccountWithSchedulerForCapability( ctx context.Context, groupID *int64, @@ -1266,13 +1660,14 @@ func (s *OpenAIGatewayService) SelectAccountWithSchedulerForCapability( requiredTransport OpenAIUpstreamTransport, requiredCapability OpenAIEndpointCapability, requireCompact bool, + previousResponseCanMove bool, platformOverride ...string, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { platform := PlatformOpenAI if len(platformOverride) > 0 { platform = platformOverride[0] } - return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, "", requireCompact, platform) + return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, "", requireCompact, platform, previousResponseCanMove) } func (s *OpenAIGatewayService) SelectAccountWithSchedulerForImages( @@ -1283,13 +1678,13 @@ func (s *OpenAIGatewayService) SelectAccountWithSchedulerForImages( excludedIDs map[int64]struct{}, requiredCapability OpenAIImagesCapability, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { - selection, decision, err := s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", requiredCapability, false, PlatformOpenAI) + selection, decision, err := s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", requiredCapability, false, PlatformOpenAI, false) if err == nil && selection != nil && selection.Account != nil { return selection, decision, nil } // 如果要求 native 能力（如指定了模型）但没有可用的 APIKey 账号，回退到 basic（OAuth 账号） if requiredCapability == OpenAIImagesCapabilityNative { - return s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", OpenAIImagesCapabilityBasic, false, PlatformOpenAI) + return s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", OpenAIImagesCapabilityBasic, false, PlatformOpenAI, false) } return selection, decision, err } @@ -1306,6 +1701,7 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( requiredImageCapability OpenAIImagesCapability, requireCompact bool, platform string, + previousResponseCanMove bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { ctx = s.withOpenAIQuotaAutoPauseContext(ctx) platform = normalizeOpenAICompatiblePlatform(platform) @@ -1378,13 +1774,23 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( stickyAccountID = accountID } } + stickyWeighted := s.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx) + subscriptionPriority := s.isOpenAIAdvancedSchedulerSubscriptionPriorityEnabled(ctx) + stickyPreviousAccountID := int64(0) + if stickyWeighted && previousResponseCanMove && strings.TrimSpace(previousResponseID) != "" && platform == PlatformOpenAI { + stickyPreviousAccountID = s.ResolveAccountIDByPreviousResponseIDForScheduler(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) + } return scheduler.Select(ctx, OpenAIAccountScheduleRequest{ GroupID: groupID, Platform: platform, SessionHash: sessionHash, StickyAccountID: stickyAccountID, + StickyPreviousAccountID: stickyPreviousAccountID, + StickyWeighted: stickyWeighted, + SubscriptionPriority: subscriptionPriority, PreviousResponseID: previousResponseID, + PreviousResponseCanMove: previousResponseCanMove, RequestedModel: requestedModel, RequiredTransport: requiredTransport, RequiredCapability: requiredCapability, @@ -1473,6 +1879,20 @@ func (s *OpenAIGatewayService) openAIWSLBTopK() int { return 7 } +func (s *OpenAIGatewayService) openAIWSLBTopKForRequest(ctx context.Context) int { + base := s.openAIWSLBTopK() + settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx) + // DB 覆盖值与 stickyWeighted/subscriptionPriority 一样受总开关门控： + // 关闭高级调度器后所有调用方（含管理页分数快照）都应回到配置/默认行为。 + if !settings.enabled { + return base + } + if settings.lbTopKOverride > 0 { + return settings.lbTopKOverride + } + return base +} + func (s *OpenAIGatewayService) openAIStickyEscapeConfig() openAIStickyEscapeConfig { if s != nil && s.cfg != nil { cfg := s.cfg.Gateway.OpenAIScheduler @@ -1514,6 +1934,8 @@ func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedul TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT, Reset: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Reset, QuotaHeadroom: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom, + Previous: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.PreviousResponse, + SessionSticky: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky, } } return GatewayOpenAIWSSchedulerScoreWeightsView{ @@ -1524,9 +1946,50 @@ func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedul TTFT: 0.5, Reset: 0.0, QuotaHeadroom: 0.0, + Previous: 5.0, + SessionSticky: 3.0, } } +func (s *OpenAIGatewayService) openAIWSSchedulerWeightsForRequest(ctx context.Context) GatewayOpenAIWSSchedulerScoreWeightsView { + weights := s.openAIWSSchedulerWeights() + settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx) + // 同 openAIWSLBTopKForRequest：总开关关闭时不应用 DB 覆盖值。 + if !settings.enabled { + return weights + } + return applyOpenAIAdvancedSchedulerWeightOverrides(weights, settings.weightOverrides) +} + +func applyOpenAIAdvancedSchedulerWeightOverrides( + weights GatewayOpenAIWSSchedulerScoreWeightsView, + overrides map[string]float64, +) GatewayOpenAIWSSchedulerScoreWeightsView { + for key, value := range overrides { + switch key { + case "priority": + weights.Priority = value + case "load": + weights.Load = value + case "queue": + weights.Queue = value + case "error_rate": + weights.ErrorRate = value + case "ttft": + weights.TTFT = value + case "reset": + weights.Reset = value + case "quota_headroom": + weights.QuotaHeadroom = value + case "previous_response": + weights.Previous = value + case "session_sticky": + weights.SessionSticky = value + } + } + return weights +} + type GatewayOpenAIWSSchedulerScoreWeightsView struct { Priority float64 Load float64 @@ -1536,6 +1999,149 @@ type GatewayOpenAIWSSchedulerScoreWeightsView struct { // Reset 倾向「会话窗口最早重置」的账号；0 表示关闭（默认）。 Reset float64 QuotaHeadroom float64 + Previous float64 + SessionSticky float64 +} + +type OpenAIAccountSchedulerScoreSnapshot struct { + BaseScore float64 + StickyScore float64 + StickyScoreInfinity bool + StickyWeightedEnabled bool +} + +func (s *RateLimitService) BuildOpenAIAccountSchedulerScoreSnapshot( + ctx context.Context, + accounts []*Account, + loadMap map[int64]*AccountLoadInfo, +) map[int64]OpenAIAccountSchedulerScoreSnapshot { + gateway := &OpenAIGatewayService{cfg: nil, rateLimitService: s} + if s != nil { + gateway.cfg = s.cfg + } + return buildOpenAIAccountSchedulerScoreSnapshot(accounts, loadMap, gateway.openAIWSSchedulerWeightsForRequest(ctx), gateway.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx)) +} + +func BuildOpenAIAccountSchedulerScoreSnapshot( + accounts []*Account, + loadMap map[int64]*AccountLoadInfo, +) map[int64]OpenAIAccountSchedulerScoreSnapshot { + gateway := &OpenAIGatewayService{} + return buildOpenAIAccountSchedulerScoreSnapshot(accounts, loadMap, gateway.openAIWSSchedulerWeights(), false) +} + +func buildOpenAIAccountSchedulerScoreSnapshot( + accounts []*Account, + loadMap map[int64]*AccountLoadInfo, + weights GatewayOpenAIWSSchedulerScoreWeightsView, + stickyWeightedEnabled bool, +) map[int64]OpenAIAccountSchedulerScoreSnapshot { + if len(accounts) == 0 { + return nil + } + candidates := make([]openAIAccountCandidateScore, 0, len(accounts)) + for _, account := range accounts { + if account == nil { + continue + } + loadInfo := loadMap[account.ID] + if loadInfo == nil { + loadInfo = &AccountLoadInfo{AccountID: account.ID} + } + candidates = append(candidates, openAIAccountCandidateScore{ + account: account, + loadInfo: loadInfo, + errorRate: 0, + ttft: 0, + hasTTFT: false, + }) + } + if len(candidates) == 0 { + return nil + } + + minPriority, maxPriority := openAIAccountSchedulingPriority(candidates[0].account), openAIAccountSchedulingPriority(candidates[0].account) + maxWaiting := 1 + for i := range candidates { + candidate := &candidates[i] + candidate.priority = openAIAccountSchedulingPriority(candidate.account) + if candidate.priority < minPriority { + minPriority = candidate.priority + } + if candidate.priority > maxPriority { + maxPriority = candidate.priority + } + if candidate.loadInfo.WaitingCount > maxWaiting { + maxWaiting = candidate.loadInfo.WaitingCount + } + } + + minResetRemaining, maxResetRemaining := 0.0, 0.0 + hasResetSample := false + now := time.Now() + if weights.Reset > 0 { + for _, candidate := range candidates { + end := candidate.account.SessionWindowEnd + if end == nil || !now.Before(*end) { + continue + } + remaining := end.Sub(now).Seconds() + if !hasResetSample { + minResetRemaining, maxResetRemaining = remaining, remaining + hasResetSample = true + continue + } + if remaining < minResetRemaining { + minResetRemaining = remaining + } + if remaining > maxResetRemaining { + maxResetRemaining = remaining + } + } + } + + result := make(map[int64]OpenAIAccountSchedulerScoreSnapshot, len(candidates)) + for _, candidate := range candidates { + priorityFactor := 1.0 + if maxPriority > minPriority { + priorityFactor = 1 - float64(candidate.priority-minPriority)/float64(maxPriority-minPriority) + } + loadFactor := 1 - clamp01(float64(candidate.loadInfo.LoadRate)/100.0) + queueFactor := 1 - clamp01(float64(candidate.loadInfo.WaitingCount)/float64(maxWaiting)) + errorFactor := 1.0 + ttftFactor := 0.5 + resetFactor := 0.0 + if weights.Reset > 0 && hasResetSample { + if end := candidate.account.SessionWindowEnd; end != nil && now.Before(*end) { + if maxResetRemaining > minResetRemaining { + resetFactor = 1 - clamp01((end.Sub(now).Seconds()-minResetRemaining)/(maxResetRemaining-minResetRemaining)) + } else { + resetFactor = 1 + } + } + } + quotaHeadroomFactor := 0.0 + if weights.QuotaHeadroom > 0 { + quotaHeadroomFactor = openAIQuotaHeadroomFactor(candidate.account, now) + } + baseScore := weights.Priority*priorityFactor + + weights.Load*loadFactor + + weights.Queue*queueFactor + + weights.ErrorRate*errorFactor + + weights.TTFT*ttftFactor + + weights.Reset*resetFactor + + weights.QuotaHeadroom*quotaHeadroomFactor + score := OpenAIAccountSchedulerScoreSnapshot{ + BaseScore: baseScore, + StickyWeightedEnabled: stickyWeightedEnabled, + StickyScoreInfinity: !stickyWeightedEnabled, + } + if stickyWeightedEnabled { + score.StickyScore = baseScore + weights.Previous + weights.SessionSticky + } + result[candidate.account.ID] = score + } + return result } func openAIQuotaHeadroomFactor(account *Account, now time.Time) float64 { diff --git a/backend/internal/service/openai_account_scheduler_reset_test.go b/backend/internal/service/openai_account_scheduler_reset_test.go index 50155f20d6..f1270c8caa 100644 --- a/backend/internal/service/openai_account_scheduler_reset_test.go +++ b/backend/internal/service/openai_account_scheduler_reset_test.go @@ -1,6 +1,7 @@ package service import ( + "context" "testing" "time" @@ -49,7 +50,7 @@ func TestBuildOpenAIAccountLoadPlan_ResetWeightPrefersSoonestReset(t *testing.T) } sched := openAIResetTestScheduler(5.0) - plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + plan := sched.buildOpenAIAccountLoadPlan(context.Background(), OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) scores := openAIPlanScores(plan) require.Greater(t, scores[2], scores[1], "重置时间最早的账号（ID=2）得分更高") } @@ -65,7 +66,7 @@ func TestBuildOpenAIAccountLoadPlan_ResetWeightZeroNoEffect(t *testing.T) { } sched := openAIResetTestScheduler(0.0) - plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + plan := sched.buildOpenAIAccountLoadPlan(context.Background(), OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) scores := openAIPlanScores(plan) require.Equal(t, scores[1], scores[2], "Reset 权重为 0 时两账号得分相同") } @@ -80,7 +81,7 @@ func TestBuildOpenAIAccountLoadPlan_ResetWeightIgnoresNilWindow(t *testing.T) { } sched := openAIResetTestScheduler(5.0) - plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + plan := sched.buildOpenAIAccountLoadPlan(context.Background(), OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) scores := openAIPlanScores(plan) require.Greater(t, scores[2], scores[1], "拥有活跃窗口的账号得分高于无窗口账号") } @@ -161,7 +162,7 @@ func TestBuildOpenAIAccountLoadPlan_QuotaHeadroomPrefersHigher7dRemaining(t *tes } sched := openAIQuotaHeadroomTestScheduler(1.0) - plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + plan := sched.buildOpenAIAccountLoadPlan(context.Background(), OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) scores := openAIPlanScores(plan) require.Greater(t, scores[2], scores[1], "7d 剩余额度更高的账号得分应更高") } @@ -190,7 +191,7 @@ func TestBuildOpenAIAccountLoadPlan_QuotaHeadroomZeroNoEffect(t *testing.T) { } sched := openAIResetTestScheduler(0) - plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + plan := sched.buildOpenAIAccountLoadPlan(context.Background(), OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) scores := openAIPlanScores(plan) require.Equal(t, scores[1], scores[2], "quota_headroom 权重为 0 时不应影响打分") } diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 0255dbbde3..2ff7c25e5d 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -183,6 +183,17 @@ func newSchedulerTestOpenAIWSV2Config() *config.Config { return cfg } +func newSchedulerTestSubscriptionPriorityConfig() *config.Config { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0 + return cfg +} + type openAIAdvancedSchedulerSettingRepoStub struct { values map[string]string } @@ -210,8 +221,14 @@ func (s *openAIAdvancedSchedulerSettingRepoStub) Set(context.Context, string, st panic("unexpected call to Set") } -func (s *openAIAdvancedSchedulerSettingRepoStub) GetMultiple(context.Context, []string) (map[string]string, error) { - panic("unexpected call to GetMultiple") +func (s *openAIAdvancedSchedulerSettingRepoStub) GetMultiple(_ context.Context, keys []string) (map[string]string, error) { + result := make(map[string]string, len(keys)) + for _, key := range keys { + if value, err := s.GetValue(context.Background(), key); err == nil { + result[key] = value + } + } + return result, nil } func (s *openAIAdvancedSchedulerSettingRepoStub) SetMultiple(context.Context, map[string]string) error { @@ -226,7 +243,7 @@ func (s *openAIAdvancedSchedulerSettingRepoStub) Delete(context.Context, string) panic("unexpected call to Delete") } -func newOpenAIAdvancedSchedulerRateLimitService(enabled string) *RateLimitService { +func newOpenAIAdvancedSchedulerRateLimitService(enabled string, values ...string) *RateLimitService { resetOpenAIAdvancedSchedulerSettingCacheForTest() repo := &openAIAdvancedSchedulerSettingRepoStub{ values: map[string]string{}, @@ -234,6 +251,12 @@ func newOpenAIAdvancedSchedulerRateLimitService(enabled string) *RateLimitServic if enabled != "" { repo.values[openAIAdvancedSchedulerSettingKey] = enabled } + if len(values) > 0 && values[0] != "" { + repo.values[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled] = values[0] + } + if len(values) > 1 && values[1] != "" { + repo.values[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled] = values[1] + } return &RateLimitService{ settingService: NewSettingService(repo, &config.Config{}), } @@ -266,6 +289,45 @@ func (s *openAISnapshotCacheStub) GetAccount(ctx context.Context, accountID int6 return &cloned, nil } +func TestOpenAIGatewayService_OpenAIAdvancedSchedulerRuntimeSettings_DBOverridesConfig(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + defer resetOpenAIAdvancedSchedulerSettingCacheForTest() + + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 11 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{ + Priority: 1, + Load: 2, + Queue: 3, + ErrorRate: 4, + TTFT: 5, + Reset: 6, + QuotaHeadroom: 7, + PreviousResponse: 8, + SessionSticky: 9, + } + repo := &openAIAdvancedSchedulerSettingRepoStub{ + values: map[string]string{ + openAIAdvancedSchedulerSettingKey: "true", + SettingKeyOpenAIAdvancedSchedulerLBTopK: "3", + SettingKeyOpenAIAdvancedSchedulerWeightPriority: "2.5", + SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: "12", + }, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + rateLimitService: &RateLimitService{settingService: NewSettingService(repo, cfg)}, + } + + ctx := context.Background() + require.Equal(t, 3, svc.openAIWSLBTopKForRequest(ctx)) + weights := svc.openAIWSSchedulerWeightsForRequest(ctx) + require.Equal(t, 2.5, weights.Priority) + require.Equal(t, 2.0, weights.Load) + require.Equal(t, 12.0, weights.Previous) + require.Equal(t, 9.0, weights.SessionSticky) +} + func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabledUsesLegacyLoadAwareness(t *testing.T) { resetOpenAIAdvancedSchedulerSettingCacheForTest() @@ -467,6 +529,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_Embeddi OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityEmbeddings, false, + false, ) require.NoError(t, err) require.NotNil(t, selection) @@ -510,6 +573,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_AllowsG OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityChatCompletions, false, + false, PlatformGrok, ) require.NoError(t, err) @@ -584,6 +648,248 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_EnabledUsesAdvancedPrev require.True(t, decision.StickyPreviousHit) } +func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedSessionInTopKUsesStickyFirst(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(101071) + accounts := []Account{ + { + ID: 37101, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 100, + GroupIDs: []int64{groupID}, + }, + { + ID: 37102, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + }, + } + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 2 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky = 3 + cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{ + "openai:session_hash_weighted_topk": 37101, + }} + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: cache, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"), + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "session_hash_weighted_topk", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + false, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(37101), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.True(t, decision.StickySessionHit) + require.Equal(t, 2, decision.TopK) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedPreviousRequiresMovableContext(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(101072) + accounts := []Account{ + { + ID: 37111, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 100, + GroupIDs: []int64{groupID}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + }, + }, + { + ID: 37112, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + }, + }, + } + cfg := newSchedulerTestOpenAIWSV2Config() + cfg.Gateway.OpenAIWS.LBTopK = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5 + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"), + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + store := svc.getOpenAIWSStateStore() + require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_weighted_unmovable", 37111, time.Hour)) + + selection, decision, err := svc.SelectAccountWithSchedulerForCapability( + ctx, + &groupID, + "resp_weighted_unmovable", + "", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + OpenAIEndpointCapabilityChatCompletions, + false, + false, + PlatformOpenAI, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(37111), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer) + require.True(t, decision.StickyPreviousHit) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + + selection, decision, err = svc.SelectAccountWithSchedulerForCapability( + ctx, + &groupID, + "resp_weighted_unmovable", + "", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + OpenAIEndpointCapabilityChatCompletions, + false, + true, + PlatformOpenAI, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(37112), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.False(t, decision.StickyPreviousHit) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseCompactUnsupportedDeletesBinding(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(101073) + accounts := []Account{ + { + ID: 37121, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + "openai_compact_mode": OpenAICompactModeForceOff, + }, + }, + { + ID: 37122, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 10, + GroupIDs: []int64{groupID}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + "openai_compact_mode": OpenAICompactModeForceOn, + }, + }, + } + cfg := newSchedulerTestOpenAIWSV2Config() + cfg.Gateway.OpenAIWS.LBTopK = 2 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5 + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + store := svc.getOpenAIWSStateStore() + require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_compact_unsupported", 37121, time.Hour)) + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "resp_compact_unsupported", + "", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + true, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(37122), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.False(t, decision.StickyPreviousHit) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + + accountID, err := store.GetResponseAccount(ctx, groupID, "resp_compact_unsupported") + require.NoError(t, err) + require.Zero(t, accountID) +} + func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkipsChatOnlyAccount(t *testing.T) { resetOpenAIAdvancedSchedulerSettingCacheForTest() @@ -635,6 +941,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkips OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityEmbeddings, false, + false, ) require.NoError(t, err) require.NotNil(t, selection) @@ -708,6 +1015,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkips OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityEmbeddings, false, + false, ) require.NoError(t, err) require.NotNil(t, selection) @@ -1560,6 +1868,217 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeDisa require.True(t, decision.StickySessionHit) } +func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityChoosesSubscriptionPoolFirst(t *testing.T) { + ctx := context.Background() + groupID := int64(10120) + accounts := []Account{ + { + ID: 21601, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 10, + GroupIDs: []int64{groupID}, + Credentials: map[string]any{"plan_type": "plus"}, + }, + { + ID: 21602, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + }, + } + concurrencyCache := schedulerTestConcurrencyCache{ + acquireResults: map[int64]bool{21601: true, 21602: true}, + loadMap: map[int64]*AccountLoadInfo{ + 21601: {AccountID: 21601, LoadRate: 90, WaitingCount: 1}, + 21602: {AccountID: 21602, LoadRate: 0, WaitingCount: 0}, + }, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: newSchedulerTestSubscriptionPriorityConfig(), + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_first", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(21601), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.Equal(t, 1, decision.TopK) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityFallsBackWhenSubscriptionFull(t *testing.T) { + ctx := context.Background() + groupID := int64(10121) + accounts := []Account{ + { + ID: 21611, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + Credentials: map[string]any{"plan_type": "team"}, + }, + { + ID: 21612, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 9, + GroupIDs: []int64{groupID}, + }, + } + concurrencyCache := schedulerTestConcurrencyCache{ + acquireResults: map[int64]bool{21611: false, 21612: true}, + loadMap: map[int64]*AccountLoadInfo{ + 21611: {AccountID: 21611, LoadRate: 0, WaitingCount: 0}, + 21612: {AccountID: 21612, LoadRate: 90, WaitingCount: 1}, + }, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: newSchedulerTestSubscriptionPriorityConfig(), + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_fallback", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(21612), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.True(t, selection.Acquired) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityDisabledUsesScore(t *testing.T) { + ctx := context.Background() + groupID := int64(10122) + accounts := []Account{ + { + ID: 21621, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 10, + GroupIDs: []int64{groupID}, + Credentials: map[string]any{"plan_type": "pro"}, + }, + { + ID: 21622, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + }, + } + concurrencyCache := schedulerTestConcurrencyCache{ + acquireResults: map[int64]bool{21621: true, 21622: true}, + loadMap: map[int64]*AccountLoadInfo{ + 21621: {AccountID: 21621, LoadRate: 90, WaitingCount: 1}, + 21622: {AccountID: 21622, LoadRate: 0, WaitingCount: 0}, + }, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: newSchedulerTestSubscriptionPriorityConfig(), + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "false"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_disabled", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(21622), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesAccountPriorityWithinGroupPool(t *testing.T) { + ctx := context.Background() + groupID := int64(10123) + accounts := []Account{ + { + ID: 21631, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 1, + AccountGroups: []AccountGroup{ + {AccountID: 21631, GroupID: groupID, Priority: 100}, + }, + GroupIDs: []int64{groupID}, + }, + { + ID: 21632, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 100000, + AccountGroups: []AccountGroup{ + {AccountID: 21632, GroupID: groupID, Priority: 1}, + }, + GroupIDs: []int64{groupID}, + }, + } + cfg := newSchedulerTestSubscriptionPriorityConfig() + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0 + svc := &OpenAIGatewayService{ + accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + + selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_group_priority", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(21631), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + func TestDefaultOpenAIAccountScheduler_ShouldEscapeStickyAccount_ThresholdBoundary(t *testing.T) { stats := newOpenAIAccountRuntimeStats() accountID := int64(21501) @@ -2421,3 +2940,141 @@ func TestDefaultOpenAIAccountScheduler_IsAccountTransportCompatible_Branches(t * func int64PtrForTest(v int64) *int64 { return &v } + +func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedFallbackSkipsOutOfGroupStickyAccount(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(101081) + otherGroupID := int64(101082) + accounts := []Account{ + { + ID: 38001, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 10, + GroupIDs: []int64{groupID}, + }, + { + // 会话粘连绑定指向的账号已被移出请求分组（绑定 TTL 内账号改组的场景）。 + ID: 38002, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{otherGroupID}, + }, + } + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 2 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky = 3 + cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{ + "openai:session_weighted_out_of_group": 38002, + }} + concurrencyCache := schedulerTestConcurrencyCache{ + acquireResults: map[int64]bool{38001: false, 38002: true}, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}}, + cache: cache, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "session_weighted_out_of_group", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + false, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + // 组内唯一候选 38001 满并发：必须返回其等待计划，绝不能把请求泄漏到组外的粘连账号 38002。 + require.Equal(t, int64(38001), selection.Account.ID) + require.False(t, selection.Acquired) + require.NotNil(t, selection.WaitPlan) + require.Equal(t, int64(38001), selection.WaitPlan.AccountID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + // 失效的粘连绑定应被清理，避免后续请求反复走同一条泄漏路径。 + require.Positive(t, cache.deletedSessions["openai:session_weighted_out_of_group"]) +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityWaitsOnBusySubscriptionWhenRegularUnusable(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(101091) + accounts := []Account{ + { + // 订阅账号：支持 compact，但并发已满（busy-but-waitable）。 + ID: 38011, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + Credentials: map[string]any{"plan_type": "team"}, + Extra: map[string]any{"openai_compact_supported": true}, + }, + { + // 常规账号：明确不支持 compact，无法服务本次请求。 + ID: 38012, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 9, + GroupIDs: []int64{groupID}, + Extra: map[string]any{"openai_compact_supported": false}, + }, + } + concurrencyCache := schedulerTestConcurrencyCache{ + acquireResults: map[int64]bool{38011: false, 38012: true}, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: newSchedulerTestSubscriptionPriorityConfig(), + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "session_subscription_wait", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + true, + ) + // 常规池无可用候选时，忙碌的订阅账号应产生等待计划，而不是直接返回 no available accounts。 + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(38011), selection.Account.ID) + require.False(t, selection.Acquired) + require.NotNil(t, selection.WaitPlan) + require.Equal(t, int64(38011), selection.WaitPlan.AccountID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) +} diff --git a/backend/internal/service/openai_apikey_responses_probe.go b/backend/internal/service/openai_apikey_responses_probe.go index 64f963ab9b..10cf050029 100644 --- a/backend/internal/service/openai_apikey_responses_probe.go +++ b/backend/internal/service/openai_apikey_responses_probe.go @@ -149,6 +149,9 @@ func (s *AccountTestService) ProbeOpenAIAPIKeyResponsesSupport(ctx context.Conte req.Header.Set("Authorization", "Bearer "+apiKey) req.Header.Set("Accept", "application/json") + // 账号级请求头覆写：能力探测与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) + proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() diff --git a/backend/internal/service/openai_client_restriction_detector.go b/backend/internal/service/openai_client_restriction_detector.go index abca88ce66..8a8097c879 100644 --- a/backend/internal/service/openai_client_restriction_detector.go +++ b/backend/internal/service/openai_client_restriction_detector.go @@ -1,6 +1,7 @@ package service import ( + "fmt" "net/http" "github.com/Wei-Shaw/sub2api/internal/config" @@ -8,6 +9,11 @@ import ( "github.com/gin-gonic/gin" ) +// CodexOfficialClientsOnlyMessage 是 codex_cli_only 拒绝时面向客户端的通用兜底文案。 +// 仅当拒绝原因不是「可解析版本但越界」（VersionTooLow/VersionTooHigh）时使用： +// 未命中官方/黑名单/缺指纹/版本无法识别都沿用这句（避免向伪装客户端泄露门控细节）。 +const CodexOfficialClientsOnlyMessage = "This account only allows Codex official clients" + const ( // CodexClientRestrictionReasonDisabled 表示账号未开启 codex_cli_only。 CodexClientRestrictionReasonDisabled = "codex_cli_only_disabled" @@ -51,6 +57,13 @@ type CodexClientRestrictionDetectionResult struct { Enabled bool Matched bool Reason string + // DetectedVersion 是从官方 UA 解析出的 Codex 引擎版本；仅在版本门拒绝 + // (VersionTooLow / VersionTooHigh) 时填充，供面向客户端的差异化文案使用。 + DetectedVersion string + // MinCodexVersion 是触发 VersionTooLow 时的最低要求版本（来自策略快照）。 + MinCodexVersion string + // MaxCodexVersion 是触发 VersionTooHigh 时的最高允许版本（来自策略快照）。 + MaxCodexVersion string } // CodexClientRestrictionDetector 定义 codex_cli_only 统一检测入口。 @@ -127,10 +140,22 @@ func (d *OpenAICodexClientRestrictionDetector) Detect(c *gin.Context, account *A return CodexClientRestrictionDetectionResult{Enabled: true, Matched: false, Reason: CodexClientRestrictionReasonVersionUndetectable} } if policy.MinCodexVersion != "" && CompareVersions(ver, policy.MinCodexVersion) < 0 { - return CodexClientRestrictionDetectionResult{Enabled: true, Matched: false, Reason: CodexClientRestrictionReasonVersionTooLow} + return CodexClientRestrictionDetectionResult{ + Enabled: true, + Matched: false, + Reason: CodexClientRestrictionReasonVersionTooLow, + DetectedVersion: ver, + MinCodexVersion: policy.MinCodexVersion, + } } if policy.MaxCodexVersion != "" && CompareVersions(ver, policy.MaxCodexVersion) > 0 { - return CodexClientRestrictionDetectionResult{Enabled: true, Matched: false, Reason: CodexClientRestrictionReasonVersionTooHigh} + return CodexClientRestrictionDetectionResult{ + Enabled: true, + Matched: false, + Reason: CodexClientRestrictionReasonVersionTooHigh, + DetectedVersion: ver, + MaxCodexVersion: policy.MaxCodexVersion, + } } } @@ -145,3 +170,22 @@ func (d *OpenAICodexClientRestrictionDetector) Detect(c *gin.Context, account *A return CodexClientRestrictionDetectionResult{Enabled: true, Matched: true, Reason: reason} } + +// CodexClientRestrictionMessage 把检测结果映射为面向客户端的 403 文案。 +// 仅版本越界（VersionTooLow/VersionTooHigh）给出带实际版本号与边界的差异化提示—— +// 这类请求其实已被识别为官方 Codex（命中官方 UA/originator），再回「只允许官方客户端」会误导； +// 其余拒绝原因统一沿用通用兜底句，不暴露门控细节。 +func CodexClientRestrictionMessage(r CodexClientRestrictionDetectionResult) string { + switch r.Reason { + case CodexClientRestrictionReasonVersionTooLow: + return fmt.Sprintf( + "Your Codex version (%s) is below the minimum required version (%s). Please update Codex.", + r.DetectedVersion, r.MinCodexVersion) + case CodexClientRestrictionReasonVersionTooHigh: + return fmt.Sprintf( + "Your Codex version (%s) exceeds the maximum allowed version (%s). Please downgrade Codex to %s or lower.", + r.DetectedVersion, r.MaxCodexVersion, r.MaxCodexVersion) + default: + return CodexOfficialClientsOnlyMessage + } +} diff --git a/backend/internal/service/openai_client_restriction_detector_test.go b/backend/internal/service/openai_client_restriction_detector_test.go index 291c79f6bf..6c79432ae2 100644 --- a/backend/internal/service/openai_client_restriction_detector_test.go +++ b/backend/internal/service/openai_client_restriction_detector_test.go @@ -284,6 +284,66 @@ func TestDetect_V3_AppServerAndSkipAndVersionScope(t *testing.T) { }) } +func TestDetect_VersionGateCarriesVersionFields(t *testing.T) { + gin.SetMode(gin.TestMode) + d := NewOpenAICodexClientRestrictionDetector(nil) + acc := func() *Account { + return &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Extra: map[string]any{"codex_cli_only": true}} + } + + t.Run("版本太低：携带 DetectedVersion + MinCodexVersion", func(t *testing.T) { + c := newCodexDetectorTestContext("codex_cli_rs/0.39.0 (x)", "") + r := d.Detect(c, acc(), CodexRestrictionPolicy{MinCodexVersion: "0.42.0"}, nil) + require.False(t, r.Matched) + require.Equal(t, CodexClientRestrictionReasonVersionTooLow, r.Reason) + require.Equal(t, "0.39.0", r.DetectedVersion) + require.Equal(t, "0.42.0", r.MinCodexVersion) + }) + + t.Run("版本太高：携带 DetectedVersion + MaxCodexVersion", func(t *testing.T) { + c := newCodexDetectorTestContext("codex_cli_rs/0.45.0 (x)", "") + r := d.Detect(c, acc(), CodexRestrictionPolicy{MaxCodexVersion: "0.42.0"}, nil) + require.False(t, r.Matched) + require.Equal(t, CodexClientRestrictionReasonVersionTooHigh, r.Reason) + require.Equal(t, "0.45.0", r.DetectedVersion) + require.Equal(t, "0.42.0", r.MaxCodexVersion) + }) +} + +func TestCodexClientRestrictionMessage(t *testing.T) { + t.Run("版本太低：带实际版本与最低要求", func(t *testing.T) { + msg := CodexClientRestrictionMessage(CodexClientRestrictionDetectionResult{ + Reason: CodexClientRestrictionReasonVersionTooLow, + DetectedVersion: "0.39.0", + MinCodexVersion: "0.42.0", + }) + require.Equal(t, "Your Codex version (0.39.0) is below the minimum required version (0.42.0). Please update Codex.", msg) + }) + + t.Run("版本太高：带实际版本与最高允许", func(t *testing.T) { + msg := CodexClientRestrictionMessage(CodexClientRestrictionDetectionResult{ + Reason: CodexClientRestrictionReasonVersionTooHigh, + DetectedVersion: "0.45.0", + MaxCodexVersion: "0.42.0", + }) + require.Equal(t, "Your Codex version (0.45.0) exceeds the maximum allowed version (0.42.0). Please downgrade Codex to 0.42.0 or lower.", msg) + }) + + t.Run("无法识别版本：保持原通用句", func(t *testing.T) { + msg := CodexClientRestrictionMessage(CodexClientRestrictionDetectionResult{ + Reason: CodexClientRestrictionReasonVersionUndetectable, + }) + require.Equal(t, "This account only allows Codex official clients", msg) + }) + + t.Run("未命中官方：保持原通用句", func(t *testing.T) { + msg := CodexClientRestrictionMessage(CodexClientRestrictionDetectionResult{ + Reason: CodexClientRestrictionReasonNotMatchedUA, + }) + require.Equal(t, "This account only allows Codex official clients", msg) + }) +} + func TestDetect_EngineFingerprintSignals(t *testing.T) { gin.SetMode(gin.TestMode) det := NewOpenAICodexClientRestrictionDetector(&config.Config{}) diff --git a/backend/internal/service/openai_codex_function_call_id_test.go b/backend/internal/service/openai_codex_function_call_id_test.go new file mode 100644 index 0000000000..2ac59e0520 --- /dev/null +++ b/backend/internal/service/openai_codex_function_call_id_test.go @@ -0,0 +1,136 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// TestFilterCodexInput_StripsFunctionCallItemID_WhenPreservingReferences +// verifies that function_call items with non-fc id (e.g. item_*) have their +// id stripped even when PreserveReferences is true. OpenAI upstream requires +// function_call ids to begin with "fc" and rejects item_* with 400: +// "Expected an ID that begins with 'fc'." (#3785) +func TestFilterCodexInput_StripsFunctionCallItemID_WhenPreservingReferences(t *testing.T) { + input := []any{ + map[string]any{ + "type": "function_call", + "id": "item_A9v0SNfS3VaLrfX0j3y4xhyK", + "call_id": "fc_abc123", + "name": "bash", + }, + map[string]any{ + "type": "function_call_output", + "call_id": "fc_abc123", + "output": "done", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 2) + + fc, ok := filtered[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "function_call", fc["type"]) + _, hasID := fc["id"] + require.False(t, hasID, "item_* id should be stripped from function_call") + require.Equal(t, "fc_abc123", fc["call_id"], "call_id must be preserved") + require.Equal(t, "bash", fc["name"]) +} + +// TestFilterCodexInput_KeepsFcID_WhenPreservingReferences +// verifies that function_call items with a valid fc* id are kept when +// PreserveReferences is true. +func TestFilterCodexInput_KeepsFcID_WhenPreservingReferences(t *testing.T) { + input := []any{ + map[string]any{ + "type": "function_call", + "id": "fc_validID123", + "call_id": "fc_validID123", + "name": "bash", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 1) + fc, ok := filtered[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "fc_validID123", fc["id"], "valid fc* id must be preserved") +} + +// TestFilterCodexInput_StripsItemIDFromAllToolCallInputTypes verifies that +// item_* ids are stripped from all call-input types (not output types). +func TestFilterCodexInput_StripsItemIDFromAllToolCallInputTypes(t *testing.T) { + types := []string{"function_call", "tool_call", "local_shell_call", "custom_tool_call", "mcp_tool_call"} + + for _, typ := range types { + input := []any{ + map[string]any{ + "type": typ, + "id": "item_xyz", + "call_id": "fc_001", + "name": "tool", + }, + } + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + require.Len(t, filtered, 1) + item, ok := filtered[0].(map[string]any) + require.True(t, ok) + _, hasID := item["id"] + require.False(t, hasID, "item_* id should be stripped from %s", typ) + } +} + +// TestFilterCodexInput_OutputTypeKeepsItemID ensures tool-output items +// (e.g. function_call_output) keep their id — only call-input types have +// the fc* constraint. +func TestFilterCodexInput_OutputTypeKeepsItemID(t *testing.T) { + input := []any{ + map[string]any{ + "type": "function_call_output", + "id": "o1", + "call_id": "fc_abc", + "output": "done", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 1) + out, ok := filtered[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "o1", out["id"], "output item id should be preserved") +} + +// TestFilterCodexInput_NonToolCallItemKeepsID ensures non-tool-call items +// (e.g. message) still keep their id when PreserveReferences is true. +func TestFilterCodexInput_NonToolCallItemKeepsID(t *testing.T) { + input := []any{ + map[string]any{ + "type": "message", + "id": "item_msg_001", + "role": "user", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 1) + msg, ok := filtered[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "item_msg_001", msg["id"], "non-tool-call items keep their id in preserve mode") +} diff --git a/backend/internal/service/openai_codex_transform.go b/backend/internal/service/openai_codex_transform.go index c33ac97f0f..44c420edbe 100644 --- a/backend/internal/service/openai_codex_transform.go +++ b/backend/internal/service/openai_codex_transform.go @@ -9,6 +9,9 @@ import ( ) var codexModelMap = map[string]string{ + "gpt-5.6-sol": "gpt-5.6-sol", + "gpt-5.6-terra": "gpt-5.6-terra", + "gpt-5.6-luna": "gpt-5.6-luna", "gpt-5.5": "gpt-5.5", "gpt-5.5-pro": "gpt-5.5-pro", "codex-auto-review": "codex-auto-review", @@ -54,6 +57,9 @@ var codexVersionModelPrefixes = []struct { prefix string target string }{ + {prefix: "gpt-5.6-sol", target: "gpt-5.6-sol"}, + {prefix: "gpt-5.6-terra", target: "gpt-5.6-terra"}, + {prefix: "gpt-5.6-luna", target: "gpt-5.6-luna"}, {prefix: "gpt-5.3-codex-spark", target: "gpt-5.3-codex-spark"}, {prefix: "gpt-5.3-codex", target: "gpt-5.3-codex"}, {prefix: "gpt-5.4-mini", target: "gpt-5.4-mini"}, @@ -1297,6 +1303,16 @@ func filterCodexInputWithOptions(input []any, opts codexInputFilterOptions) []an if !opts.PreserveReferences { ensureCopy() delete(newItem, "id") + } else if isCodexToolCallInputType(typ) { + // 续链模式下保留 id 以维持上下文引用，但 function_call 等 + // call-input 类 item 的 id 必须以 "fc" 开头（上游校验 + // "Expected an ID that begins with 'fc'"）。item_* 形式的 id + // 来自客户端回放，需要删除。 + // 注意：function_call_output 等 output 类的 id 无此约束，不动。 + if id, ok := m["id"].(string); ok && id != "" && !strings.HasPrefix(id, "fc") { + ensureCopy() + delete(newItem, "id") + } } filtered = append(filtered, newItem) @@ -1322,6 +1338,22 @@ func isCodexToolCallItemType(typ string) bool { } } +// isCodexToolCallInputType 仅匹配 call-input 类型（不含 output），这些类型的 +// id 必须以 "fc" 开头，上游会校验 "Expected an ID that begins with 'fc'."。 +func isCodexToolCallInputType(typ string) bool { + switch typ { + case "function_call", + "tool_call", + "local_shell_call", + "tool_search_call", + "custom_tool_call", + "mcp_tool_call": + return true + default: + return false + } +} + func codexInputItemRequiresName(typ string) bool { switch strings.TrimSpace(typ) { case "function_call", "custom_tool_call", "mcp_tool_call": diff --git a/backend/internal/service/openai_embeddings.go b/backend/internal/service/openai_embeddings.go index 0fb3fff1f7..fb2dc5ccbb 100644 --- a/backend/internal/service/openai_embeddings.go +++ b/backend/internal/service/openai_embeddings.go @@ -82,6 +82,9 @@ func (s *OpenAIGatewayService) ForwardEmbeddings( upstreamReq.Header.Set("user-agent", customUA) } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效） + account.ApplyHeaderOverrides(upstreamReq.Header) + proxyURL := "" if account.Proxy != nil { proxyURL = account.Proxy.URL() diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 6bcb6718b7..023440f94f 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -105,6 +105,33 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( return nil, policyErr } upstreamBody = updatedBody + + // Grok Composer does not accept image_url parts directly, but Grok Build + // can describe the images first. Bridge only this exact failure mode. + token, tokenKind, err := s.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + if strings.TrimSpace(token) == "" { + return nil, fmt.Errorf("account %d missing %s credential", account.ID, tokenKind) + } + + var bridgeUsage OpenAIUsage + if account.Platform == PlatformGrok { + bridgedBody, usage, bridged, bridgeErr := s.bridgeGrokComposerImageInputs(ctx, c, account, upstreamBody, token) + if bridgeErr != nil { + var failoverErr *UpstreamFailoverError + if !errors.As(bridgeErr, &failoverErr) && c != nil && c.Writer != nil && !c.Writer.Written() { + writeChatCompletionsError(c, http.StatusBadGateway, "upstream_error", bridgeErr.Error()) + } + return nil, bridgeErr + } + if bridged { + upstreamBody = bridgedBody + addOpenAIUsage(&bridgeUsage, usage) + } + } + if clientStream { var usageErr error upstreamBody, usageErr = ensureOpenAIChatStreamUsage(upstreamBody) @@ -122,14 +149,6 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( ) // 5. Build upstream request - token, tokenKind, err := s.GetAccessToken(ctx, account) - if err != nil { - return nil, err - } - if strings.TrimSpace(token) == "" { - return nil, fmt.Errorf("account %d missing %s credential", account.ID, tokenKind) - } - targetURL, err := s.rawChatCompletionsURL(account) if err != nil { return nil, err @@ -166,6 +185,9 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( upstreamReq.Header.Set("user-agent", "sub2api-grok/1.0") } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效） + account.ApplyHeaderOverrides(upstreamReq.Header) + // 6. Send request proxyURL := "" if account.Proxy != nil { @@ -242,10 +264,17 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } // 8. Forward response + var result *OpenAIForwardResult + var forwardErr error if clientStream { - return s.streamRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime, len(body)) + result, forwardErr = s.streamRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime, len(body)) + } else { + result, forwardErr = s.bufferRawChatCompletions(c, resp, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime) } - return s.bufferRawChatCompletions(c, resp, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime) + if result != nil { + addOpenAIUsage(&result.Usage, bridgeUsage) + } + return result, forwardErr } func (s *OpenAIGatewayService) rawChatCompletionsURL(account *Account) (string, error) { diff --git a/backend/internal/service/openai_gateway_count_tokens.go b/backend/internal/service/openai_gateway_count_tokens.go index 4a01b143e9..7518a6073a 100644 --- a/backend/internal/service/openai_gateway_count_tokens.go +++ b/backend/internal/service/openai_gateway_count_tokens.go @@ -231,6 +231,9 @@ func (s *OpenAIGatewayService) buildInputTokensUpstreamRequest( } } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op） + account.ApplyHeaderOverrides(req.Header) + return req, nil } diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 4961b9c589..4a0ad06d46 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -10,12 +10,18 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) +const ( + grokComposerImageBridgeVisionModel = "grok-build-0.1" + grokComposerImageBridgeMaxOutputTokens = 512 +) + func (s *OpenAIGatewayService) forwardGrokResponses( ctx context.Context, c *gin.Context, @@ -309,6 +315,303 @@ func shouldDropGrokToolChoice(toolChoice gjson.Result, tools []json.RawMessage) return false } +func (s *OpenAIGatewayService) bridgeGrokComposerImageInputs( + ctx context.Context, + c *gin.Context, + account *Account, + body []byte, + token string, +) ([]byte, OpenAIUsage, bool, error) { + if !shouldBridgeGrokComposerImageInputs(body) { + return body, OpenAIUsage{}, false, nil + } + + var reqBody map[string]any + if err := json.Unmarshal(body, &reqBody); err != nil { + return body, OpenAIUsage{}, false, fmt.Errorf("parse grok composer image bridge request: %w", err) + } + + imageURLs := collectGrokComposerImageURLs(reqBody) + if len(imageURLs) == 0 { + return body, OpenAIUsage{}, false, nil + } + + descriptions := make([]string, 0, len(imageURLs)) + var bridgeUsage OpenAIUsage + for index, imageURL := range imageURLs { + description, usage, err := s.describeGrokComposerImage(ctx, c, account, token, imageURL, index+1) + if err != nil { + return body, bridgeUsage, false, err + } + descriptions = append(descriptions, description) + addOpenAIUsage(&bridgeUsage, usage) + } + + if !rewriteGrokComposerImagesAsText(reqBody, descriptions) { + return body, bridgeUsage, false, nil + } + bridgedBody, err := marshalOpenAIUpstreamJSON(reqBody) + if err != nil { + return body, bridgeUsage, false, fmt.Errorf("serialize grok composer image bridge request: %w", err) + } + return bridgedBody, bridgeUsage, true, nil +} + +func shouldBridgeGrokComposerImageInputs(body []byte) bool { + if len(body) == 0 || !isGrokComposerModel(gjson.GetBytes(body, "model").String()) { + return false + } + messages := gjson.GetBytes(body, "messages") + if !messages.Exists() { + return false + } + return openAIJSONValueMayContainImageInput(messages) +} + +func isGrokComposerModel(model string) bool { + model = strings.TrimSpace(strings.ToLower(model)) + if model == "" { + return false + } + if strings.Contains(model, "/") { + parts := strings.Split(model, "/") + model = strings.TrimSpace(parts[len(parts)-1]) + } + return strings.Contains(model, "composer") +} + +func collectGrokComposerImageURLs(reqBody map[string]any) []string { + messages, ok := reqBody["messages"].([]any) + if !ok { + return nil + } + + var imageURLs []string + for _, msg := range messages { + msgMap, ok := msg.(map[string]any) + if !ok { + continue + } + parts, ok := msgMap["content"].([]any) + if !ok { + continue + } + for _, part := range parts { + if imageURL := grokComposerImageURLFromPart(part); imageURL != "" { + imageURLs = append(imageURLs, imageURL) + } + } + } + return imageURLs +} + +func grokComposerImageURLFromPart(part any) string { + partMap, ok := part.(map[string]any) + if !ok { + return "" + } + if strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"]))) != "image_url" { + return "" + } + switch imageURL := partMap["image_url"].(type) { + case string: + return normalizeGrokComposerImageURL(imageURL) + case map[string]any: + raw, _ := imageURL["url"].(string) + return normalizeGrokComposerImageURL(raw) + default: + return "" + } +} + +func normalizeGrokComposerImageURL(raw string) string { + trimmed := strings.TrimSpace(raw) + if trimmed == "" || isEmptyBase64DataURI(trimmed) { + return "" + } + return trimmed +} + +func (s *OpenAIGatewayService) describeGrokComposerImage( + ctx context.Context, + c *gin.Context, + account *Account, + token string, + imageURL string, + index int, +) (string, OpenAIUsage, error) { + body, err := buildGrokComposerImageDescriptionBody(imageURL, index) + if err != nil { + return "", OpenAIUsage{}, err + } + + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token) + releaseUpstreamCtx() + if err != nil { + return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err) + } + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + if err != nil { + return "", OpenAIUsage{}, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody)) + if upstreamMsg == "" { + upstreamMsg = fmt.Sprintf("xAI image bridge upstream returned status %d", resp.StatusCode) + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")), + Kind: "failover", + Message: upstreamMsg, + }) + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + if s.shouldFailoverUpstreamError(resp.StatusCode) { + return "", OpenAIUsage{}, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg) + } + + s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil) + if err != nil { + return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err) + } + + var parsed apicompat.ResponsesResponse + if err := json.Unmarshal(respBody, &parsed); err != nil { + return "", OpenAIUsage{}, fmt.Errorf("parse grok composer image bridge response: %w", err) + } + description := strings.TrimSpace(grokResponsesOutputText(&parsed)) + if description == "" { + return "", copyOpenAIUsageFromResponsesUsage(parsed.Usage), fmt.Errorf("grok composer image bridge returned empty description") + } + return description, copyOpenAIUsageFromResponsesUsage(parsed.Usage), nil +} + +func buildGrokComposerImageDescriptionBody(imageURL string, index int) ([]byte, error) { + prompt := fmt.Sprintf("Describe image %d in concise, factual text for a downstream coding/composer model. Include visible text, UI elements, diagrams, errors, and spatial relationships. Do not mention that you are an image analysis bridge.", index) + req := map[string]any{ + "model": grokComposerImageBridgeVisionModel, + "stream": false, + "store": false, + "max_output_tokens": grokComposerImageBridgeMaxOutputTokens, + "input": []any{ + map[string]any{ + "type": "message", + "role": "user", + "content": []any{ + map[string]any{"type": "input_text", "text": prompt}, + map[string]any{"type": "input_image", "image_url": imageURL}, + }, + }, + }, + } + return marshalOpenAIUpstreamJSON(req) +} + +func grokResponsesOutputText(resp *apicompat.ResponsesResponse) string { + if resp == nil { + return "" + } + var parts []string + for _, output := range resp.Output { + for _, content := range output.Content { + if content.Type == "output_text" || content.Type == "text" || content.Type == "input_text" { + if text := strings.TrimSpace(content.Text); text != "" { + parts = append(parts, text) + } + } + } + } + return strings.Join(parts, "\n\n") +} + +func rewriteGrokComposerImagesAsText(reqBody map[string]any, descriptions []string) bool { + messages, ok := reqBody["messages"].([]any) + if !ok { + return false + } + + imageIndex := 0 + changed := false + for _, msg := range messages { + msgMap, ok := msg.(map[string]any) + if !ok { + continue + } + parts, ok := msgMap["content"].([]any) + if !ok { + continue + } + var textParts []string + messageChanged := false + for _, part := range parts { + if imageURL := grokComposerImageURLFromPart(part); imageURL != "" { + if imageIndex < len(descriptions) { + textParts = append(textParts, fmt.Sprintf("Image %d description: %s", imageIndex+1, strings.TrimSpace(descriptions[imageIndex]))) + } + imageIndex++ + messageChanged = true + continue + } + if text := grokComposerTextFromPart(part); text != "" { + textParts = append(textParts, text) + } + } + if messageChanged { + msgMap["content"] = strings.Join(textParts, "\n\n") + changed = true + } + } + return changed +} + +func grokComposerTextFromPart(part any) string { + partMap, ok := part.(map[string]any) + if !ok { + return "" + } + partType := strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"]))) + switch partType { + case "text", "input_text": + text, _ := partMap["text"].(string) + return strings.TrimSpace(text) + default: + return "" + } +} + +func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) { + if dst == nil { + return + } + dst.InputTokens += usage.InputTokens + dst.ImageInputTokens += usage.ImageInputTokens + dst.OutputTokens += usage.OutputTokens + dst.CacheCreationInputTokens += usage.CacheCreationInputTokens + dst.CacheReadInputTokens += usage.CacheReadInputTokens + dst.ImageOutputTokens += usage.ImageOutputTokens +} + func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) { targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) if err != nil { diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index eae424ce9a..d012ab533a 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -651,6 +651,76 @@ func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *te require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey]) } +func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-composer-2.5-fast","messages":[{"role":"system","content":"You are concise."},{"role":"user","content":[{"type":"text","text":"What is shown?"},{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}]}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 55, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{55: account}, + }, + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}, "xai-request-id": []string{"vision-req"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_vision","object":"response","model":"grok-build-0.1","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"A small diagram with ABC letters."}]}],"usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18}}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "X-Request-Id": []string{"composer-req"}, + "X-Ratelimit-Limit-Requests": []string{"10"}, + "X-Ratelimit-Remaining-Requests": []string{"9"}, + "X-Ratelimit-Limit-Tokens": []string{"1000"}, + "X-Ratelimit-Remaining-Tokens": []string{"980"}, + }, + Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl_composer","object":"chat.completion","model":"grok-composer-2.5-fast","choices":[{"index":0,"message":{"role":"assistant","content":"It shows ABC."},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`)), + }, + }} + svc := &OpenAIGatewayService{ + cfg: rawChatCompletionsTestConfig(), + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.requests, 2) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.requests[0].URL.String()) + require.Equal(t, "grok-build-0.1", gjson.GetBytes(upstream.bodies[0], "model").String()) + require.Equal(t, "input_image", gjson.GetBytes(upstream.bodies[0], "input.0.content.1.type").String()) + require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.requests[1].URL.String()) + require.Equal(t, "grok-composer-2.5-fast", gjson.GetBytes(upstream.bodies[1], "model").String()) + require.False(t, strings.Contains(string(upstream.bodies[1]), "image_url")) + require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "Image 1 description") + require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "A small diagram with ABC letters.") + require.Equal(t, 14, result.Usage.InputTokens) + require.Equal(t, 12, result.Usage.OutputTokens) + require.Equal(t, "It shows ABC.", gjson.Get(recorder.Body.String(), "choices.0.message.content").String()) + require.NotNil(t, repo.updates[55][grokQuotaSnapshotExtraKey]) +} + func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index 8323035ffe..697d89e81c 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -1803,6 +1803,52 @@ func TestOpenAIGatewayServiceRecordUsage_ImageIndependentMultiplierUsesImageRate require.Equal(t, string(BillingModeImage), *usageRepo.lastLog.BillingMode) } +func TestGrokVideoMediaBillingUsesImageRateMultiplier(t *testing.T) { + mediaPrice2K := 0.4 + groupID := int64(126) + + usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} + svc := newOpenAIRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{}, nil) + + err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{ + Result: &OpenAIForwardResult{ + RequestID: "video-request-123", + ResponseID: "video-request-123", + Model: "grok-imagine-video-1.5", + BillingModel: "grok-imagine-video-1.5", + // The usage schema has no separate video count; video generation is billed as one media unit. + ImageCount: 1, + ImageSize: ImageBillingSize2K, + Duration: time.Second, + }, + APIKey: &APIKey{ + ID: 10126, + GroupID: i64p(groupID), + Group: &Group{ + ID: groupID, + Platform: PlatformGrok, + RateMultiplier: 0.15, + ImageRateIndependent: true, + ImageRateMultiplier: 0.5, + ImagePrice2K: &mediaPrice2K, + }, + }, + User: &User{ID: 20126}, + Account: &Account{ID: 30126, Platform: PlatformGrok}, + }) + + require.NoError(t, err) + require.NotNil(t, usageRepo.lastLog) + require.Equal(t, "grok-imagine-video-1.5", usageRepo.lastLog.Model) + require.Equal(t, 1, usageRepo.lastLog.ImageCount) + require.Equal(t, ImageBillingSize2K, *usageRepo.lastLog.ImageSize) + require.InDelta(t, 0.4, usageRepo.lastLog.TotalCost, 1e-12) + require.InDelta(t, 0.2, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, 0.5, usageRepo.lastLog.RateMultiplier, 1e-12) + require.NotNil(t, usageRepo.lastLog.BillingMode) + require.Equal(t, string(BillingModeImage), *usageRepo.lastLog.BillingMode) +} + func TestOpenAIGatewayServiceRecordUsage_ChannelImageBillingUsesImageCountAndSharedMultiplier(t *testing.T) { groupID := int64(123) usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} diff --git a/backend/internal/service/openai_gateway_responses_chat_fallback.go b/backend/internal/service/openai_gateway_responses_chat_fallback.go index d33df4c19d..c499bec778 100644 --- a/backend/internal/service/openai_gateway_responses_chat_fallback.go +++ b/backend/internal/service/openai_gateway_responses_chat_fallback.go @@ -138,6 +138,9 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions( upstreamReq.Header.Set("user-agent", customUA) } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效） + account.ApplyHeaderOverrides(upstreamReq.Header) + proxyURL := "" if account.Proxy != nil { proxyURL = account.Proxy.URL() diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index dd3d48aed0..fd104db8b6 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -2617,7 +2617,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco c.JSON(http.StatusForbidden, gin.H{ "error": gin.H{ "type": "forbidden_error", - "message": "This account only allows Codex official clients", + "message": CodexClientRestrictionMessage(restrictionResult), }, }) return nil, errors.New("codex_cli_only restriction: only codex official clients are allowed") @@ -3783,6 +3783,9 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough( req.Header.Set("content-type", "application/json") } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op） + account.ApplyHeaderOverrides(req.Header) + return req, nil } @@ -4568,6 +4571,9 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin. req.Header.Set("content-type", "application/json") } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op） + account.ApplyHeaderOverrides(req.Header) + return req, nil } @@ -5662,7 +5668,13 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r if isEventStreamResponse(resp.Header) { return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel) } - bodyLooksLikeSSE := bytes.Contains(body, []byte("data:")) || bytes.Contains(body, []byte("event:")) + // bodyLooksLikeSSE is a line-level heuristic: real SSE framing requires + // "data:"/"event:" field names at the very start of a physical line. A + // plain bytes.Contains scan would also match ordinary JSON responses + // whose string content merely echoes the literal text "data:" or + // "event:" (e.g. compact tool output), causing those JSON bodies to be + // misrouted into handleSSEToJSON and lose their usage accounting. + bodyLooksLikeSSE := bodyHasSSEFraming(body) // For OAuth accounts, also fall back to a body-content heuristic because // the upstream may omit the Content-Type header while still sending SSE. @@ -5712,6 +5724,22 @@ func isEventStreamResponse(header http.Header) bool { return strings.Contains(contentType, "text/event-stream") } +// bodyHasSSEFraming reports whether body contains genuine SSE framing by +// scanning for physical lines that begin with the "data:" or "event:" +// field names, per the SSE spec. Unlike a raw substring scan, this does not +// match when those strings only appear embedded inside JSON string values +// (e.g. "data: foo" quoted as part of an assistant text field), since such +// occurrences never start a physical line in a valid JSON encoding. +func bodyHasSSEFraming(body []byte) bool { + for _, line := range bytes.Split(body, []byte("\n")) { + line = bytes.TrimRight(line, "\r") + if bytes.HasPrefix(line, []byte("data:")) || bytes.HasPrefix(line, []byte("event:")) { + return true + } + } + return false +} + func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Context, body []byte, originalModel, mappedModel string) (*openaiNonStreamingResult, error) { bodyText := string(body) finalResponse, ok := extractCodexFinalResponse(bodyText) diff --git a/backend/internal/service/openai_gateway_service_codex_cli_only_test.go b/backend/internal/service/openai_gateway_service_codex_cli_only_test.go index 23a1750021..7eb133c125 100644 --- a/backend/internal/service/openai_gateway_service_codex_cli_only_test.go +++ b/backend/internal/service/openai_gateway_service_codex_cli_only_test.go @@ -59,6 +59,52 @@ func TestOpenAIGatewayService_GetCodexClientRestrictionDetector(t *testing.T) { }) } +func TestOpenAIGatewayService_Forward_VersionGateMessage(t *testing.T) { + gin.SetMode(gin.TestMode) + + newCtx := func() (*httptest.ResponseRecorder, *gin.Context) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + return rec, c + } + account := func() *Account { + return &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Extra: map[string]any{"codex_cli_only": true}} + } + body := []byte(`{"model":"gpt-5.1-codex"}`) + + t.Run("版本太低：返回带版本号的差异化文案", func(t *testing.T) { + rec, c := newCtx() + svc := &OpenAIGatewayService{codexDetector: &stubCodexRestrictionDetector{result: CodexClientRestrictionDetectionResult{ + Enabled: true, + Matched: false, + Reason: CodexClientRestrictionReasonVersionTooLow, + DetectedVersion: "0.39.0", + MinCodexVersion: "0.42.0", + }}} + + _, err := svc.Forward(context.Background(), c, account(), body) + require.Error(t, err) + require.Equal(t, http.StatusForbidden, rec.Code) + require.Contains(t, rec.Body.String(), "Your Codex version (0.39.0) is below the minimum required version (0.42.0)") + require.NotContains(t, rec.Body.String(), "This account only allows Codex official clients") + }) + + t.Run("未命中官方：仍返回通用兜底文案", func(t *testing.T) { + rec, c := newCtx() + svc := &OpenAIGatewayService{codexDetector: &stubCodexRestrictionDetector{result: CodexClientRestrictionDetectionResult{ + Enabled: true, + Matched: false, + Reason: CodexClientRestrictionReasonNotMatchedUA, + }}} + + _, err := svc.Forward(context.Background(), c, account(), body) + require.Error(t, err) + require.Equal(t, http.StatusForbidden, rec.Code) + require.Contains(t, rec.Body.String(), "This account only allows Codex official clients") + }) +} + func TestGetAPIKeyIDFromContext(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index c11d78e55c..b3e9889a7d 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -2739,6 +2739,41 @@ func TestHandleNonStreamingResponse_APIKeyFallsBackToSSEBodyWhenContentTypeIsWro require.Equal(t, "hello", gjson.Get(rec.Body.String(), "output.0.content.0.text").String()) } +func TestHandleNonStreamingResponse_OAuthJSONBodyWithDataEventTextKeepsJSONUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil) + + svc := &OpenAIGatewayService{cfg: &config.Config{}} + // Plain JSON compact response whose output text happens to contain the + // literal substrings "data:" and "event:" (e.g. echoing shell/log output). + // This must NOT be misdetected as SSE framing: it has a top-level usage + // object and no upstream text/event-stream Content-Type. + jsonBody := `{"id":"resp_oauth_compact","object":"response","model":"gpt-5.4","status":"completed",` + + `"output":[{"type":"message","content":[{"type":"output_text",` + + `"text":"processing data: 1,2,3 then event: click finished"}]}],` + + `"usage":{"input_tokens":11,"output_tokens":22,"total_tokens":33}}` + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(jsonBody)), + } + account := &Account{ID: 146, Type: AccountTypeOAuth} + + result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 11, result.InputTokens) + require.Equal(t, 22, result.OutputTokens) + // Response must remain the original JSON body (not routed through the SSE + // path, which would rewrite/lose the body or usage). + require.Equal(t, "application/json", rec.Header().Get("Content-Type")) + require.Equal(t, "resp_oauth_compact", gjson.Get(rec.Body.String(), "id").String()) + require.Equal(t, int64(33), gjson.Get(rec.Body.String(), "usage.total_tokens").Int()) + require.Contains(t, rec.Body.String(), "processing data: 1,2,3 then event: click finished") +} + func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 7081653d80..09472fbaf1 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -760,6 +760,8 @@ func (s *OpenAIGatewayService) buildOpenAIImagesRequest( if strings.TrimSpace(contentType) != "" { req.Header.Set("Content-Type", contentType) } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op） + account.ApplyHeaderOverrides(req.Header) return req, nil } diff --git a/backend/internal/service/openai_model_alias.go b/backend/internal/service/openai_model_alias.go index ac2a8cf942..4e3d3b2d9a 100644 --- a/backend/internal/service/openai_model_alias.go +++ b/backend/internal/service/openai_model_alias.go @@ -65,6 +65,12 @@ func normalizeKnownOpenAICodexModel(model string) string { } switch { + case strings.Contains(normalized, "gpt-5.6-sol"): + return "gpt-5.6-sol" + case strings.Contains(normalized, "gpt-5.6-terra"): + return "gpt-5.6-terra" + case strings.Contains(normalized, "gpt-5.6-luna"): + return "gpt-5.6-luna" case strings.Contains(normalized, "gpt-5.5-pro"): return "gpt-5.5-pro" case strings.Contains(normalized, "gpt-5.5"): diff --git a/backend/internal/service/openai_tool_continuation.go b/backend/internal/service/openai_tool_continuation.go index 6515c0c4e5..a507213701 100644 --- a/backend/internal/service/openai_tool_continuation.go +++ b/backend/internal/service/openai_tool_continuation.go @@ -215,6 +215,84 @@ func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValid return result } +// ToolCallOutputContextCoverage 描述 input 中工具输出与可重建上下文的覆盖关系， +// 用于判断剥离 previous_response_id 后上游能否仅凭 input 重建工具续链。 +type ToolCallOutputContextCoverage struct { + HasFunctionCallOutput bool + // ContextCoversAllCallIDs 表示每个工具输出的 call_id 都能在 input 内找到 + // 同 call_id 的工具调用上下文项或同 id 的 item_reference，且不存在缺失 call_id 的输出。 + // 任一输出无法由 input 自身重建时为 false，此时剥离 previous_response_id 会导致 + // 上游以 "No tool call found for function call output" 拒绝请求。 + ContextCoversAllCallIDs bool +} + +// AnalyzeToolCallOutputContextCoverageBytes 全量扫描 input，按 call_id 精确匹配工具输出 +// 与可重建上下文。不能复用 ValidateFunctionCallOutputContextBytes 的 HasToolCallContext： +// 该标志只代表"存在某一个上下文项"，部分覆盖的续链仍会被上游拒绝。 +func AnalyzeToolCallOutputContextCoverageBytes(body []byte) ToolCallOutputContextCoverage { + coverage := ToolCallOutputContextCoverage{} + if len(body) == 0 { + return coverage + } + input := parseRawJSONView(body).Get("input") + if !input.IsArray() { + return coverage + } + + missingCallID := false + var outputCallIDs map[string]struct{} + var contextIDs map[string]struct{} + input.ForEach(func(_, item gjson.Result) bool { + if !item.IsObject() { + return true + } + itemType := item.Get("type").String() + switch { + case isCodexToolCallOutputItemType(itemType): + coverage.HasFunctionCallOutput = true + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID == "" { + missingCallID = true + return true + } + if outputCallIDs == nil { + outputCallIDs = make(map[string]struct{}) + } + outputCallIDs[callID] = struct{}{} + case isCodexToolCallContextItemType(itemType): + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID == "" { + return true + } + if contextIDs == nil { + contextIDs = make(map[string]struct{}) + } + contextIDs[callID] = struct{}{} + case itemType == "item_reference": + idValue := strings.TrimSpace(item.Get("id").String()) + if idValue == "" { + return true + } + if contextIDs == nil { + contextIDs = make(map[string]struct{}) + } + contextIDs[idValue] = struct{}{} + } + return true + }) + + if !coverage.HasFunctionCallOutput || missingCallID { + return coverage + } + for callID := range outputCallIDs { + if _, ok := contextIDs[callID]; !ok { + return coverage + } + } + coverage.ContextCoversAllCallIDs = true + return coverage +} + // ValidateFunctionCallOutputContext 为 handler 提供低开销校验结果： // 1) 无工具输出直接返回 // 2) 若已存在工具调用上下文则提前返回 diff --git a/backend/internal/service/openai_tool_continuation_test.go b/backend/internal/service/openai_tool_continuation_test.go index 4610652b6c..569d89eff0 100644 --- a/backend/internal/service/openai_tool_continuation_test.go +++ b/backend/internal/service/openai_tool_continuation_test.go @@ -184,3 +184,109 @@ func TestValidateFunctionCallOutputContextBytesMatchesMapValidation(t *testing.T }) } } + +func TestAnalyzeToolCallOutputContextCoverageBytes(t *testing.T) { + cases := []struct { + name string + body map[string]any + hasOutput bool + coversAllIDs bool + }{ + { + name: "no_input", + body: map[string]any{"model": "gpt-5.1"}, + hasOutput: false, + coversAllIDs: false, + }, + { + name: "no_tool_output", + body: map[string]any{"input": []any{ + map[string]any{"type": "message", "content": "hi"}, + }}, + hasOutput: false, + coversAllIDs: false, + }, + { + name: "all_outputs_covered_by_context", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call", "call_id": "call_a"}, + map[string]any{"type": "function_call_output", "call_id": "call_a"}, + }}, + hasOutput: true, + coversAllIDs: true, + }, + { + name: "all_outputs_covered_by_item_reference", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_a"}, + map[string]any{"type": "item_reference", "id": "call_a"}, + }}, + hasOutput: true, + coversAllIDs: true, + }, + { + // 关键回归用例：input 内存在某一个上下文项，但另一个输出的 call_id + // 只能由上游会话链（previous_response_id）解析——不可剥离。 + name: "partial_coverage_not_movable", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call", "call_id": "call_a"}, + map[string]any{"type": "function_call_output", "call_id": "call_a"}, + map[string]any{"type": "function_call_output", "call_id": "call_b"}, + }}, + hasOutput: true, + coversAllIDs: false, + }, + { + name: "unrelated_context_does_not_cover", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call", "call_id": "call_x"}, + map[string]any{"type": "function_call_output", "call_id": "call_b"}, + }}, + hasOutput: true, + coversAllIDs: false, + }, + { + name: "output_missing_call_id_not_movable", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call", "call_id": "call_a"}, + map[string]any{"type": "function_call_output"}, + map[string]any{"type": "function_call_output", "call_id": "call_a"}, + }}, + hasOutput: true, + coversAllIDs: false, + }, + { + name: "mixed_context_and_reference_cover_all", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call", "call_id": "call_a"}, + map[string]any{"type": "function_call_output", "call_id": "call_a"}, + map[string]any{"type": "function_call_output", "call_id": "call_b"}, + map[string]any{"type": "item_reference", "id": "call_b"}, + }}, + hasOutput: true, + coversAllIDs: true, + }, + { + name: "all_codex_output_types_covered", + body: map[string]any{"input": []any{ + map[string]any{"type": "tool_search_output", "call_id": "call_s"}, + map[string]any{"type": "tool_search_call", "call_id": "call_s"}, + map[string]any{"type": "mcp_tool_call_output", "call_id": "call_m"}, + map[string]any{"type": "mcp_tool_call", "call_id": "call_m"}, + }}, + hasOutput: true, + coversAllIDs: true, + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + bodyBytes, err := json.Marshal(tt.body) + require.NoError(t, err) + + coverage := AnalyzeToolCallOutputContextCoverageBytes(bodyBytes) + require.Equal(t, tt.hasOutput, coverage.HasFunctionCallOutput, "HasFunctionCallOutput") + require.Equal(t, tt.coversAllIDs, coverage.ContextCoversAllCallIDs, "ContextCoversAllCallIDs") + }) + } +} diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index 9261db5544..bbca9776ab 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -1183,6 +1183,10 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders( headers.Set("user-agent", codexCLIUserAgent) } + // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op）。 + // 覆盖所有 WS 模式（ctx_pool/dedicated/passthrough）的握手头。 + account.ApplyHeaderOverrides(headers) + return headers, sessionResolution, nil } @@ -4325,87 +4329,8 @@ func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( if s == nil { return nil, nil } - responseID := strings.TrimSpace(previousResponseID) - if responseID == "" { - return nil, nil - } - store := s.getOpenAIWSStateStore() - if store == nil { - return nil, nil - } - - accountID, err := store.GetResponseAccount(ctx, derefGroupID(groupID), responseID) - if err != nil || accountID <= 0 { - return nil, nil - } - if excludedIDs != nil { - if _, excluded := excludedIDs[accountID]; excluded { - return nil, nil - } - } - - account, err := s.getSchedulableAccount(ctx, accountID) - if err != nil || account == nil { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - // 非 WSv2 场景（如 force_http/全局关闭）不应使用 previous_response_id 粘连， - // 以保持“回滚到 HTTP”后的历史行为一致性。 - if s.getOpenAIWSProtocolResolver().Resolve(account).Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - return nil, nil - } - if shouldClearStickySession(account, requestedModel) || !account.IsOpenAI() || !account.IsSchedulable() { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - if requestedModel != "" && !account.IsModelSupported(requestedModel) { - return nil, nil - } - if !account.SupportsOpenAIEndpointCapability(requiredCapability) { - return nil, nil - } - // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an - // account over its 5h/7d threshold keeps serving the same response chain even though - // normal scheduling skips it. Pause is transient, so fall through to normal scheduling - // without deleting the binding (the window may reset before the next turn). - if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { - return nil, nil - } - if s.schedulerSnapshot != nil && s.accountRepo != nil { - latest, latestErr := s.accountRepo.GetByID(ctx, account.ID) - if latestErr != nil || latest == nil { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - if shouldClearStickySession(latest, requestedModel) || !latest.IsOpenAI() || !latest.IsSchedulable() { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - if requestedModel != "" && !latest.IsModelSupported(requestedModel) { - return nil, nil - } - if !latest.SupportsOpenAIEndpointCapability(requiredCapability) { - return nil, nil - } - if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { - return nil, nil - } - if s.isOpenAIAccountRuntimeBlocked(latest) { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return nil, nil - } - account = latest - } - if requireCompact && openAICompactSupportTier(account) == 0 { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + accountID, account, responseID, store := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) + if accountID <= 0 || account == nil || store == nil { return nil, nil } @@ -4439,6 +4364,117 @@ func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( return nil, nil } +func (s *OpenAIGatewayService) ResolveAccountIDByPreviousResponseIDForScheduler( + ctx context.Context, + groupID *int64, + previousResponseID string, + requestedModel string, + excludedIDs map[int64]struct{}, + requiredCapability OpenAIEndpointCapability, + requireCompact bool, +) int64 { + accountID, _, _, _ := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) + return accountID +} + +func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability( + ctx context.Context, + groupID *int64, + previousResponseID string, + requestedModel string, + excludedIDs map[int64]struct{}, + requiredCapability OpenAIEndpointCapability, + requireCompact bool, +) (int64, *Account, string, OpenAIWSStateStore) { + if s == nil { + return 0, nil, "", nil + } + responseID := strings.TrimSpace(previousResponseID) + if responseID == "" { + return 0, nil, "", nil + } + store := s.getOpenAIWSStateStore() + if store == nil { + return 0, nil, "", nil + } + + accountID, err := store.GetResponseAccount(ctx, derefGroupID(groupID), responseID) + if err != nil || accountID <= 0 { + return 0, nil, "", nil + } + if excludedIDs != nil { + if _, excluded := excludedIDs[accountID]; excluded { + return 0, nil, "", nil + } + } + + account, err := s.getSchedulableAccount(ctx, accountID) + if err != nil || account == nil { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + // 非 WSv2 场景（如 force_http/全局关闭）不应使用 previous_response_id 粘连， + // 以保持“回滚到 HTTP”后的历史行为一致性。 + if s.getOpenAIWSProtocolResolver().Resolve(account).Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { + return 0, nil, "", nil + } + if shouldClearStickySession(account, requestedModel) || !account.IsOpenAI() || !account.IsSchedulable() { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if requestedModel != "" && !account.IsModelSupported(requestedModel) { + return 0, nil, "", nil + } + if !account.SupportsOpenAIEndpointCapability(requiredCapability) { + return 0, nil, "", nil + } + // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an + // account over its 5h/7d threshold keeps serving the same response chain even though + // normal scheduling skips it. Pause is transient, so fall through to normal scheduling + // without deleting the binding (the window may reset before the next turn). + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + return 0, nil, "", nil + } + if s.schedulerSnapshot != nil && s.accountRepo != nil { + latest, latestErr := s.accountRepo.GetByID(ctx, account.ID) + if latestErr != nil || latest == nil { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if shouldClearStickySession(latest, requestedModel) || !latest.IsOpenAI() || !latest.IsSchedulable() { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if requestedModel != "" && !latest.IsModelSupported(requestedModel) { + return 0, nil, "", nil + } + if !latest.SupportsOpenAIEndpointCapability(requiredCapability) { + return 0, nil, "", nil + } + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { + return 0, nil, "", nil + } + if s.isOpenAIAccountRuntimeBlocked(latest) { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + account = latest + } + if requireCompact && openAICompactSupportTier(account) == 0 { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + return accountID, account, responseID, store +} + func classifyOpenAIWSAcquireError(err error) string { if err == nil { return "acquire_conn" diff --git a/backend/internal/service/payment_amounts.go b/backend/internal/service/payment_amounts.go index a7f620d33e..2fd00c5957 100644 --- a/backend/internal/service/payment_amounts.go +++ b/backend/internal/service/payment_amounts.go @@ -16,6 +16,15 @@ func normalizeBalanceRechargeMultiplier(multiplier float64) float64 { return multiplier } +// normalizeSubscriptionUSDToCNYRate 将非法值归一为 0（换算关闭）。 +// 与余额倍率不同，0 是合法状态：表示订阅保持 price 直付的存量行为。 +func normalizeSubscriptionUSDToCNYRate(rate float64) float64 { + if math.IsNaN(rate) || math.IsInf(rate, 0) || rate < 0 { + return 0 + } + return rate +} + func calculateCreditedBalance(paymentAmount, multiplier float64) float64 { return decimal.NewFromFloat(paymentAmount). Mul(decimal.NewFromFloat(normalizeBalanceRechargeMultiplier(multiplier))). diff --git a/backend/internal/service/payment_config_limits.go b/backend/internal/service/payment_config_limits.go index 45b24bfce7..202eea9f26 100644 --- a/backend/internal/service/payment_config_limits.go +++ b/backend/internal/service/payment_config_limits.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" dbent "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/ent/paymentproviderinstance" @@ -31,6 +32,7 @@ func (s *PaymentConfigService) GetAvailableMethodLimits(ctx context.Context) (*M continue } ml := pcAggregateMethodLimits(pt, insts) + ml.DisplayName = s.pcAggregateMethodDisplayName(pt, insts) ml.Currency = currency resp.Methods[ml.PaymentType] = ml } @@ -93,6 +95,7 @@ func (s *PaymentConfigService) GetMethodLimits(ctx context.Context, types []stri continue } ml := pcAggregateMethodLimits(pt, matching) + ml.DisplayName = s.pcAggregateMethodDisplayName(pt, matching) ml.Currency = currency result = append(result, ml) } @@ -163,6 +166,53 @@ func (s *PaymentConfigService) pcInstancePaymentCurrency(inst *dbent.PaymentProv return paymentProviderConfigCurrency(inst.ProviderKey, cfg) } +type easyPayCustomMethodDisplayConfig struct { + Type string `json:"type"` + DisplayName string `json:"displayName"` +} + +func (s *PaymentConfigService) pcAggregateMethodDisplayName(pt string, instances []*dbent.PaymentProviderInstance) string { + pt = strings.TrimSpace(pt) + if pt == "" { + return "" + } + for _, inst := range instances { + displayName := s.pcInstanceEasyPayCustomMethodDisplayName(inst, pt) + if displayName != "" { + return displayName + } + } + return "" +} + +func (s *PaymentConfigService) pcInstanceEasyPayCustomMethodDisplayName(inst *dbent.PaymentProviderInstance, pt string) string { + if inst == nil || inst.ProviderKey != payment.TypeEasyPay { + return "" + } + cfg := map[string]string{} + if s != nil { + decrypted, err := s.decryptConfig(inst.Config) + if err == nil && decrypted != nil { + cfg = decrypted + } + } + raw := strings.TrimSpace(cfg["customMethods"]) + if raw == "" { + return "" + } + + var methods []easyPayCustomMethodDisplayConfig + if err := json.Unmarshal([]byte(raw), &methods); err != nil { + return "" + } + for _, method := range methods { + if strings.TrimSpace(method.Type) == pt { + return strings.TrimSpace(method.DisplayName) + } + } + return "" +} + // pcGroupByPaymentType groups instances by user-facing payment type. // For Stripe providers, ALL sub-types (card, link, alipay, wxpay) map to "stripe" // because the user sees a single "Stripe" button, not individual sub-methods. diff --git a/backend/internal/service/payment_config_limits_test.go b/backend/internal/service/payment_config_limits_test.go index c0aa2b27a5..a70bc90a29 100644 --- a/backend/internal/service/payment_config_limits_test.go +++ b/backend/internal/service/payment_config_limits_test.go @@ -255,6 +255,28 @@ func TestGetAvailableMethodLimitsOmitsMixedCurrencyMethod(t *testing.T) { require.Equal(t, "PAYMENT_METHOD_CURRENCY_CONFLICT", appErr.Reason) } +func TestGetAvailableMethodLimitsIncludesEasyPayCustomMethodDisplayName(t *testing.T) { + ctx := context.Background() + client := newPaymentConfigServiceTestClient(t) + + _, err := client.PaymentProviderInstance.Create(). + SetProviderKey(payment.TypeEasyPay). + SetName("EasyPay Custom"). + SetConfig(`{"customMethods":"[{\"type\":\"ldc\",\"upstreamType\":\"ldc\",\"displayName\":\"LDC Pay\"}]"}`). + SetSupportedTypes("alipay,wxpay,ldc"). + SetEnabled(true). + Save(ctx) + require.NoError(t, err) + + svc := &PaymentConfigService{entClient: client} + resp, err := svc.GetAvailableMethodLimits(ctx) + require.NoError(t, err) + + limits, ok := resp.Methods["ldc"] + require.True(t, ok, "expected custom EasyPay method limits to be visible") + require.Equal(t, "LDC Pay", limits.DisplayName) +} + func TestPcComputeGlobalRange(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/payment_config_providers.go b/backend/internal/service/payment_config_providers.go index 7e92558568..d1bf2de7aa 100644 --- a/backend/internal/service/payment_config_providers.go +++ b/backend/internal/service/payment_config_providers.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log/slog" + "regexp" "strconv" "strings" @@ -185,6 +186,11 @@ func (s *PaymentConfigService) CreateProviderInstance(ctx context.Context, req C if err := validateProviderRequest(req.ProviderKey, req.Name, typesStr); err != nil { return nil, err } + if req.ProviderKey == payment.TypeEasyPay { + if err := validateEasyPayCustomMethods(req.Config, typesStr); err != nil { + return nil, err + } + } if err := s.validateVisibleMethodEnablementConflicts(ctx, 0, req.ProviderKey, typesStr, req.Enabled); err != nil { return nil, err } @@ -217,6 +223,67 @@ func validateProviderRequest(providerKey, name, supportedTypes string) error { return nil } +var easyPayCustomMethodCodePattern = regexp.MustCompile(`^[a-z0-9_-]+$`) + +type easyPayCustomMethodConfig struct { + Type string `json:"type"` + UpstreamType string `json:"upstreamType"` + DisplayName string `json:"displayName"` +} + +func validateEasyPayCustomMethods(config map[string]string, supportedTypes string) error { + if config == nil { + config = map[string]string{} + } + raw := strings.TrimSpace(config["customMethods"]) + methods := make([]easyPayCustomMethodConfig, 0) + if raw != "" { + if err := json.Unmarshal([]byte(raw), &methods); err != nil { + return infraerrors.BadRequest("VALIDATION_ERROR", "customMethods must be a JSON array") + } + } + + customTypes := make(map[string]struct{}, len(methods)) + for _, method := range methods { + method.Type = strings.TrimSpace(method.Type) + method.UpstreamType = strings.TrimSpace(method.UpstreamType) + if method.Type == "" || method.UpstreamType == "" { + return infraerrors.BadRequest("VALIDATION_ERROR", "customMethods upstreamType is required") + } + if !easyPayCustomMethodCodePattern.MatchString(method.Type) { + return infraerrors.BadRequest("VALIDATION_ERROR", "customMethods type may only contain lowercase letters, digits, underscores, and hyphens") + } + if !easyPayCustomMethodCodePattern.MatchString(method.UpstreamType) { + return infraerrors.BadRequest("VALIDATION_ERROR", "customMethods upstreamType may only contain lowercase letters, digits, underscores, and hyphens") + } + if easyPayCustomMethodTypeConflictsWithBuiltin(method.Type) { + return infraerrors.BadRequest("VALIDATION_ERROR", "customMethods type cannot start with alipay or wxpay") + } + if _, exists := customTypes[method.Type]; exists { + return infraerrors.BadRequest("VALIDATION_ERROR", "duplicate customMethods type") + } + customTypes[method.Type] = struct{}{} + } + + for _, supportedType := range splitTypes(supportedTypes) { + supportedType = strings.TrimSpace(supportedType) + if supportedType == "" || supportedType == payment.TypeAlipay || supportedType == payment.TypeWxpay { + continue + } + if !easyPayCustomMethodCodePattern.MatchString(supportedType) { + return infraerrors.BadRequest("VALIDATION_ERROR", fmt.Sprintf("supported EasyPay custom type %s may only contain lowercase letters, digits, underscores, and hyphens", supportedType)) + } + if _, exists := customTypes[supportedType]; !exists { + return infraerrors.BadRequest("VALIDATION_ERROR", fmt.Sprintf("supported EasyPay custom type %s has no customMethods mapping", supportedType)) + } + } + return nil +} + +func easyPayCustomMethodTypeConflictsWithBuiltin(methodType string) bool { + return strings.HasPrefix(methodType, payment.TypeAlipay) || strings.HasPrefix(methodType, payment.TypeWxpay) +} + // UpdateProviderInstance updates a provider instance by ID (patch semantics). // NOTE: This function exceeds 30 lines due to per-field nil-check patch update // boilerplate and pending-order safety checks. @@ -279,6 +346,18 @@ func (s *PaymentConfigService) UpdateProviderInstance(ctx context.Context, id in WithMetadata(map[string]string{"count": strconv.Itoa(count)}) } } + configToValidate := mergedConfig + if configToValidate == nil { + configToValidate, err = s.decryptConfig(current.Config) + if err != nil { + return nil, fmt.Errorf("decrypt existing config: %w", err) + } + } + if current.ProviderKey == payment.TypeEasyPay { + if err := validateEasyPayCustomMethods(configToValidate, nextSupportedTypes); err != nil { + return nil, err + } + } // Validate merged config when the instance will end up enabled. // This surfaces provider-level errors (e.g. wxpay missing certSerial) at save time, // so admins see them in the dialog instead of only when an order is created. @@ -287,13 +366,6 @@ func (s *PaymentConfigService) UpdateProviderInstance(ctx context.Context, id in finalEnabled = *req.Enabled } if finalEnabled { - configToValidate := mergedConfig - if configToValidate == nil { - configToValidate, err = s.decryptConfig(current.Config) - if err != nil { - return nil, fmt.Errorf("decrypt existing config: %w", err) - } - } if err := s.validateProviderConfig(current.ProviderKey, configToValidate); err != nil { return nil, err } diff --git a/backend/internal/service/payment_config_providers_test.go b/backend/internal/service/payment_config_providers_test.go index 43708de73d..74fd2a3467 100644 --- a/backend/internal/service/payment_config_providers_test.go +++ b/backend/internal/service/payment_config_providers_test.go @@ -114,6 +114,92 @@ func TestValidateProviderRequest(t *testing.T) { } } +func TestValidateEasyPayCustomMethods(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + config map[string]string + supportedTypes string + wantErr string + }{ + { + name: "valid custom methods", + config: map[string]string{"customMethods": `[{"type":"ldc","upstreamType":"epay","displayName":"LDC"}]`}, + supportedTypes: "alipay,wxpay,ldc", + }, + { + name: "malformed custom methods json", + config: map[string]string{"customMethods": `not-json`}, + supportedTypes: "alipay,wxpay,ldc", + wantErr: "customMethods must be a JSON array", + }, + { + name: "missing upstream type", + config: map[string]string{"customMethods": `[{"type":"ldc","displayName":"LDC"}]`}, + supportedTypes: "alipay,wxpay,ldc", + wantErr: "customMethods upstreamType is required", + }, + { + name: "duplicate custom type", + config: map[string]string{"customMethods": `[{"type":"ldc","upstreamType":"epay"},{"type":"ldc","upstreamType":"epay2"}]`}, + supportedTypes: "alipay,wxpay,ldc", + wantErr: "duplicate customMethods type", + }, + { + name: "custom type must already be lowercase", + config: map[string]string{"customMethods": `[{"type":"LDC","upstreamType":"epay"}]`}, + supportedTypes: "alipay,wxpay,ldc", + wantErr: "customMethods type may only contain lowercase letters", + }, + { + name: "upstream type must already be lowercase", + config: map[string]string{"customMethods": `[{"type":"ldc","upstreamType":"ALIPAY"}]`}, + supportedTypes: "alipay,wxpay,ldc", + wantErr: "customMethods upstreamType may only contain lowercase letters", + }, + { + name: "custom type uses alipay prefix", + config: map[string]string{"customMethods": `[{"type":"alipay_hk","upstreamType":"hkpay"}]`}, + supportedTypes: "alipay,wxpay,alipay_hk", + wantErr: "customMethods type cannot start with alipay or wxpay", + }, + { + name: "custom type uses wxpay prefix", + config: map[string]string{"customMethods": `[{"type":"wxpay_usdt","upstreamType":"usdt"}]`}, + supportedTypes: "alipay,wxpay,wxpay_usdt", + wantErr: "customMethods type cannot start with alipay or wxpay", + }, + { + name: "supported custom type missing mapping", + config: map[string]string{"customMethods": `[{"type":"ldc","upstreamType":"epay"}]`}, + supportedTypes: "alipay,wxpay,ldc,usdt_trc20", + wantErr: "supported EasyPay custom type usdt_trc20 has no customMethods mapping", + }, + { + name: "supported custom type must already be lowercase", + config: map[string]string{"customMethods": `[{"type":"ldc","upstreamType":"epay"}]`}, + supportedTypes: "alipay,wxpay,LDC", + wantErr: "supported EasyPay custom type LDC may only contain lowercase letters", + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + err := validateEasyPayCustomMethods(tc.config, tc.supportedTypes) + if tc.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + require.Contains(t, err.Error(), tc.wantErr) + }) + } +} + func TestIsSensitiveProviderConfigField(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/payment_config_service.go b/backend/internal/service/payment_config_service.go index 022b1b0156..52a7ddc67e 100644 --- a/backend/internal/service/payment_config_service.go +++ b/backend/internal/service/payment_config_service.go @@ -24,17 +24,20 @@ const ( SettingLoadBalanceStrategy = "LOAD_BALANCE_STRATEGY" SettingBalancePayDisabled = "BALANCE_PAYMENT_DISABLED" SettingBalanceRechargeMult = "BALANCE_RECHARGE_MULTIPLIER" - SettingRechargeFeeRate = "RECHARGE_FEE_RATE" - SettingProductNamePrefix = "PRODUCT_NAME_PREFIX" - SettingProductNameSuffix = "PRODUCT_NAME_SUFFIX" - SettingHelpImageURL = "PAYMENT_HELP_IMAGE_URL" - SettingHelpText = "PAYMENT_HELP_TEXT" - SettingCancelRateLimitOn = "CANCEL_RATE_LIMIT_ENABLED" - SettingCancelRateLimitMax = "CANCEL_RATE_LIMIT_MAX" - SettingCancelWindowSize = "CANCEL_RATE_LIMIT_WINDOW" - SettingCancelWindowUnit = "CANCEL_RATE_LIMIT_UNIT" - SettingCancelWindowMode = "CANCEL_RATE_LIMIT_WINDOW_MODE" - SettingAlipayForceQRCode = "ALIPAY_FORCE_QRCODE" + // SettingSubscriptionUSDToCNYRate 是订阅 CNY 换算汇率（1 USD = X CNY）。 + // 0/未配置 = 关闭换算（订阅按 price 数值直付），显式配置后 CNY 通道订阅按 price × rate 收款。 + SettingSubscriptionUSDToCNYRate = "SUBSCRIPTION_USD_TO_CNY_RATE" + SettingRechargeFeeRate = "RECHARGE_FEE_RATE" + SettingProductNamePrefix = "PRODUCT_NAME_PREFIX" + SettingProductNameSuffix = "PRODUCT_NAME_SUFFIX" + SettingHelpImageURL = "PAYMENT_HELP_IMAGE_URL" + SettingHelpText = "PAYMENT_HELP_TEXT" + SettingCancelRateLimitOn = "CANCEL_RATE_LIMIT_ENABLED" + SettingCancelRateLimitMax = "CANCEL_RATE_LIMIT_MAX" + SettingCancelWindowSize = "CANCEL_RATE_LIMIT_WINDOW" + SettingCancelWindowUnit = "CANCEL_RATE_LIMIT_UNIT" + SettingCancelWindowMode = "CANCEL_RATE_LIMIT_WINDOW_MODE" + SettingAlipayForceQRCode = "ALIPAY_FORCE_QRCODE" ) // Default values for payment configuration settings. @@ -54,13 +57,15 @@ type PaymentConfig struct { EnabledTypes []string `json:"enabled_payment_types"` BalanceDisabled bool `json:"balance_disabled"` BalanceRechargeMultiplier float64 `json:"balance_recharge_multiplier"` - RechargeFeeRate float64 `json:"recharge_fee_rate"` - LoadBalanceStrategy string `json:"load_balance_strategy"` - ProductNamePrefix string `json:"product_name_prefix"` - ProductNameSuffix string `json:"product_name_suffix"` - HelpImageURL string `json:"help_image_url"` - HelpText string `json:"help_text"` - StripePublishableKey string `json:"stripe_publishable_key,omitempty"` + // SubscriptionUSDToCNYRate 为 0 时订阅换算关闭（兼容存量行为）。 + SubscriptionUSDToCNYRate float64 `json:"subscription_usd_to_cny_rate"` + RechargeFeeRate float64 `json:"recharge_fee_rate"` + LoadBalanceStrategy string `json:"load_balance_strategy"` + ProductNamePrefix string `json:"product_name_prefix"` + ProductNameSuffix string `json:"product_name_suffix"` + HelpImageURL string `json:"help_image_url"` + HelpText string `json:"help_text"` + StripePublishableKey string `json:"stripe_publishable_key,omitempty"` // Cancel rate limit settings CancelRateLimitEnabled bool `json:"cancel_rate_limit_enabled"` @@ -84,6 +89,7 @@ type UpdatePaymentConfigRequest struct { EnabledTypes []string `json:"enabled_payment_types"` BalanceDisabled *bool `json:"balance_disabled"` BalanceRechargeMultiplier *float64 `json:"balance_recharge_multiplier"` + SubscriptionUSDToCNYRate *float64 `json:"subscription_usd_to_cny_rate"` RechargeFeeRate *float64 `json:"recharge_fee_rate"` LoadBalanceStrategy *string `json:"load_balance_strategy"` ProductNamePrefix *string `json:"product_name_prefix"` @@ -110,6 +116,7 @@ type UpdatePaymentConfigRequest struct { // MethodLimits holds per-payment-type limits. type MethodLimits struct { PaymentType string `json:"payment_type"` + DisplayName string `json:"display_name,omitempty"` Currency string `json:"currency"` FeeRate float64 `json:"fee_rate"` DailyLimit float64 `json:"daily_limit"` @@ -204,7 +211,7 @@ func (s *PaymentConfigService) GetPaymentConfig(ctx context.Context) (*PaymentCo keys := []string{ SettingPaymentEnabled, SettingMinRechargeAmount, SettingMaxRechargeAmount, SettingDailyRechargeLimit, SettingOrderTimeoutMinutes, SettingMaxPendingOrders, - SettingEnabledPaymentTypes, SettingBalancePayDisabled, SettingBalanceRechargeMult, SettingRechargeFeeRate, SettingLoadBalanceStrategy, + SettingEnabledPaymentTypes, SettingBalancePayDisabled, SettingBalanceRechargeMult, SettingSubscriptionUSDToCNYRate, SettingRechargeFeeRate, SettingLoadBalanceStrategy, SettingProductNamePrefix, SettingProductNameSuffix, SettingHelpImageURL, SettingHelpText, SettingCancelRateLimitOn, SettingCancelRateLimitMax, @@ -233,6 +240,7 @@ func (s *PaymentConfigService) parsePaymentConfig(vals map[string]string) *Payme MaxPendingOrders: pcParseInt(vals[SettingMaxPendingOrders], defaultMaxPendingOrders), BalanceDisabled: vals[SettingBalancePayDisabled] == "true", BalanceRechargeMultiplier: normalizeBalanceRechargeMultiplier(pcParseFloat(vals[SettingBalanceRechargeMult], defaultBalanceRechargeMultiplier)), + SubscriptionUSDToCNYRate: normalizeSubscriptionUSDToCNYRate(pcParseFloat(vals[SettingSubscriptionUSDToCNYRate], 0)), RechargeFeeRate: pcParseFloat(vals[SettingRechargeFeeRate], 0), LoadBalanceStrategy: vals[SettingLoadBalanceStrategy], ProductNamePrefix: vals[SettingProductNamePrefix], @@ -294,6 +302,12 @@ func (s *PaymentConfigService) UpdatePaymentConfig(ctx context.Context, req Upda return infraerrors.BadRequest("INVALID_BALANCE_RECHARGE_MULTIPLIER", "balance recharge multiplier must be greater than 0") } } + if req.SubscriptionUSDToCNYRate != nil { + v := *req.SubscriptionUSDToCNYRate + if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 { + return infraerrors.BadRequest("INVALID_SUBSCRIPTION_USD_TO_CNY_RATE", "subscription USD to CNY rate must be 0 (disabled) or a positive number") + } + } if req.RechargeFeeRate != nil { v := *req.RechargeFeeRate if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 || v > 100 { @@ -313,6 +327,7 @@ func (s *PaymentConfigService) UpdatePaymentConfig(ctx context.Context, req Upda SettingMaxPendingOrders: formatPositiveInt(req.MaxPendingOrders), SettingBalancePayDisabled: formatBoolOrEmpty(req.BalanceDisabled), SettingBalanceRechargeMult: formatPositiveFloat(req.BalanceRechargeMultiplier), + SettingSubscriptionUSDToCNYRate: formatPositiveFloatExact(req.SubscriptionUSDToCNYRate), SettingRechargeFeeRate: formatNonNegativeFloat(req.RechargeFeeRate), SettingLoadBalanceStrategy: derefStr(req.LoadBalanceStrategy), SettingProductNamePrefix: derefStr(req.ProductNamePrefix), @@ -352,6 +367,14 @@ func formatPositiveFloat(v *float64) string { return strconv.FormatFloat(*v, 'f', 2, 64) } +// formatPositiveFloatExact 保留完整精度，用于汇率等对小数位敏感的配置。 +func formatPositiveFloatExact(v *float64) string { + if v == nil || *v <= 0 { + return "" // empty → parsePaymentConfig 视为未配置（换算关闭） + } + return strconv.FormatFloat(*v, 'f', -1, 64) +} + func formatNonNegativeFloat(v *float64) string { if v == nil || *v < 0 { return "" diff --git a/backend/internal/service/payment_config_service_test.go b/backend/internal/service/payment_config_service_test.go index f04f4697b1..bfc69d1705 100644 --- a/backend/internal/service/payment_config_service_test.go +++ b/backend/internal/service/payment_config_service_test.go @@ -187,6 +187,23 @@ func TestParsePaymentConfig(t *testing.T) { } }) + t.Run("custom enabled types are preserved", func(t *testing.T) { + t.Parallel() + vals := map[string]string{ + SettingEnabledPaymentTypes: "alipay,ldc,usdt_trc20", + } + cfg := svc.parsePaymentConfig(vals) + want := []string{"alipay", "ldc", "usdt_trc20"} + if len(cfg.EnabledTypes) != len(want) { + t.Fatalf("EnabledTypes len = %d, want %d (%v)", len(cfg.EnabledTypes), len(want), cfg.EnabledTypes) + } + for i := range want { + if cfg.EnabledTypes[i] != want[i] { + t.Fatalf("EnabledTypes[%d] = %q, want %q (full=%v)", i, cfg.EnabledTypes[i], want[i], cfg.EnabledTypes) + } + } + }) + t.Run("empty enabled types string", func(t *testing.T) { t.Parallel() vals := map[string]string{ diff --git a/backend/internal/service/payment_fulfillment_test.go b/backend/internal/service/payment_fulfillment_test.go index b46d6a1fc8..a8c78d713c 100644 --- a/backend/internal/service/payment_fulfillment_test.go +++ b/backend/internal/service/payment_fulfillment_test.go @@ -602,8 +602,8 @@ func TestExecuteSubscriptionFulfillmentAppliesAffiliateRebate(t *testing.T) { SetUserID(user.ID). SetUserEmail(user.Email). SetUserName(user.Username). - SetAmount(120). - SetPayAmount(120). + SetAmount(9.99). + SetPayAmount(71.36). SetFeeRate(0). SetRechargeCode("PAY-SUB-AFFILIATE"). SetOutTradeNo("sub2_subscription_affiliate"). @@ -636,7 +636,7 @@ func TestExecuteSubscriptionFulfillmentAppliesAffiliateRebate(t *testing.T) { } settingSvc := NewSettingService(&paymentFulfillmentSettingRepoStub{values: map[string]string{ SettingKeyAffiliateEnabled: "true", - SettingKeyAffiliateRebateRate: "20", + SettingKeyAffiliateRebateRate: "15", SettingKeyAffiliateRebateFreezeHours: "0", }}, nil) subRepo := newSubscriptionUserSubRepoStub() @@ -659,7 +659,7 @@ func TestExecuteSubscriptionFulfillmentAppliesAffiliateRebate(t *testing.T) { require.Len(t, affiliateRepo.accrueCalls, 1) require.Equal(t, inviterID, affiliateRepo.accrueCalls[0].inviterID) require.Equal(t, user.ID, affiliateRepo.accrueCalls[0].inviteeUserID) - require.Equal(t, 24.0, affiliateRepo.accrueCalls[0].amount) + require.InDelta(t, 1.4985, affiliateRepo.accrueCalls[0].amount, 0.00000001) require.NotNil(t, affiliateRepo.accrueCalls[0].sourceOrderID) require.Equal(t, order.ID, *affiliateRepo.accrueCalls[0].sourceOrderID) require.Equal(t, 1, subRepo.createCalls) @@ -668,8 +668,8 @@ func TestExecuteSubscriptionFulfillmentAppliesAffiliateRebate(t *testing.T) { Where(paymentauditlog.OrderIDEQ(strconv.FormatInt(order.ID, 10)), paymentauditlog.ActionEQ("AFFILIATE_REBATE_APPLIED")). Only(ctx) require.NoError(t, err) - require.Contains(t, applied.Detail, `"baseAmount":120`) - require.Contains(t, applied.Detail, `"rebateAmount":24`) + require.Contains(t, applied.Detail, `"baseAmount":9.99`) + require.Contains(t, applied.Detail, `"rebateAmount":1.4985`) } func TestExecuteSubscriptionFulfillmentDoesNotDuplicateWorkAfterLegacySuccessAudit(t *testing.T) { diff --git a/backend/internal/service/payment_order.go b/backend/internal/service/payment_order.go index 29fe40b1b6..04feb8002a 100644 --- a/backend/internal/service/payment_order.go +++ b/backend/internal/service/payment_order.go @@ -16,6 +16,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/payment" "github.com/Wei-Shaw/sub2api/internal/payment/provider" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/shopspring/decimal" ) // --- Order Creation --- @@ -67,8 +68,7 @@ func (s *PaymentService) CreateOrder(ctx context.Context, req CreateOrderRequest return nil, err } } - // 订阅套餐 price 是直付价，余额充值倍率只影响余额充值到账，不参与订阅 pay_amount 计算。 - payAmountStr, payAmount, err := calculateCreateOrderPayAmount(limitAmount, feeRate, methodCurrency) + payAmountStr, payAmount, err := calculateCreateOrderPayAmountForOrderType(limitAmount, feeRate, methodCurrency, req.OrderType, cfg.SubscriptionUSDToCNYRate) if err != nil { return nil, err } @@ -84,7 +84,7 @@ func (s *PaymentService) CreateOrder(ctx context.Context, req CreateOrderRequest selectedCurrency = paymentProviderConfigCurrency(sel.ProviderKey, sel.Config) } if selectedCurrency != methodCurrency { - payAmountStr, payAmount, err = calculateCreateOrderPayAmount(limitAmount, feeRate, selectedCurrency) + payAmountStr, payAmount, err = calculateCreateOrderPayAmountForOrderType(limitAmount, feeRate, selectedCurrency, req.OrderType, cfg.SubscriptionUSDToCNYRate) if err != nil { return nil, err } @@ -453,6 +453,7 @@ func (s *PaymentService) invokeProvider(ctx context.Context, order *dbent.Paymen } return nil, classifyCreatePaymentError(req, sel.ProviderKey, err) } + sanitizeCreatePaymentResponseDetails(pr) _, err = s.entClient.PaymentOrder.UpdateOneID(order.ID). SetNillablePaymentTradeNo(psNilIfEmpty(pr.TradeNo)). SetNillablePayURL(psNilIfEmpty(pr.PayURL)). @@ -480,6 +481,22 @@ func (s *PaymentService) invokeProvider(ctx context.Context, order *dbent.Paymen return resp, nil } +func sanitizeCreatePaymentResponseDetails(pr *payment.CreatePaymentResponse) { + if pr == nil { + return + } + pr.TradeNo = removePostgresTextNUL(pr.TradeNo) + pr.PayURL = removePostgresTextNUL(pr.PayURL) + pr.QRCode = removePostgresTextNUL(pr.QRCode) +} + +func removePostgresTextNUL(value string) string { + if !strings.ContainsRune(value, 0) { + return value + } + return strings.ReplaceAll(value, "\x00", "") +} + func buildProviderCreatePaymentRequest(req CreateOrderRequest, sel *payment.InstanceSelection, orderID, amount, subject string) payment.CreatePaymentRequest { return payment.CreatePaymentRequest{ OrderID: orderID, @@ -613,6 +630,28 @@ func calculateCreateOrderPayAmount(limitAmount, feeRate float64, currency string return payAmountStr, payAmount, nil } +func calculateCreateOrderPayAmountForOrderType(limitAmount, feeRate float64, currency, orderType string, usdToCnyRate float64) (string, float64, error) { + paymentAmount := limitAmount + if orderType == payment.OrderTypeSubscription { + paymentAmount = calculateSubscriptionGatewayBaseAmount(limitAmount, usdToCnyRate, currency) + } + return calculateCreateOrderPayAmount(paymentAmount, feeRate, currency) +} + +// calculateSubscriptionGatewayBaseAmount 计算订阅订单的网关扣款基数。 +// 换算是显式 opt-in：仅当管理员配置了订阅汇率（rate > 0，1 USD = rate CNY） +// 且网关币种为 CNY 时，按 price × rate 换算；未配置时保持 price 直付的存量行为。 +func calculateSubscriptionGatewayBaseAmount(amount, usdToCnyRate float64, currency string) float64 { + rate := normalizeSubscriptionUSDToCNYRate(usdToCnyRate) + if rate <= 0 || currency != payment.DefaultPaymentCurrency { + return amount + } + return decimal.NewFromFloat(amount). + Mul(decimal.NewFromFloat(rate)). + Round(int32(payment.CurrencyMaxFractionDigits(currency))). + InexactFloat64() +} + func validateCreateOrderAmountCurrency(amount float64, currency string) error { amountStr := strconv.FormatFloat(amount, 'f', -1, 64) if _, err := payment.AmountToMinorUnit(amountStr, currency); err != nil { diff --git a/backend/internal/service/payment_order_result_test.go b/backend/internal/service/payment_order_result_test.go index b7545ee45d..ac439ee6f2 100644 --- a/backend/internal/service/payment_order_result_test.go +++ b/backend/internal/service/payment_order_result_test.go @@ -91,6 +91,41 @@ func TestBuildCreateOrderResponseCopiesJSAPIPayload(t *testing.T) { } } +func TestSanitizeCreatePaymentResponseDetailsRemovesNULBytes(t *testing.T) { + t.Parallel() + + resp := &payment.CreatePaymentResponse{ + TradeNo: "trade\x00-no", + PayURL: "https://pay.example.com/\x00checkout", + QRCode: "wxp://payment-token\x00", + ClientSecret: "secret\x00unchanged", + } + + sanitizeCreatePaymentResponseDetails(resp) + + if strings.ContainsRune(resp.TradeNo, 0) { + t.Fatalf("trade_no still contains NUL: %q", resp.TradeNo) + } + if strings.ContainsRune(resp.PayURL, 0) { + t.Fatalf("pay_url still contains NUL: %q", resp.PayURL) + } + if strings.ContainsRune(resp.QRCode, 0) { + t.Fatalf("qr_code still contains NUL: %q", resp.QRCode) + } + if resp.TradeNo != "trade-no" { + t.Fatalf("trade_no = %q, want trade-no", resp.TradeNo) + } + if resp.PayURL != "https://pay.example.com/checkout" { + t.Fatalf("pay_url = %q, want sanitized URL", resp.PayURL) + } + if resp.QRCode != "wxp://payment-token" { + t.Fatalf("qr_code = %q, want sanitized QR code", resp.QRCode) + } + if resp.ClientSecret != "secret\x00unchanged" { + t.Fatalf("client_secret = %q, should not be touched by payment detail sanitization", resp.ClientSecret) + } +} + func TestValidateSelectedCreateOrderAmountCurrencyRejectsFractionalZeroDecimal(t *testing.T) { t.Parallel() @@ -126,27 +161,66 @@ func TestCalculateCreateOrderPayAmountUsesCurrencyPrecision(t *testing.T) { } } -func TestCalculateCreateOrderPayAmountForSubscriptionKeepsDirectPrice(t *testing.T) { +func TestCalculateCreateOrderPayAmountForSubscriptionConvertsCNYPriceWhenRateConfigured(t *testing.T) { t.Parallel() - amountStr, amount, err := calculateCreateOrderPayAmount(5, 0, "CNY") + amountStr, amount, err := calculateCreateOrderPayAmountForOrderType(9.99, 0, "CNY", payment.OrderTypeSubscription, 7.15) if err != nil { t.Fatalf("unexpected error: %v", err) } - if amountStr != "5.00" || amount != 5 { - t.Fatalf("subscription CNY pay amount = (%q, %v), want (5.00, 5)", amountStr, amount) + if amountStr != "71.43" || amount != 71.43 { + t.Fatalf("subscription CNY pay amount = (%q, %v), want (71.43, 71.43)", amountStr, amount) } } -func TestCalculateCreateOrderPayAmountForSubscriptionAppliesFeeToDirectPrice(t *testing.T) { +func TestCalculateCreateOrderPayAmountForSubscriptionAppliesFeeAfterCNYConversion(t *testing.T) { t.Parallel() - amountStr, amount, err := calculateCreateOrderPayAmount(5, 2.5, "CNY") + amountStr, amount, err := calculateCreateOrderPayAmountForOrderType(9.99, 2.5, "CNY", payment.OrderTypeSubscription, 7.15) if err != nil { t.Fatalf("unexpected error: %v", err) } - if amountStr != "5.13" || amount != 5.13 { - t.Fatalf("subscription CNY pay amount with fee = (%q, %v), want (5.13, 5.13)", amountStr, amount) + if amountStr != "73.22" || amount != 73.22 { + t.Fatalf("subscription CNY pay amount with fee = (%q, %v), want (73.22, 73.22)", amountStr, amount) + } +} + +func TestCalculateCreateOrderPayAmountForSubscriptionKeepsNonCNYPrice(t *testing.T) { + t.Parallel() + + amountStr, amount, err := calculateCreateOrderPayAmountForOrderType(9.99, 0, "USD", payment.OrderTypeSubscription, 7.15) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if amountStr != "9.99" || amount != 9.99 { + t.Fatalf("subscription USD pay amount = (%q, %v), want (9.99, 9.99)", amountStr, amount) + } +} + +// 换算是 opt-in：未配置汇率（rate=0）时，CNY 订阅保持 price 直付的存量行为。 +// 该测试锁住存量部署升级后行为不变的兼容承诺。 +func TestCalculateCreateOrderPayAmountForSubscriptionKeepsDirectPriceWhenRateDisabled(t *testing.T) { + t.Parallel() + + amountStr, amount, err := calculateCreateOrderPayAmountForOrderType(9.99, 0, "CNY", payment.OrderTypeSubscription, 0) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if amountStr != "9.99" || amount != 9.99 { + t.Fatalf("subscription CNY pay amount without rate = (%q, %v), want (9.99, 9.99)", amountStr, amount) + } +} + +// 汇率只作用于订阅订单，余额充值订单不受影响。 +func TestCalculateCreateOrderPayAmountForBalanceIgnoresSubscriptionRate(t *testing.T) { + t.Parallel() + + amountStr, amount, err := calculateCreateOrderPayAmountForOrderType(50, 0, "CNY", payment.OrderTypeBalance, 7.15) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if amountStr != "50.00" || amount != 50 { + t.Fatalf("balance CNY pay amount = (%q, %v), want (50.00, 50)", amountStr, amount) } } diff --git a/backend/internal/service/payment_resume_service_test.go b/backend/internal/service/payment_resume_service_test.go index 7e0adc2de8..17b637fa23 100644 --- a/backend/internal/service/payment_resume_service_test.go +++ b/backend/internal/service/payment_resume_service_test.go @@ -26,9 +26,10 @@ func TestNormalizeVisibleMethods(t *testing.T) { " wxpay_direct ", "wxpay", "stripe", + "ldc", }) - want := []string{"alipay", "wxpay", "stripe"} + want := []string{"alipay", "wxpay", "stripe", "ldc"} if len(got) != len(want) { t.Fatalf("NormalizeVisibleMethods len = %d, want %d (%v)", len(got), len(want), got) } @@ -39,6 +40,21 @@ func TestNormalizeVisibleMethods(t *testing.T) { } } +func TestEnabledVisibleMethodsForEasyPayIncludesCustomSupportedTypes(t *testing.T) { + t.Parallel() + + got := enabledVisibleMethodsForProvider(payment.TypeEasyPay, "alipay,ldc,usdt_trc20") + want := []string{"alipay", "ldc", "usdt_trc20"} + if len(got) != len(want) { + t.Fatalf("enabledVisibleMethodsForProvider len = %d, want %d (%v)", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("enabledVisibleMethodsForProvider[%d] = %q, want %q (full=%v)", i, got[i], want[i], got) + } + } +} + func TestNormalizePaymentSource(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/payment_visible_method_instances.go b/backend/internal/service/payment_visible_method_instances.go index 899bd7a020..97b3b1ef66 100644 --- a/backend/internal/service/payment_visible_method_instances.go +++ b/backend/internal/service/payment_visible_method_instances.go @@ -16,8 +16,7 @@ func enabledVisibleMethodsForProvider(providerKey, supportedTypes string) []stri methodSet := make(map[string]struct{}, 2) addMethod := func(method string) { method = NormalizeVisibleMethod(method) - switch method { - case payment.TypeAlipay, payment.TypeWxpay: + if method != "" { methodSet[method] = struct{}{} } } @@ -55,6 +54,14 @@ func enabledVisibleMethodsForProvider(providerKey, supportedTypes string) []stri for _, method := range []string{payment.TypeAlipay, payment.TypeWxpay} { if _, ok := methodSet[method]; ok { methods = append(methods, method) + delete(methodSet, method) + } + } + for _, supportedType := range splitTypes(supportedTypes) { + method := NormalizeVisibleMethod(supportedType) + if _, ok := methodSet[method]; ok { + methods = append(methods, method) + delete(methodSet, method) } } return methods @@ -215,7 +222,7 @@ func (s *PaymentConfigService) resolveEnabledVisibleMethodInstance( } method = NormalizeVisibleMethod(method) - if method != payment.TypeAlipay && method != payment.TypeWxpay { + if method == "" { return nil, nil } diff --git a/backend/internal/service/pricing_service.go b/backend/internal/service/pricing_service.go index bd0c30df45..f3fb36d73a 100644 --- a/backend/internal/service/pricing_service.go +++ b/backend/internal/service/pricing_service.go @@ -319,6 +319,7 @@ func (s *PricingService) downloadPricingData() error { if err != nil { return fmt.Errorf("parse pricing data: %w", err) } + data = s.mergeFallbackPricingData(data) // 保存到本地文件 pricingFile := s.getPricingFilePath() @@ -373,7 +374,7 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel } // 只保留有有效价格的条目 - if entry.InputCostPerToken == nil && entry.OutputCostPerToken == nil { + if entry.InputCostPerToken == nil && entry.OutputCostPerToken == nil && entry.OutputCostPerImage == nil && entry.OutputCostPerImageToken == nil { continue } @@ -441,6 +442,7 @@ func (s *PricingService) loadPricingData(filePath string) error { if err != nil { return fmt.Errorf("parse pricing data: %w", err) } + pricingData = s.mergeFallbackPricingData(pricingData) // 计算哈希 hash := sha256.Sum256(data) @@ -462,6 +464,37 @@ func (s *PricingService) loadPricingData(filePath string) error { return nil } +func (s *PricingService) mergeFallbackPricingData(data map[string]*LiteLLMModelPricing) map[string]*LiteLLMModelPricing { + if data == nil { + data = make(map[string]*LiteLLMModelPricing) + } + if s == nil || s.cfg == nil || strings.TrimSpace(s.cfg.Pricing.FallbackFile) == "" { + return data + } + fallbackBody, err := os.ReadFile(s.cfg.Pricing.FallbackFile) + if err != nil { + logger.LegacyPrintf("service.pricing", "[Pricing] Fallback merge skipped: %v", err) + return data + } + fallbackData, err := s.parsePricingData(fallbackBody) + if err != nil { + logger.LegacyPrintf("service.pricing", "[Pricing] Fallback merge parse skipped: %v", err) + return data + } + merged := 0 + for modelName, pricing := range fallbackData { + if _, ok := data[modelName]; ok { + continue + } + data[modelName] = pricing + merged++ + } + if merged > 0 { + logger.LegacyPrintf("service.pricing", "[Pricing] Merged %d fallback-only models", merged) + } + return data +} + // useFallbackPricing 使用回退价格文件 func (s *PricingService) useFallbackPricing() error { fallbackFile := s.cfg.Pricing.FallbackFile @@ -798,6 +831,13 @@ func (s *PricingService) matchOpenAIModel(model string) *LiteLLMModelPricing { } } + // GPT-5.6（sol / terra / luna）回退到 GPT-5.4 定价 + if strings.HasPrefix(model, "gpt-5.6") { + logger.With(zap.String("component", "service.pricing")). + Info(fmt.Sprintf("[Pricing] OpenAI fallback matched %s -> %s", model, "gpt-5.4(static)")) + return openAIGPT54FallbackPricing + } + // GPT-5.5 回退到 GPT-5.4 定价 if strings.HasPrefix(model, "gpt-5.5") { logger.With(zap.String("component", "service.pricing")). diff --git a/backend/internal/service/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index f4252f9540..11c1b58da9 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -6,6 +6,7 @@ import ( "path/filepath" "testing" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/stretchr/testify/require" ) @@ -37,6 +38,57 @@ func TestParsePricingData_ParsesPriorityAndServiceTierFields(t *testing.T) { require.True(t, pricing.SupportsServiceTier) } +func TestParsePricingData_KeepsImageOnlyPricing(t *testing.T) { + svc := &PricingService{} + body := []byte(`{ + "image-only-model": { + "output_cost_per_image": 0.034, + "litellm_provider": "vertex_ai-language-models", + "mode": "image_generation" + } + }`) + + data, err := svc.parsePricingData(body) + require.NoError(t, err) + pricing := data["image-only-model"] + require.NotNil(t, pricing) + require.InDelta(t, 0.034, pricing.OutputCostPerImage, 1e-12) + require.Equal(t, "image_generation", pricing.Mode) +} + +func TestPricingService_MergesFallbackOnlyModels(t *testing.T) { + dir := t.TempDir() + fallbackFile := filepath.Join(dir, "fallback.json") + require.NoError(t, os.WriteFile(fallbackFile, []byte(`{ + "remote-model": { + "input_cost_per_token": 0.000001, + "litellm_provider": "test", + "mode": "chat" + }, + "gemini-3.1-flash-lite-image": { + "output_cost_per_image": 0.034, + "litellm_provider": "vertex_ai-language-models", + "mode": "image_generation" + } + }`), 0644)) + + svc := &PricingService{cfg: &config.Config{}} + svc.cfg.Pricing.FallbackFile = fallbackFile + remoteData, err := svc.parsePricingData([]byte(`{ + "remote-model": { + "input_cost_per_token": 0.000002, + "litellm_provider": "test", + "mode": "chat" + } + }`)) + require.NoError(t, err) + + merged := svc.mergeFallbackPricingData(remoteData) + require.InDelta(t, 0.000002, merged["remote-model"].InputCostPerToken, 1e-12) + require.NotNil(t, merged["gemini-3.1-flash-lite-image"]) + require.InDelta(t, 0.034, merged["gemini-3.1-flash-lite-image"].OutputCostPerImage, 1e-12) +} + func TestGetModelPricing_Gpt53CodexSparkUsesGpt51CodexPricing(t *testing.T) { sparkPricing := &LiteLLMModelPricing{InputCostPerToken: 1} gpt53Pricing := &LiteLLMModelPricing{InputCostPerToken: 9} diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index 2b9c85a2fa..50d38e7d13 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -116,6 +116,14 @@ func (s *RateLimitService) SetAccountRuntimeBlocker(blocker AccountRuntimeBlocke s.runtimeBlocker = blocker } +func (s *RateLimitService) IsOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx context.Context) bool { + if s == nil || s.settingService == nil { + return false + } + gateway := &OpenAIGatewayService{rateLimitService: s} + return gateway.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx) +} + func (s *RateLimitService) notifyAccountSchedulingBlocked(account *Account, until time.Time, reason string) { if s == nil || s.runtimeBlocker == nil || account == nil { return @@ -292,6 +300,20 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc if upstreamMsg != "" { msg = "OAuth 401: " + upstreamMsg } + if authAccount.Platform == PlatformAntigravity { + extraUpdates := antigravityForceTokenRefreshExtra("401_invalid") + if err := s.accountRepo.UpdateExtra(ctx, authAccount.ID, extraUpdates); err != nil { + slog.Warn("antigravity_401_force_refresh_mark_failed", "account_id", authAccount.ID, "error", err) + } else { + if authAccount.Extra == nil { + authAccount.Extra = make(map[string]any, len(extraUpdates)) + } + for k, v := range extraUpdates { + authAccount.Extra[k] = v + } + slog.Info("antigravity_401_force_refresh_marked", "account_id", authAccount.ID) + } + } cooldownMinutes := s.cfg.RateLimit.OAuth401CooldownMinutes if cooldownMinutes <= 0 { cooldownMinutes = 10 diff --git a/backend/internal/service/ratelimit_service_401_test.go b/backend/internal/service/ratelimit_service_401_test.go index 09afb9314b..48e6a41def 100644 --- a/backend/internal/service/ratelimit_service_401_test.go +++ b/backend/internal/service/ratelimit_service_401_test.go @@ -18,7 +18,9 @@ type rateLimitAccountRepoStub struct { setErrorCalls int tempCalls int updateCredentialsCalls int + updateExtraCalls int lastCredentials map[string]any + lastExtraUpdates map[string]any lastErrorMsg string lastTempReason string lastErrorID int64 @@ -45,6 +47,12 @@ func (r *rateLimitAccountRepoStub) UpdateCredentials(ctx context.Context, id int return nil } +func (r *rateLimitAccountRepoStub) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error { + r.updateExtraCalls++ + r.lastExtraUpdates = shallowCopyMap(updates) + return nil +} + type tokenCacheInvalidatorRecorder struct { accounts []*Account err error @@ -133,6 +141,10 @@ func TestRateLimitService_HandleUpstreamError_OAuth401SetsTempUnschedulable(t *t require.Equal(t, 1, repo.tempCalls) require.Equal(t, int64(100), repo.lastTempID) require.Contains(t, repo.lastTempReason, "invalid or expired credentials") + require.Equal(t, 1, repo.updateExtraCalls) + require.Equal(t, true, repo.lastExtraUpdates[antigravityForceTokenRefreshExtraKey]) + require.Equal(t, "401_invalid", repo.lastExtraUpdates[antigravityForceTokenRefreshReasonExtraKey]) + require.Equal(t, true, account.Extra[antigravityForceTokenRefreshExtraKey]) require.Len(t, invalidator.accounts, 1) require.Equal(t, int64(100), invalidator.accounts[0].ID) }) @@ -245,6 +257,7 @@ func TestRateLimitService_HandleUpstreamError_OAuth401DoesNotOverwriteCredential require.True(t, shouldDisable) require.Equal(t, 0, repo.updateCredentialsCalls, "401 handler must not write credentials back from the request-start snapshot") + require.Equal(t, 0, repo.updateExtraCalls, "OpenAI 401 must not set Antigravity force-refresh marker") require.Equal(t, 1, repo.tempCalls, "401 handler should still set temp-unschedulable cooldown") require.Nil(t, repo.lastCredentials, "no credentials should have been persisted") } diff --git a/backend/internal/service/ratelimit_session_window_test.go b/backend/internal/service/ratelimit_session_window_test.go index cb19227e54..279a31ccdc 100644 --- a/backend/internal/service/ratelimit_session_window_test.go +++ b/backend/internal/service/ratelimit_session_window_test.go @@ -87,6 +87,9 @@ func (m *sessionWindowMockRepo) List(context.Context, pagination.PaginationParam func (m *sessionWindowMockRepo) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string, string, int64, string) ([]Account, *pagination.PaginationResult, error) { panic("unexpected") } +func (m *sessionWindowMockRepo) ListAllWithFilters(context.Context, string, string, string, string, int64, string) ([]Account, error) { + panic("unexpected") +} func (m *sessionWindowMockRepo) ListByGroup(context.Context, int64) ([]Account, error) { panic("unexpected") } diff --git a/backend/internal/service/setting_service.go b/backend/internal/service/setting_service.go index 1a5e677a12..3aeb611418 100644 --- a/backend/internal/service/setting_service.go +++ b/backend/internal/service/setting_service.go @@ -1921,6 +1921,9 @@ func (s *SettingService) buildSystemSettingsUpdates(ctx context.Context, setting if err != nil { return nil, err } + if err := s.normalizeOpenAIAdvancedSchedulerOverrides(settings); err != nil { + return nil, err + } settings.PaymentVisibleMethodAlipaySource = alipaySource settings.PaymentVisibleMethodWxpaySource = wxpaySource settings.WeChatConnectAppID = strings.TrimSpace(settings.WeChatConnectAppID) @@ -2229,6 +2232,18 @@ func (s *SettingService) buildSystemSettingsUpdates(ctx context.Context, setting updates[SettingPaymentVisibleMethodAlipayEnabled] = strconv.FormatBool(settings.PaymentVisibleMethodAlipayEnabled) updates[SettingPaymentVisibleMethodWxpayEnabled] = strconv.FormatBool(settings.PaymentVisibleMethodWxpayEnabled) updates[openAIAdvancedSchedulerSettingKey] = strconv.FormatBool(settings.OpenAIAdvancedSchedulerEnabled) + updates[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled] = strconv.FormatBool(settings.OpenAIAdvancedSchedulerStickyWeightedEnabled) + updates[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled] = strconv.FormatBool(settings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled) + updates[SettingKeyOpenAIAdvancedSchedulerLBTopK] = settings.OpenAIAdvancedSchedulerLBTopK + updates[SettingKeyOpenAIAdvancedSchedulerWeightPriority] = settings.OpenAIAdvancedSchedulerWeightPriority + updates[SettingKeyOpenAIAdvancedSchedulerWeightLoad] = settings.OpenAIAdvancedSchedulerWeightLoad + updates[SettingKeyOpenAIAdvancedSchedulerWeightQueue] = settings.OpenAIAdvancedSchedulerWeightQueue + updates[SettingKeyOpenAIAdvancedSchedulerWeightErrorRate] = settings.OpenAIAdvancedSchedulerWeightErrorRate + updates[SettingKeyOpenAIAdvancedSchedulerWeightTTFT] = settings.OpenAIAdvancedSchedulerWeightTTFT + updates[SettingKeyOpenAIAdvancedSchedulerWeightReset] = settings.OpenAIAdvancedSchedulerWeightReset + updates[SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom] = settings.OpenAIAdvancedSchedulerWeightQuotaHeadroom + updates[SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse] = settings.OpenAIAdvancedSchedulerWeightPreviousResponse + updates[SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky] = settings.OpenAIAdvancedSchedulerWeightSessionSticky // 余额、订阅到期与账号限额通知 updates[SettingKeyBalanceLowNotifyEnabled] = strconv.FormatBool(settings.BalanceLowNotifyEnabled) @@ -2376,7 +2391,21 @@ func (s *SettingService) refreshCachedSettings(settings *SystemSettings) { }) openAIAdvancedSchedulerSettingSF.Forget(openAIAdvancedSchedulerSettingKey) openAIAdvancedSchedulerSettingCache.Store(&cachedOpenAIAdvancedSchedulerSetting{ - enabled: settings.OpenAIAdvancedSchedulerEnabled, + enabled: settings.OpenAIAdvancedSchedulerEnabled, + stickyWeightedEnabled: settings.OpenAIAdvancedSchedulerStickyWeightedEnabled, + subscriptionPriorityEnabled: settings.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled, + lbTopKOverride: parsePositiveIntOverride(settings.OpenAIAdvancedSchedulerLBTopK), + weightOverrides: parseOpenAIAdvancedSchedulerWeightOverrides(map[string]string{ + SettingKeyOpenAIAdvancedSchedulerWeightPriority: settings.OpenAIAdvancedSchedulerWeightPriority, + SettingKeyOpenAIAdvancedSchedulerWeightLoad: settings.OpenAIAdvancedSchedulerWeightLoad, + SettingKeyOpenAIAdvancedSchedulerWeightQueue: settings.OpenAIAdvancedSchedulerWeightQueue, + SettingKeyOpenAIAdvancedSchedulerWeightErrorRate: settings.OpenAIAdvancedSchedulerWeightErrorRate, + SettingKeyOpenAIAdvancedSchedulerWeightTTFT: settings.OpenAIAdvancedSchedulerWeightTTFT, + SettingKeyOpenAIAdvancedSchedulerWeightReset: settings.OpenAIAdvancedSchedulerWeightReset, + SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom: settings.OpenAIAdvancedSchedulerWeightQuotaHeadroom, + SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: settings.OpenAIAdvancedSchedulerWeightPreviousResponse, + SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky: settings.OpenAIAdvancedSchedulerWeightSessionSticky, + }), expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(), }) // Invalidate the quota auto-pause cache and let the next read trigger a fresh load. @@ -3196,17 +3225,29 @@ func (s *SettingService) InitializeDefaultSettings(ctx context.Context) error { SettingKeyCodexCLIOnlyEngineFingerprintSignals: openai.DefaultEngineFingerprintSignalsJSON(), // 分组隔离（默认不允许未分组 Key 调度） - SettingKeyAllowUngroupedKeyScheduling: "false", - SettingKeyEnableAnthropicCacheTTL1hInjection: "false", - SettingKeyRewriteMessageCacheControl: strconv.FormatBool(s.defaultRewriteMessageCacheControl()), - SettingKeyEnableClientDatelineNormalization: "true", - SettingKeyAntigravityUserAgentVersion: "", - SettingKeyOpenAICodexUserAgent: "", - SettingPaymentVisibleMethodAlipaySource: "", - SettingPaymentVisibleMethodWxpaySource: "", - SettingPaymentVisibleMethodAlipayEnabled: "false", - SettingPaymentVisibleMethodWxpayEnabled: "false", - openAIAdvancedSchedulerSettingKey: "false", + SettingKeyAllowUngroupedKeyScheduling: "false", + SettingKeyEnableAnthropicCacheTTL1hInjection: "false", + SettingKeyRewriteMessageCacheControl: strconv.FormatBool(s.defaultRewriteMessageCacheControl()), + SettingKeyEnableClientDatelineNormalization: "true", + SettingKeyAntigravityUserAgentVersion: "", + SettingKeyOpenAICodexUserAgent: "", + SettingPaymentVisibleMethodAlipaySource: "", + SettingPaymentVisibleMethodWxpaySource: "", + SettingPaymentVisibleMethodAlipayEnabled: "false", + SettingPaymentVisibleMethodWxpayEnabled: "false", + openAIAdvancedSchedulerSettingKey: "false", + SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled: "false", + SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled: "false", + SettingKeyOpenAIAdvancedSchedulerLBTopK: "", + SettingKeyOpenAIAdvancedSchedulerWeightPriority: "", + SettingKeyOpenAIAdvancedSchedulerWeightLoad: "", + SettingKeyOpenAIAdvancedSchedulerWeightQueue: "", + SettingKeyOpenAIAdvancedSchedulerWeightErrorRate: "", + SettingKeyOpenAIAdvancedSchedulerWeightTTFT: "", + SettingKeyOpenAIAdvancedSchedulerWeightReset: "", + SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom: "", + SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: "", + SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky: "", SettingKeyAllowUserViewErrorRequests: "false", } @@ -3769,6 +3810,29 @@ func (s *SettingService) parseSettings(settings map[string]string) *SystemSettin result.PaymentVisibleMethodAlipayEnabled = settings[SettingPaymentVisibleMethodAlipayEnabled] == "true" result.PaymentVisibleMethodWxpayEnabled = settings[SettingPaymentVisibleMethodWxpayEnabled] == "true" result.OpenAIAdvancedSchedulerEnabled = settings[openAIAdvancedSchedulerSettingKey] == "true" + result.OpenAIAdvancedSchedulerStickyWeightedEnabled = settings[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled] == "true" + result.OpenAIAdvancedSchedulerSubscriptionPriorityEnabled = settings[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled] == "true" + result.OpenAIAdvancedSchedulerLBTopK = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerLBTopK]) + result.OpenAIAdvancedSchedulerWeightPriority = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightPriority]) + result.OpenAIAdvancedSchedulerWeightLoad = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightLoad]) + result.OpenAIAdvancedSchedulerWeightQueue = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightQueue]) + result.OpenAIAdvancedSchedulerWeightErrorRate = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightErrorRate]) + result.OpenAIAdvancedSchedulerWeightTTFT = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightTTFT]) + result.OpenAIAdvancedSchedulerWeightReset = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightReset]) + result.OpenAIAdvancedSchedulerWeightQuotaHeadroom = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom]) + result.OpenAIAdvancedSchedulerWeightPreviousResponse = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse]) + result.OpenAIAdvancedSchedulerWeightSessionSticky = strings.TrimSpace(settings[SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky]) + result.OpenAIAdvancedSchedulerEffectiveLBTopK = s.openAIAdvancedSchedulerEffectiveLBTopK() + effectiveWeights := s.openAIAdvancedSchedulerEffectiveWeights() + result.OpenAIAdvancedSchedulerEffectiveWeightPriority = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.Priority) + result.OpenAIAdvancedSchedulerEffectiveWeightLoad = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.Load) + result.OpenAIAdvancedSchedulerEffectiveWeightQueue = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.Queue) + result.OpenAIAdvancedSchedulerEffectiveWeightErrorRate = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.ErrorRate) + result.OpenAIAdvancedSchedulerEffectiveWeightTTFT = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.TTFT) + result.OpenAIAdvancedSchedulerEffectiveWeightReset = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.Reset) + result.OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.QuotaHeadroom) + result.OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.PreviousResponse) + result.OpenAIAdvancedSchedulerEffectiveWeightSessionSticky = formatOpenAIAdvancedSchedulerFloat(effectiveWeights.SessionSticky) // 余额、订阅到期与账号限额通知 result.BalanceLowNotifyEnabled = settings[SettingKeyBalanceLowNotifyEnabled] == "true" @@ -3841,6 +3905,119 @@ func normalizeVisibleMethodSettingSource(method, source string, enabled bool) (s return normalized, nil } +func (s *SettingService) openAIAdvancedSchedulerEffectiveLBTopK() string { + if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.LBTopK > 0 { + return strconv.Itoa(s.cfg.Gateway.OpenAIWS.LBTopK) + } + return "7" +} + +func (s *SettingService) openAIAdvancedSchedulerEffectiveWeights() config.GatewayOpenAIWSSchedulerScoreWeights { + defaults := config.GatewayOpenAIWSSchedulerScoreWeights{ + Priority: 1.0, + Load: 1.0, + Queue: 0.7, + ErrorRate: 0.8, + TTFT: 0.5, + Reset: 0.0, + QuotaHeadroom: 0.0, + PreviousResponse: 5.0, + SessionSticky: 3.0, + } + if s == nil || s.cfg == nil { + return defaults + } + + weights := s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights + baseSum := weights.Priority + weights.Load + weights.Queue + weights.ErrorRate + weights.TTFT + weights.QuotaHeadroom + if baseSum <= 0 { + return defaults + } + return weights +} + +func formatOpenAIAdvancedSchedulerFloat(value float64) string { + return strconv.FormatFloat(value, 'f', -1, 64) +} + +func (s *SettingService) normalizeOpenAIAdvancedSchedulerOverrides(settings *SystemSettings) error { + lbTopK, err := normalizeOptionalPositiveIntString(settings.OpenAIAdvancedSchedulerLBTopK) + if err != nil { + return infraerrors.BadRequest("INVALID_OPENAI_ADVANCED_SCHEDULER_LB_TOP_K", "openai advanced scheduler TopK must be a positive integer or empty") + } + settings.OpenAIAdvancedSchedulerLBTopK = lbTopK + + weights := []*string{ + &settings.OpenAIAdvancedSchedulerWeightPriority, + &settings.OpenAIAdvancedSchedulerWeightLoad, + &settings.OpenAIAdvancedSchedulerWeightQueue, + &settings.OpenAIAdvancedSchedulerWeightErrorRate, + &settings.OpenAIAdvancedSchedulerWeightTTFT, + &settings.OpenAIAdvancedSchedulerWeightReset, + &settings.OpenAIAdvancedSchedulerWeightQuotaHeadroom, + &settings.OpenAIAdvancedSchedulerWeightPreviousResponse, + &settings.OpenAIAdvancedSchedulerWeightSessionSticky, + } + for _, target := range weights { + normalized, err := normalizeOptionalNonNegativeFloatString(*target) + if err != nil { + return infraerrors.BadRequest("INVALID_OPENAI_ADVANCED_SCHEDULER_WEIGHT", "openai advanced scheduler weights must be non-negative numbers or empty") + } + *target = normalized + } + + // 与 config.Validate 的 "scheduler_score_weights must not all be zero" 保持一致： + // 覆盖值（空则回退到生效的配置值）叠加后的基础权重和不允许为 0， + // 否则调度会静默退化为 TopK 内均匀随机。 + effective := s.openAIAdvancedSchedulerEffectiveWeights() + baseSum := resolveOpenAIAdvancedSchedulerWeight(settings.OpenAIAdvancedSchedulerWeightPriority, effective.Priority) + + resolveOpenAIAdvancedSchedulerWeight(settings.OpenAIAdvancedSchedulerWeightLoad, effective.Load) + + resolveOpenAIAdvancedSchedulerWeight(settings.OpenAIAdvancedSchedulerWeightQueue, effective.Queue) + + resolveOpenAIAdvancedSchedulerWeight(settings.OpenAIAdvancedSchedulerWeightErrorRate, effective.ErrorRate) + + resolveOpenAIAdvancedSchedulerWeight(settings.OpenAIAdvancedSchedulerWeightTTFT, effective.TTFT) + + resolveOpenAIAdvancedSchedulerWeight(settings.OpenAIAdvancedSchedulerWeightQuotaHeadroom, effective.QuotaHeadroom) + if baseSum <= 0 { + return infraerrors.BadRequest("INVALID_OPENAI_ADVANCED_SCHEDULER_WEIGHT", "openai advanced scheduler base weights must not all be zero") + } + return nil +} + +// resolveOpenAIAdvancedSchedulerWeight 返回覆盖值（已归一化的非空字符串），空则回退默认值。 +func resolveOpenAIAdvancedSchedulerWeight(normalized string, fallback float64) float64 { + if normalized == "" { + return fallback + } + value, err := strconv.ParseFloat(normalized, 64) + if err != nil { + return fallback + } + return value +} + +func normalizeOptionalPositiveIntString(raw string) (string, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return "", nil + } + value, err := strconv.Atoi(raw) + if err != nil || value <= 0 { + return "", fmt.Errorf("invalid positive integer") + } + return strconv.Itoa(value), nil +} + +func normalizeOptionalNonNegativeFloatString(raw string) (string, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return "", nil + } + value, err := strconv.ParseFloat(raw, 64) + if err != nil || value < 0 || math.IsNaN(value) || math.IsInf(value, 0) { + return "", fmt.Errorf("invalid non-negative float") + } + return strconv.FormatFloat(value, 'f', -1, 64), nil +} + func parseDefaultSubscriptions(raw string) []DefaultSubscriptionSetting { raw = strings.TrimSpace(raw) if raw == "" { diff --git a/backend/internal/service/setting_service_update_test.go b/backend/internal/service/setting_service_update_test.go index 379bf9bc04..a47e1379f7 100644 --- a/backend/internal/service/setting_service_update_test.go +++ b/backend/internal/service/setting_service_update_test.go @@ -49,6 +49,42 @@ func (s *settingUpdateRepoStub) Delete(ctx context.Context, key string) error { panic("unexpected Delete call") } +type settingGetAllRepoStub struct { + values map[string]string +} + +func (s *settingGetAllRepoStub) Get(ctx context.Context, key string) (*Setting, error) { + panic("unexpected Get call") +} + +func (s *settingGetAllRepoStub) GetValue(ctx context.Context, key string) (string, error) { + panic("unexpected GetValue call") +} + +func (s *settingGetAllRepoStub) Set(ctx context.Context, key, value string) error { + panic("unexpected Set call") +} + +func (s *settingGetAllRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) { + panic("unexpected GetMultiple call") +} + +func (s *settingGetAllRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error { + panic("unexpected SetMultiple call") +} + +func (s *settingGetAllRepoStub) GetAll(ctx context.Context) (map[string]string, error) { + out := make(map[string]string, len(s.values)) + for key, value := range s.values { + out[key] = value + } + return out, nil +} + +func (s *settingGetAllRepoStub) Delete(ctx context.Context, key string) error { + panic("unexpected Delete call") +} + type settingAntigravityUARepoStub struct { values map[string]string } @@ -261,15 +297,30 @@ func TestSettingService_UpdateSettings_TablePreferences(t *testing.T) { } func TestSettingService_UpdateSettings_PaymentVisibleMethodsAndAdvancedScheduler(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + defer resetOpenAIAdvancedSchedulerSettingCacheForTest() + repo := &settingUpdateRepoStub{} svc := NewSettingService(repo, &config.Config{}) err := svc.UpdateSettings(context.Background(), &SystemSettings{ - PaymentVisibleMethodAlipaySource: "alipay", - PaymentVisibleMethodWxpaySource: "easypay", - PaymentVisibleMethodAlipayEnabled: true, - PaymentVisibleMethodWxpayEnabled: false, - OpenAIAdvancedSchedulerEnabled: true, + PaymentVisibleMethodAlipaySource: "alipay", + PaymentVisibleMethodWxpaySource: "easypay", + PaymentVisibleMethodAlipayEnabled: true, + PaymentVisibleMethodWxpayEnabled: false, + OpenAIAdvancedSchedulerEnabled: true, + OpenAIAdvancedSchedulerStickyWeightedEnabled: true, + OpenAIAdvancedSchedulerSubscriptionPriorityEnabled: true, + OpenAIAdvancedSchedulerLBTopK: " 3 ", + OpenAIAdvancedSchedulerWeightPriority: "2.50", + OpenAIAdvancedSchedulerWeightLoad: "0", + OpenAIAdvancedSchedulerWeightQueue: "0.75", + OpenAIAdvancedSchedulerWeightErrorRate: "1.25", + OpenAIAdvancedSchedulerWeightTTFT: "0.5", + OpenAIAdvancedSchedulerWeightReset: "", + OpenAIAdvancedSchedulerWeightQuotaHeadroom: "0.2", + OpenAIAdvancedSchedulerWeightPreviousResponse: "8", + OpenAIAdvancedSchedulerWeightSessionSticky: "4", }) require.NoError(t, err) require.Equal(t, VisibleMethodSourceOfficialAlipay, repo.updates[SettingPaymentVisibleMethodAlipaySource]) @@ -277,6 +328,49 @@ func TestSettingService_UpdateSettings_PaymentVisibleMethodsAndAdvancedScheduler require.Equal(t, "true", repo.updates[SettingPaymentVisibleMethodAlipayEnabled]) require.Equal(t, "false", repo.updates[SettingPaymentVisibleMethodWxpayEnabled]) require.Equal(t, "true", repo.updates[openAIAdvancedSchedulerSettingKey]) + require.Equal(t, "true", repo.updates[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]) + require.Equal(t, "true", repo.updates[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]) + require.Equal(t, "3", repo.updates[SettingKeyOpenAIAdvancedSchedulerLBTopK]) + require.Equal(t, "2.5", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightPriority]) + require.Equal(t, "0", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightLoad]) + require.Equal(t, "0.75", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightQueue]) + require.Equal(t, "1.25", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightErrorRate]) + require.Equal(t, "0.5", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightTTFT]) + require.Equal(t, "", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightReset]) + require.Equal(t, "0.2", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom]) + require.Equal(t, "8", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse]) + require.Equal(t, "4", repo.updates[SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky]) +} + +func TestSettingService_GetAllSettings_OpenAIAdvancedSchedulerEffectiveValuesUseConfig(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 13 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{ + Priority: 2, + Load: 3, + Queue: 4, + ErrorRate: 5, + TTFT: 6, + Reset: 7, + QuotaHeadroom: 8, + PreviousResponse: 9, + SessionSticky: 10, + } + svc := NewSettingService(&settingGetAllRepoStub{values: map[string]string{ + SettingKeyOpenAIAdvancedSchedulerLBTopK: "3", + SettingKeyOpenAIAdvancedSchedulerWeightPriority: "99", + SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky: "88", + }}, cfg) + + settings, err := svc.GetAllSettings(context.Background()) + require.NoError(t, err) + require.Equal(t, "3", settings.OpenAIAdvancedSchedulerLBTopK) + require.Equal(t, "99", settings.OpenAIAdvancedSchedulerWeightPriority) + require.Equal(t, "88", settings.OpenAIAdvancedSchedulerWeightSessionSticky) + require.Equal(t, "13", settings.OpenAIAdvancedSchedulerEffectiveLBTopK) + require.Equal(t, "2", settings.OpenAIAdvancedSchedulerEffectiveWeightPriority) + require.Equal(t, "3", settings.OpenAIAdvancedSchedulerEffectiveWeightLoad) + require.Equal(t, "10", settings.OpenAIAdvancedSchedulerEffectiveWeightSessionSticky) } func TestSettingService_UpdateSettings_AntigravityUserAgentVersion(t *testing.T) { diff --git a/backend/internal/service/settings_view.go b/backend/internal/service/settings_view.go index ac225e1b14..465760a445 100644 --- a/backend/internal/service/settings_view.go +++ b/backend/internal/service/settings_view.go @@ -218,7 +218,29 @@ type SystemSettings struct { PaymentVisibleMethodWxpayEnabled bool // OpenAI 账号调度 - OpenAIAdvancedSchedulerEnabled bool + OpenAIAdvancedSchedulerEnabled bool + OpenAIAdvancedSchedulerStickyWeightedEnabled bool + OpenAIAdvancedSchedulerSubscriptionPriorityEnabled bool + OpenAIAdvancedSchedulerLBTopK string + OpenAIAdvancedSchedulerWeightPriority string + OpenAIAdvancedSchedulerWeightLoad string + OpenAIAdvancedSchedulerWeightQueue string + OpenAIAdvancedSchedulerWeightErrorRate string + OpenAIAdvancedSchedulerWeightTTFT string + OpenAIAdvancedSchedulerWeightReset string + OpenAIAdvancedSchedulerWeightQuotaHeadroom string + OpenAIAdvancedSchedulerWeightPreviousResponse string + OpenAIAdvancedSchedulerWeightSessionSticky string + OpenAIAdvancedSchedulerEffectiveLBTopK string + OpenAIAdvancedSchedulerEffectiveWeightPriority string + OpenAIAdvancedSchedulerEffectiveWeightLoad string + OpenAIAdvancedSchedulerEffectiveWeightQueue string + OpenAIAdvancedSchedulerEffectiveWeightErrorRate string + OpenAIAdvancedSchedulerEffectiveWeightTTFT string + OpenAIAdvancedSchedulerEffectiveWeightReset string + OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom string + OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse string + OpenAIAdvancedSchedulerEffectiveWeightSessionSticky string // 余额不足提醒 BalanceLowNotifyEnabled bool diff --git a/backend/internal/service/token_refresh_service.go b/backend/internal/service/token_refresh_service.go index 9ca369a52a..2d4a026692 100644 --- a/backend/internal/service/token_refresh_service.go +++ b/backend/internal/service/token_refresh_service.go @@ -312,6 +312,7 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc if isNonRetryableRefreshError(err) { errorMsg := "Token refresh failed (non-retryable): " + logredact.RedactText(err.Error()) s.notifyAccountSchedulingBlocked(account, time.Time{}, "token_refresh_non_retryable") + s.clearAntigravityForceTokenRefresh(ctx, account, "non_retryable") if setErr := s.accountRepo.SetError(ctx, account.ID, errorMsg); setErr != nil { slog.Error("token_refresh.set_error_status_failed", "account_id", account.ID, @@ -369,6 +370,8 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc // postRefreshActions 刷新成功后的后续动作（清除错误状态、缓存失效、调度器同步等） func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *Account) { + s.clearAntigravityForceTokenRefresh(ctx, account, "success") + // Antigravity 账户：如果之前是因为缺少 project_id 而标记为 error，现在成功获取到了，清除错误状态 if account.Platform == PlatformAntigravity && account.Status == StatusError && @@ -432,6 +435,30 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A s.ensureAntigravityPrivacy(ctx, account) } +func (s *TokenRefreshService) clearAntigravityForceTokenRefresh(ctx context.Context, account *Account, outcome string) { + if s == nil || account == nil || !accountNeedsAntigravityForceTokenRefresh(account) { + return + } + updates := clearAntigravityForceTokenRefreshExtra() + if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil { + slog.Warn("token_refresh.clear_antigravity_force_refresh_failed", + "account_id", account.ID, + "outcome", outcome, + "error", err, + ) + return + } + if account.Extra != nil { + for k, v := range updates { + account.Extra[k] = v + } + } + slog.Info("token_refresh.cleared_antigravity_force_refresh", + "account_id", account.ID, + "outcome", outcome, + ) +} + // errRefreshSkipped 表示刷新被跳过（锁竞争或已被其他路径刷新），不计入 failed 或 refreshed var errRefreshSkipped = fmt.Errorf("refresh skipped") diff --git a/backend/internal/service/token_refresh_service_test.go b/backend/internal/service/token_refresh_service_test.go index 6f227352f9..bd6a521576 100644 --- a/backend/internal/service/token_refresh_service_test.go +++ b/backend/internal/service/token_refresh_service_test.go @@ -20,8 +20,10 @@ type tokenRefreshAccountRepo struct { setErrorCalls int clearTempCalls int setTempUnschedCalls int + updateExtraCalls int lastErrorMessage string lastTempUnschedReason string + lastExtraUpdates map[string]any lastAccount *Account updateErr error } @@ -68,6 +70,22 @@ func (r *tokenRefreshAccountRepo) SetTempUnschedulable(ctx context.Context, id i return nil } +func (r *tokenRefreshAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error { + r.updateExtraCalls++ + r.lastExtraUpdates = shallowCopyMap(updates) + if r.accountsByID != nil { + if acc, ok := r.accountsByID[id]; ok && acc != nil { + if acc.Extra == nil { + acc.Extra = make(map[string]any, len(updates)) + } + for k, v := range updates { + acc.Extra[k] = v + } + } + } + return nil +} + type tokenCacheInvalidatorStub struct { calls int err error @@ -233,6 +251,121 @@ func TestTokenRefreshService_RefreshWithRetry_Antigravity(t *testing.T) { require.Equal(t, 1, invalidator.calls) // Antigravity 也应触发缓存失效 } +func TestAntigravityTokenRefresher_NeedsRefresh_ForceRefreshMarker(t *testing.T) { + refresher := NewAntigravityTokenRefresher(nil) + account := &Account{ + ID: 3675, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + Extra: map[string]any{ + antigravityForceTokenRefreshExtraKey: true, + }, + } + + require.True(t, refresher.NeedsRefresh(account, 0), "server-invalidated token must refresh even before expires_at") +} + +func TestAntigravityTokenRefresher_NeedsRefresh_NormalExpiryRulesUnchanged(t *testing.T) { + refresher := NewAntigravityTokenRefresher(nil) + + t.Run("normal_unexpired_without_marker_does_not_refresh", func(t *testing.T) { + account := &Account{ + ID: 3707, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + } + + require.False(t, refresher.NeedsRefresh(account, 0)) + }) + + t.Run("normal_expiring_refreshes", func(t *testing.T) { + account := &Account{ + ID: 3708, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "expires_at": time.Now().Add(5 * time.Minute).Format(time.RFC3339), + }, + } + + require.True(t, refresher.NeedsRefresh(account, 0)) + }) +} + +func TestTokenRefreshService_RefreshWithRetry_AntigravityClearsForceRefreshOnSuccess(t *testing.T) { + repo := &tokenRefreshAccountRepo{} + cfg := &config.Config{ + TokenRefresh: config.TokenRefreshConfig{ + MaxRetries: 1, + RetryBackoffSeconds: 0, + }, + } + service := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, nil, cfg, nil) + until := time.Now().Add(10 * time.Minute) + account := &Account{ + ID: 3709, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + TempUnschedulableUntil: &until, + Extra: map[string]any{ + antigravityForceTokenRefreshExtraKey: true, + antigravityForceTokenRefreshReasonExtraKey: "401_invalid", + "privacy_mode": AntigravityPrivacySet, + }, + } + refresher := &tokenRefresherStub{ + credentials: map[string]any{ + "access_token": "new-ag-token", + }, + } + + err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour) + require.NoError(t, err) + require.Equal(t, 1, repo.updateCredentialsCalls) + require.Equal(t, 1, repo.updateExtraCalls) + require.Equal(t, false, repo.lastExtraUpdates[antigravityForceTokenRefreshExtraKey]) + require.Equal(t, "", repo.lastExtraUpdates[antigravityForceTokenRefreshReasonExtraKey]) + require.Equal(t, false, account.Extra[antigravityForceTokenRefreshExtraKey]) + require.Equal(t, 1, repo.clearTempCalls, "successful refresh should restore schedulability") +} + +func TestTokenRefreshService_RefreshWithRetry_AntigravityForceRefreshInvalidGrantSetsError(t *testing.T) { + repo := &tokenRefreshAccountRepo{} + cfg := &config.Config{ + TokenRefresh: config.TokenRefreshConfig{ + MaxRetries: 3, + RetryBackoffSeconds: 0, + }, + } + service := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, nil, cfg, nil) + account := &Account{ + ID: 3710, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Extra: map[string]any{ + antigravityForceTokenRefreshExtraKey: true, + antigravityForceTokenRefreshReasonExtraKey: "401_invalid", + }, + } + refresher := &tokenRefresherStub{ + err: errors.New("invalid_grant: token revoked"), + } + + err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour) + require.Error(t, err) + require.Equal(t, 1, repo.setErrorCalls) + require.Equal(t, 0, repo.setTempUnschedCalls) + require.Equal(t, 1, repo.updateExtraCalls) + require.Equal(t, false, repo.lastExtraUpdates[antigravityForceTokenRefreshExtraKey]) + require.Contains(t, repo.lastErrorMessage, "non-retryable") +} + // TestTokenRefreshService_RefreshWithRetry_NonOAuthAccount 测试非 OAuth 账号不触发缓存失效 func TestTokenRefreshService_RefreshWithRetry_NonOAuthAccount(t *testing.T) { repo := &tokenRefreshAccountRepo{} diff --git a/backend/internal/service/upstream_models.go b/backend/internal/service/upstream_models.go index ee3e6bfc04..4f6a305b25 100644 --- a/backend/internal/service/upstream_models.go +++ b/backend/internal/service/upstream_models.go @@ -208,6 +208,8 @@ func (s *AccountTestService) buildAnthropicUpstreamModelsRequest(ctx context.Con } else { setAnthropicAPIKeyAuthHeader(req.Header, account, apiKeyAuthToken) } + // 账号级请求头覆写：模型列表探测与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) return req, nil } @@ -277,6 +279,8 @@ func (s *AccountTestService) buildOpenAIUpstreamModelsRequest(ctx context.Contex } req.Header.Set("Accept", "application/json") req.Header.Set("Authorization", "Bearer "+apiKey) + // 账号级请求头覆写：模型列表探测与真实转发保持一致的最终头 + account.ApplyHeaderOverrides(req.Header) return req, nil } @@ -387,14 +391,7 @@ func buildV1ModelsURL(base string) string { } func buildOpenAIModelsURL(base string) string { - normalized := strings.TrimRight(strings.TrimSpace(base), "/") - if strings.HasSuffix(normalized, "/v1/models") { - return normalized - } - if strings.HasSuffix(normalized, "/v1") { - return normalized + "/models" - } - return normalized + "/v1/models" + return buildOpenAIEndpointURL(base, "/v1/models") } func buildGeminiModelsURL(base string) string { diff --git a/backend/internal/service/upstream_models_test.go b/backend/internal/service/upstream_models_test.go index 1fe9415d34..3904194ffa 100644 --- a/backend/internal/service/upstream_models_test.go +++ b/backend/internal/service/upstream_models_test.go @@ -29,6 +29,61 @@ func TestBuildV1ModelsURL(t *testing.T) { require.Equal(t, "https://gateway.example.com/antigravity/v1/models", buildV1ModelsURL("https://gateway.example.com/antigravity/")) } +func TestBuildOpenAIModelsURL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + base string + want string + }{ + { + name: "zhipu v4 coding base url", + base: "https://open.bigmodel.cn/api/coding/paas/v4", + want: "https://open.bigmodel.cn/api/coding/paas/v4/models", + }, + { + name: "openai v1 base url", + base: "https://api.openai.com/v1", + want: "https://api.openai.com/v1/models", + }, + { + name: "models url unchanged", + base: "https://api.openai.com/v1/models", + want: "https://api.openai.com/v1/models", + }, + { + name: "host fallback uses v1", + base: "https://api.openai.com", + want: "https://api.openai.com/v1/models", + }, + { + name: "trailing slash on v4", + base: "https://open.bigmodel.cn/api/coding/paas/v4/", + want: "https://open.bigmodel.cn/api/coding/paas/v4/models", + }, + { + name: "v2 base url", + base: "https://gateway.example.com/openai/v2", + want: "https://gateway.example.com/openai/v2/models", + }, + { + name: "v3 base url", + base: "https://gateway.example.com/openai/v3", + want: "https://gateway.example.com/openai/v3/models", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + require.Equal(t, tt.want, buildOpenAIModelsURL(tt.base)) + }) + } +} + func TestBuildGeminiModelsURL(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/usage_billing.go b/backend/internal/service/usage_billing.go index accc7cb2cb..8d52c92d26 100644 --- a/backend/internal/service/usage_billing.go +++ b/backend/internal/service/usage_billing.go @@ -119,6 +119,57 @@ type UsageBillingApplyResult struct { QuotaState *AccountQuotaState // post-increment quota state (nil = no quota increment) } +// BatchImageBalanceHoldCommand describes an idempotent balance hold operation. +type BatchImageBalanceHoldCommand struct { + RequestID string + APIKeyID int64 + RequestFingerprint string + RequestPayloadHash string + UserID int64 + BatchID string + HoldAmount float64 + ActualAmount float64 +} + +func (c *BatchImageBalanceHoldCommand) Normalize() { + if c == nil { + return + } + c.RequestID = strings.TrimSpace(c.RequestID) + c.BatchID = strings.TrimSpace(c.BatchID) + if strings.TrimSpace(c.RequestFingerprint) == "" { + c.RequestFingerprint = buildBatchImageBalanceHoldFingerprint(c) + } +} + +func buildBatchImageBalanceHoldFingerprint(c *BatchImageBalanceHoldCommand) string { + if c == nil { + return "" + } + raw := fmt.Sprintf( + "%d|%d|%s|%0.10f|%0.10f", + c.UserID, + c.APIKeyID, + strings.TrimSpace(c.BatchID), + c.HoldAmount, + c.ActualAmount, + ) + if payloadHash := strings.TrimSpace(c.RequestPayloadHash); payloadHash != "" { + raw += "|" + payloadHash + } + sum := sha256.Sum256([]byte(raw)) + return hex.EncodeToString(sum[:]) +} + +type BatchImageBalanceHoldResult struct { + Applied bool + NewBalance *float64 + FrozenBalance *float64 +} + type UsageBillingRepository interface { Apply(ctx context.Context, cmd *UsageBillingCommand) (*UsageBillingApplyResult, error) + ReserveBatchImageBalance(ctx context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) + CaptureBatchImageBalance(ctx context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) + ReleaseBatchImageBalance(ctx context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) } diff --git a/backend/internal/service/user.go b/backend/internal/service/user.go index edb944ee05..22a21a4634 100644 --- a/backend/internal/service/user.go +++ b/backend/internal/service/user.go @@ -19,6 +19,7 @@ type User struct { PasswordHash string Role string Balance float64 + FrozenBalance float64 Concurrency int Status string AllowedGroups []int64 diff --git a/backend/internal/service/user_msg_queue_service.go b/backend/internal/service/user_msg_queue_service.go index f3f105ac93..72027d387a 100644 --- a/backend/internal/service/user_msg_queue_service.go +++ b/backend/internal/service/user_msg_queue_service.go @@ -25,10 +25,8 @@ type UserMsgQueueCache interface { GetLastCompletedMs(ctx context.Context, accountID int64) (int64, error) // GetCurrentTimeMs 获取 Redis 服务器当前时间（毫秒），与 ReleaseLock 记录的时间源一致 GetCurrentTimeMs(ctx context.Context) (int64, error) - // ForceReleaseLock 强制释放锁（孤儿锁清理） - ForceReleaseLock(ctx context.Context, accountID int64) error - // ScanLockKeys 扫描 PTTL == -1 的孤儿锁 key，返回 accountID 列表 - ScanLockKeys(ctx context.Context, maxCount int) ([]int64, error) + // ReconcileExpiredLockCandidates 处理锁索引中的到期候选，按真实 PTTL 清理或刷新索引 + ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (cleaned int, err error) } // QueueLockResult 锁获取结果 @@ -246,8 +244,8 @@ func (s *UserMessageQueueService) CalculateRPMAwareDelay(ctx context.Context, ac return applyJitter(baseDelay, 0.15) } -// StartCleanupWorker 启动孤儿锁清理 worker -// 定期 SCAN umq:*:lock 并清理 PTTL == -1 的异常锁（PTTL 检查在 cache.ScanLockKeys 内完成） +// StartCleanupWorker 启动孤儿锁清理 worker。 +// worker 只处理锁索引中的到期候选，真正删除前由 cache 层再次校验锁 PTTL。 func (s *UserMessageQueueService) StartCleanupWorker(interval time.Duration) { if s == nil || s.cache == nil || interval <= 0 { return @@ -257,23 +255,13 @@ func (s *UserMessageQueueService) StartCleanupWorker(interval time.Duration) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - accountIDs, err := s.cache.ScanLockKeys(ctx, 1000) + // 每轮限制处理数量，避免清理任务在大量过期候选时长时间占用 Redis。 + cleaned, err := s.cache.ReconcileExpiredLockCandidates(ctx, 1000) if err != nil { - logger.LegacyPrintf("service.umq", "Cleanup scan failed: %v", err) + logger.LegacyPrintf("service.umq", "Cleanup reconcile failed: %v", err) return } - cleaned := 0 - for _, accountID := range accountIDs { - cleanCtx, cleanCancel := context.WithTimeout(context.Background(), 2*time.Second) - if err := s.cache.ForceReleaseLock(cleanCtx, accountID); err != nil { - logger.LegacyPrintf("service.umq", "Cleanup force release failed for account %d: %v", accountID, err) - } else { - cleaned++ - } - cleanCancel() - } - if cleaned > 0 { logger.LegacyPrintf("service.umq", "Cleanup completed: released %d orphaned locks", cleaned) } diff --git a/backend/internal/service/user_msg_queue_service_test.go b/backend/internal/service/user_msg_queue_service_test.go new file mode 100644 index 0000000000..f554269b88 --- /dev/null +++ b/backend/internal/service/user_msg_queue_service_test.go @@ -0,0 +1,54 @@ +//go:build unit + +package service + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type cleanupWorkerUserMsgQueueCache struct { + reconcileCalls atomic.Int64 + maxCount atomic.Int64 +} + +var _ UserMsgQueueCache = (*cleanupWorkerUserMsgQueueCache)(nil) + +func (c *cleanupWorkerUserMsgQueueCache) AcquireLock(context.Context, int64, string, int) (bool, error) { + return true, nil +} + +func (c *cleanupWorkerUserMsgQueueCache) ReleaseLock(context.Context, int64, string) (bool, error) { + return true, nil +} + +func (c *cleanupWorkerUserMsgQueueCache) GetLastCompletedMs(context.Context, int64) (int64, error) { + return 0, nil +} + +func (c *cleanupWorkerUserMsgQueueCache) GetCurrentTimeMs(context.Context) (int64, error) { + return time.Now().UnixMilli(), nil +} + +func (c *cleanupWorkerUserMsgQueueCache) ReconcileExpiredLockCandidates(_ context.Context, maxCount int) (int, error) { + c.reconcileCalls.Add(1) + c.maxCount.Store(int64(maxCount)) + return 1, nil +} + +func TestStartCleanupWorker_ReconcilesExpiredLockCandidates(t *testing.T) { + cache := &cleanupWorkerUserMsgQueueCache{} + svc := NewUserMessageQueueService(cache, nil, nil) + defer svc.Stop() + + svc.StartCleanupWorker(time.Millisecond) + + require.Eventually(t, func() bool { + return cache.reconcileCalls.Load() > 0 + }, time.Second, 10*time.Millisecond) + require.EqualValues(t, 1000, cache.maxCount.Load()) +} diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7278a4e0c5..7258ff05a3 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -45,6 +45,16 @@ func ProvideOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiToke return NewOAuthRefreshAPI(accountRepo, tokenCache) } +func ProvideBatchImageModelPricingResolver(resolver *ModelPricingResolver) *BatchImageModelPricingResolver { + return &BatchImageModelPricingResolver{Resolver: resolver} +} + +func ProvideBatchImageCleanupService(repo BatchImageRepository, accountRepo AccountRepository, cfg *config.Config) *BatchImageCleanupService { + svc := NewBatchImageCleanupService(repo, accountRepo, cfg) + svc.Start() + return svc +} + // ProvideOpenAIOAuthService creates OpenAIOAuthService with privacy/account enrichment support. func ProvideOpenAIOAuthService( proxyRepo ProxyRepository, @@ -537,9 +547,11 @@ func ProvideAPIKeyService( cache APIKeyCache, cfg *config.Config, billingCacheService *BillingCacheService, + concurrencyService *ConcurrencyService, ) *APIKeyService { svc := NewAPIKeyService(apiKeyRepo, userRepo, groupRepo, userSubRepo, userGroupRateRepo, cache, cfg) svc.SetRateLimitCacheInvalidator(billingCacheService) + svc.SetConcurrencyService(concurrencyService) return svc } @@ -564,6 +576,11 @@ var ProviderSet = wire.NewSet( NewAdminService, NewGatewayService, NewOpenAIGatewayService, + ProvideBatchImageModelPricingResolver, + NewBatchImagePublicService, + NewBatchImageDownloadService, + ProvideBatchImageCleanupService, + ProvideBatchImageWorkerRuntime, wire.Bind(new(AccountRuntimeBlocker), new(*OpenAIGatewayService)), NewOAuthService, ProvideOpenAIOAuthService, 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/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index e88ed2da22..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, @@ -4886,6 +4960,150 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-5.6-sol": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "gpt-5.6-terra": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "gpt-5.6-luna": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.5": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, diff --git a/deploy/.env.example b/deploy/.env.example index d8892dbecf..5925f0abb4 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -359,11 +359,11 @@ DASHBOARD_AGGREGATION_RETENTION_DAILY_DAYS=730 # 启用 URL 白名单验证（false 则跳过白名单检查，仅做基本格式校验） SECURITY_URL_ALLOWLIST_ENABLED=false -# 关闭白名单时，是否允许 http:// URL（默认 false，只允许 https://） -# ⚠️ 警告：允许 HTTP 存在安全风险（明文传输），仅建议在开发/测试环境或可信内网中使用 -# Allow insecure HTTP URLs when allowlist is disabled (default: false, requires https) +# 关闭白名单时，是否允许 http:// URL（默认 true，设为 false 则只允许 https://） +# ⚠️ 警告：允许 HTTP 存在安全风险（明文传输），生产环境建议设为 false +# Allow insecure HTTP URLs when allowlist is disabled (default: true; set to false to require https) # ⚠️ WARNING: Allowing HTTP has security risks (plaintext transmission) -# Only recommended for dev/test environments or trusted networks +# Recommended to set false in production SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=true # 是否允许本地/私有 IP 地址用于上游/定价/CRS（仅在可信网络中使用） diff --git a/deploy/Dockerfile b/deploy/Dockerfile index c9222da50f..d83b0e25b8 100644 --- a/deploy/Dockerfile +++ b/deploy/Dockerfile @@ -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/config.example.yaml b/deploy/config.example.yaml index 5a3afb0315..eff4bfb598 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -108,8 +108,8 @@ security: # Allow localhost/private IPs for upstream/pricing/CRS (use only in trusted networks) # 允许本地/私有 IP 地址用于上游/定价/CRS（仅在可信网络中使用） allow_private_hosts: true - # Allow http:// URLs when allowlist is disabled (default: false, require https) - # 白名单禁用时是否允许 http:// URL（默认: false，要求 https） + # Allow http:// URLs when allowlist is disabled (default: true; set to false to require https) + # 白名单禁用时是否允许 http:// URL（默认: true，设为 false 则仅允许 https） allow_insecure_http: true response_headers: # Enable configurable response header filtering (default: true) 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/deploy/docker-compose.local.yml b/deploy/docker-compose.local.yml index 21f46e3760..042752e857 100644 --- a/deploy/docker-compose.local.yml +++ b/deploy/docker-compose.local.yml @@ -139,10 +139,10 @@ services: # ======================================================================= # Enable URL allowlist validation (false to skip allowlist checks) - SECURITY_URL_ALLOWLIST_ENABLED=${SECURITY_URL_ALLOWLIST_ENABLED:-false} - # Allow insecure HTTP URLs when allowlist is disabled (default: false, requires https) - - SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=${SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP:-false} - # Allow private IP addresses for upstream/pricing/CRS (for internal deployments) - - SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS=${SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS:-false} + # Allow insecure HTTP URLs when allowlist is disabled (default: true; set to false to require https) + - SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=${SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP:-true} + # Allow private IP addresses for upstream/pricing/CRS (default: true; set to false to block private hosts) + - SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS=${SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS:-true} # Upstream hosts whitelist (comma-separated, only used when enabled=true) - SECURITY_URL_ALLOWLIST_UPSTREAM_HOSTS=${SECURITY_URL_ALLOWLIST_UPSTREAM_HOSTS:-} diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 3da3bcbfb0..22713c59aa 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -135,10 +135,10 @@ services: # ======================================================================= # Enable URL allowlist validation (false to skip allowlist checks) - SECURITY_URL_ALLOWLIST_ENABLED=${SECURITY_URL_ALLOWLIST_ENABLED:-false} - # Allow insecure HTTP URLs when allowlist is disabled (default: false, requires https) - - SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=${SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP:-false} - # Allow private IP addresses for upstream/pricing/CRS (for internal deployments) - - SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS=${SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS:-false} + # Allow insecure HTTP URLs when allowlist is disabled (default: true; set to false to require https) + - SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP=${SECURITY_URL_ALLOWLIST_ALLOW_INSECURE_HTTP:-true} + # Allow private IP addresses for upstream/pricing/CRS (default: true; set to false to block private hosts) + - SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS=${SECURITY_URL_ALLOWLIST_ALLOW_PRIVATE_HOSTS:-true} # Upstream hosts whitelist (comma-separated, only used when enabled=true) - SECURITY_URL_ALLOWLIST_UPSTREAM_HOSTS=${SECURITY_URL_ALLOWLIST_UPSTREAM_HOSTS:-} 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/admin/payment.ts b/frontend/src/api/admin/payment.ts index 49efcc355d..1d4305948e 100644 --- a/frontend/src/api/admin/payment.ts +++ b/frontend/src/api/admin/payment.ts @@ -24,6 +24,8 @@ export interface AdminPaymentConfig { enabled_payment_types: string[] balance_disabled: boolean balance_recharge_multiplier: number + subscription_usd_to_cny_rate: number + recharge_fee_rate: number load_balance_strategy: string product_name_prefix: string product_name_suffix: string @@ -42,6 +44,8 @@ export interface UpdatePaymentConfigRequest { enabled_payment_types?: string[] balance_disabled?: boolean balance_recharge_multiplier?: number + subscription_usd_to_cny_rate?: number + recharge_fee_rate?: number load_balance_strategy?: string product_name_prefix?: string product_name_suffix?: string diff --git a/frontend/src/api/admin/settings.ts b/frontend/src/api/admin/settings.ts index 44fbe29187..f5da990930 100644 --- a/frontend/src/api/admin/settings.ts +++ b/frontend/src/api/admin/settings.ts @@ -589,6 +589,7 @@ export interface SystemSettings { payment_enabled_types: string[]; payment_balance_disabled: boolean; payment_balance_recharge_multiplier: number; + payment_subscription_usd_to_cny_rate: number; payment_recharge_fee_rate: number; payment_load_balance_strategy: string; payment_product_name_prefix: string; @@ -606,6 +607,28 @@ export interface SystemSettings { payment_visible_method_alipay_enabled?: boolean; payment_visible_method_wxpay_enabled?: boolean; openai_advanced_scheduler_enabled?: boolean; + openai_advanced_scheduler_sticky_weighted_enabled?: boolean; + openai_advanced_scheduler_subscription_priority_enabled?: boolean; + openai_advanced_scheduler_lb_top_k?: string; + openai_advanced_scheduler_weight_priority?: string; + openai_advanced_scheduler_weight_load?: string; + openai_advanced_scheduler_weight_queue?: string; + openai_advanced_scheduler_weight_error_rate?: string; + openai_advanced_scheduler_weight_ttft?: string; + openai_advanced_scheduler_weight_reset?: string; + openai_advanced_scheduler_weight_quota_headroom?: string; + openai_advanced_scheduler_weight_previous_response?: string; + openai_advanced_scheduler_weight_session_sticky?: string; + openai_advanced_scheduler_effective_lb_top_k?: string; + openai_advanced_scheduler_effective_weight_priority?: string; + openai_advanced_scheduler_effective_weight_load?: string; + openai_advanced_scheduler_effective_weight_queue?: string; + openai_advanced_scheduler_effective_weight_error_rate?: string; + openai_advanced_scheduler_effective_weight_ttft?: string; + openai_advanced_scheduler_effective_weight_reset?: string; + openai_advanced_scheduler_effective_weight_quota_headroom?: string; + openai_advanced_scheduler_effective_weight_previous_response?: string; + openai_advanced_scheduler_effective_weight_session_sticky?: string; // 余额、订阅到期与账号限额通知 balance_low_notify_enabled: boolean; @@ -838,6 +861,7 @@ export interface UpdateSettingsRequest { payment_enabled_types?: string[]; payment_balance_disabled?: boolean; payment_balance_recharge_multiplier?: number; + payment_subscription_usd_to_cny_rate?: number; payment_recharge_fee_rate?: number; payment_load_balance_strategy?: string; payment_product_name_prefix?: string; @@ -855,6 +879,18 @@ export interface UpdateSettingsRequest { payment_visible_method_alipay_enabled?: boolean; payment_visible_method_wxpay_enabled?: boolean; openai_advanced_scheduler_enabled?: boolean; + openai_advanced_scheduler_sticky_weighted_enabled?: boolean; + openai_advanced_scheduler_subscription_priority_enabled?: boolean; + openai_advanced_scheduler_lb_top_k?: string; + openai_advanced_scheduler_weight_priority?: string; + openai_advanced_scheduler_weight_load?: string; + openai_advanced_scheduler_weight_queue?: string; + openai_advanced_scheduler_weight_error_rate?: string; + openai_advanced_scheduler_weight_ttft?: string; + openai_advanced_scheduler_weight_reset?: string; + openai_advanced_scheduler_weight_quota_headroom?: string; + openai_advanced_scheduler_weight_previous_response?: string; + openai_advanced_scheduler_weight_session_sticky?: string; // 余额、订阅到期与账号限额通知 balance_low_notify_enabled?: boolean; balance_low_notify_threshold?: 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/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/assets/icons/payment.svg b/frontend/src/assets/icons/payment.svg new file mode 100644 index 0000000000..c78bea4cf7 --- /dev/null +++ b/frontend/src/assets/icons/payment.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/frontend/src/components/account/BulkEditAccountModal.vue b/frontend/src/components/account/BulkEditAccountModal.vue index d24488b47a..91656b196d 100644 --- a/frontend/src/components/account/BulkEditAccountModal.vue +++ b/frontend/src/components/account/BulkEditAccountModal.vue @@ -486,6 +486,129 @@ + +
+
+
+ +

+ {{ t('admin.accounts.headerOverride.hint') }} +

+
+ +
+
+ + +
+
+

+ + {{ t('admin.accounts.headerOverride.info') }} +

+
+ +

+ {{ t('admin.accounts.headerOverride.bulkReplaceHint') }} +

+ +
+
+ + + +
+
+ + + +
+ +
+ +

+ {{ t('admin.accounts.headerOverride.emptyValueHint') }} +

+
+

+ {{ t('admin.accounts.headerOverride.bulkDisableHint') }} +

+
+
+
@@ -1149,6 +1272,16 @@ import { buildModelMappingObject as buildModelMappingPayload, getPresetMappingsByPlatform } from '@/composables/useModelWhitelist' +import { + buildHeaderOverridesObject, + getHeaderOverrideTemplate, + isHeaderOverridePlatform, + validateHeaderOverrideRows, + HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY, + HEADER_OVERRIDES_CREDENTIAL_KEY, + type HeaderOverrideRow +} from '@/components/account/credentialsBuilder' +import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' import { OPENAI_WS_MODE_CTX_POOL, OPENAI_WS_MODE_OFF, @@ -1217,6 +1350,16 @@ const allOpenAIAPIKey = computed(() => { ) }) +// 是否全部为 anthropic/openai 平台的 apikey 账号（请求头覆写仅在此条件下显示） +const allHeaderOverrideCapable = computed(() => { + return ( + targetSelectedPlatforms.value.length > 0 && + targetSelectedPlatforms.value.every(p => isHeaderOverridePlatform(p)) && + targetSelectedTypes.value.length > 0 && + targetSelectedTypes.value.every(t => t === 'apikey') + ) +}) + // 是否全部为 Anthropic OAuth/SetupToken（RPM 配置仅在此条件下显示） const allAnthropicOAuthOrSetupToken = computed(() => { return ( @@ -1253,6 +1396,7 @@ const enableBaseUrl = ref(false) const enableModelRestriction = ref(false) const enableCustomErrorCodes = ref(false) const enableInterceptWarmup = ref(false) +const enableHeaderOverride = ref(false) const enableProxy = ref(false) const enableConcurrency = ref(false) const enableLoadFactor = ref(false) @@ -1281,6 +1425,39 @@ const modelMappings = ref([]) const selectedErrorCodes = ref([]) const customErrorCodeInput = ref(null) const interceptWarmupRequests = ref(false) +const headerOverrideEnabled = ref(false) +const headerOverrideRows = ref([]) +const getHeaderOverrideRowKey = createStableObjectKeyResolver('bulk-header-override-row') + +const addHeaderOverrideRow = () => { + headerOverrideRows.value.push({ name: '', value: '' }) +} + +const removeHeaderOverrideRow = (index: number) => { + headerOverrideRows.value.splice(index, 1) +} + +// 模板仅在所选账号平台唯一时可用：混合 anthropic+openai 选择无法确定用哪套模板， +// 误填会把另一平台的专有头写进所有所选账号 +const headerOverrideTemplatePlatform = computed(() => { + return targetSelectedPlatforms.value.length === 1 ? targetSelectedPlatforms.value[0] : null +}) + +// 模板按钮：填入所选平台的标准客户端请求头名称（值留空），跳过已存在的同名行 +const fillHeaderOverrideTemplate = () => { + const platform = headerOverrideTemplatePlatform.value + if (!platform) return + const existing = new Set( + headerOverrideRows.value.map((row) => row.name.trim().toLowerCase()).filter(Boolean) + ) + const rows = headerOverrideRows.value.filter((row) => row.name.trim() || row.value.trim()) + for (const row of getHeaderOverrideTemplate(platform)) { + if (!existing.has(row.name)) { + rows.push(row) + } + } + headerOverrideRows.value = rows +} const proxyId = ref(null) const concurrency = ref(1) const loadFactor = ref(null) @@ -1523,6 +1700,15 @@ const buildUpdatePayload = (): Record | null => { credentialsChanged = true } + if (enableHeaderOverride.value) { + // 后端使用 JSONB || merge 语义：关闭时显式写入 false + 空对象以清除旧配置 + credentials[HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY] = headerOverrideEnabled.value + credentials[HEADER_OVERRIDES_CREDENTIAL_KEY] = headerOverrideEnabled.value + ? buildHeaderOverridesObject(headerOverrideRows.value) + : {} + credentialsChanged = true + } + if (enableOpenAIWSMode.value) { const extra = ensureExtra() extra.openai_oauth_responses_websockets_v2_mode = openaiOAuthResponsesWebSocketV2Mode.value @@ -1651,6 +1837,7 @@ const handleSubmit = async () => { enableModelRestriction.value || enableCustomErrorCodes.value || enableInterceptWarmup.value || + enableHeaderOverride.value || enableProxy.value || enableConcurrency.value || enableLoadFactor.value || @@ -1672,6 +1859,20 @@ const handleSubmit = async () => { return } + if (enableHeaderOverride.value && headerOverrideEnabled.value) { + // 批量保存对 header_overrides 是整键替换：开启但没有任何有效行会把所选账号的 + // 既有覆写配置静默清空，必须显式拦截（清空请走关闭开关的路径，有专门提示） + if (!headerOverrideRows.value.some((row) => row.name.trim())) { + appStore.showError(t('admin.accounts.headerOverride.bulkEmptyRows')) + return + } + const headerError = validateHeaderOverrideRows(headerOverrideRows.value) + if (headerError) { + appStore.showError(t(`admin.accounts.headerOverride.${headerError}`)) + return + } + } + const built = buildUpdatePayload() if (!built) { appStore.showError(t('admin.accounts.bulkEdit.noFieldsSelected')) @@ -1753,6 +1954,7 @@ watch( enableModelRestriction.value = false enableCustomErrorCodes.value = false enableInterceptWarmup.value = false + enableHeaderOverride.value = false enableProxy.value = false enableConcurrency.value = false enableLoadFactor.value = false @@ -1778,6 +1980,8 @@ watch( selectedErrorCodes.value = [] customErrorCodeInput.value = null interceptWarmupRequests.value = false + headerOverrideEnabled.value = false + headerOverrideRows.value = [] proxyId.value = null concurrency.value = 1 loadFactor.value = null diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 3514153c82..67750e4e34 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -1468,6 +1468,110 @@
+ +
+
+
+ +

+ {{ t('admin.accounts.headerOverride.hint') }} +

+
+ +
+ +
+
+

+ + {{ t('admin.accounts.headerOverride.info') }} +

+
+ +
+
+ + + +
+
+ + + +
+ +
+ +

+ {{ t('admin.accounts.headerOverride.emptyValueHint') }} +

+
+
+ @@ -3328,7 +3432,12 @@ import ModelWhitelistSelector from '@/components/account/ModelWhitelistSelector. import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' import { applyAntigravityProjectID, - applyInterceptWarmup + applyHeaderOverride, + applyInterceptWarmup, + getHeaderOverrideTemplate, + isHeaderOverridePlatform, + validateHeaderOverrideRows, + type HeaderOverrideRow } from '@/components/account/credentialsBuilder' import { formatDateTimeLocalInput, parseDateTimeLocalInput } from '@/utils/format' import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' @@ -3512,6 +3621,30 @@ function parsePoolModeRetryStatusCodes(input: string): number[] { const customErrorCodesEnabled = ref(false) const selectedErrorCodes = ref([]) const customErrorCodeInput = ref(null) +const headerOverrideEnabled = ref(false) +const headerOverrideRows = ref([]) + +const addHeaderOverrideRow = () => { + headerOverrideRows.value.push({ name: '', value: '' }) +} + +const removeHeaderOverrideRow = (index: number) => { + headerOverrideRows.value.splice(index, 1) +} + +// 模板按钮：填入标准客户端请求头名称（值留空），跳过已存在的同名行 +const fillHeaderOverrideTemplate = () => { + const existing = new Set( + headerOverrideRows.value.map((row) => row.name.trim().toLowerCase()).filter(Boolean) + ) + const rows = headerOverrideRows.value.filter((row) => row.name.trim() || row.value.trim()) + for (const row of getHeaderOverrideTemplate(form.platform)) { + if (!existing.has(row.name)) { + rows.push(row) + } + } + headerOverrideRows.value = rows +} const interceptWarmupRequests = ref(false) const autoPauseOnExpired = ref(true) const openaiPassthroughEnabled = ref(false) @@ -3569,6 +3702,7 @@ const vertexServiceAccountDragActive = ref(false) const tempUnschedEnabled = ref(false) const tempUnschedRules = ref([]) const getModelMappingKey = createStableObjectKeyResolver('create-model-mapping') +const getHeaderOverrideRowKey = createStableObjectKeyResolver('create-header-override-row') const getOpenAICompactModelMappingKey = createStableObjectKeyResolver('create-openai-compact-model-mapping') const getAntigravityModelMappingKey = createStableObjectKeyResolver('create-antigravity-model-mapping') const getTempUnschedRuleKey = createStableObjectKeyResolver('create-temp-unsched-rule') @@ -3970,6 +4104,10 @@ watch( anthropicAPIKeyAuthScheme.value = 'x_api_key' webSearchEmulationMode.value = 'default' } + // 请求头覆写为平台相关配置（模板/常用头集合不同），切换平台时清空， + // 避免上一平台的模板行被提交到新平台账号 + headerOverrideEnabled.value = false + headerOverrideRows.value = [] // Reset OAuth states oauth.resetState() openaiOAuth.resetState() @@ -4359,6 +4497,8 @@ const resetForm = () => { customErrorCodesEnabled.value = false selectedErrorCodes.value = [] customErrorCodeInput.value = null + headerOverrideEnabled.value = false + headerOverrideRows.value = [] interceptWarmupRequests.value = false autoPauseOnExpired.value = true openaiPassthroughEnabled.value = false @@ -4789,6 +4929,18 @@ const handleSubmit = async () => { credentials.custom_error_codes = [...selectedErrorCodes.value] } + // Add header override if enabled (anthropic/openai apikey only) + if (isHeaderOverridePlatform(form.platform)) { + if (headerOverrideEnabled.value) { + const headerError = validateHeaderOverrideRows(headerOverrideRows.value) + if (headerError) { + appStore.showError(t(`admin.accounts.headerOverride.${headerError}`)) + return + } + } + applyHeaderOverride(credentials, headerOverrideEnabled.value, headerOverrideRows.value, 'create') + } + applyInterceptWarmup(credentials, interceptWarmupRequests.value, 'create') if (!applyTempUnschedConfig(credentials)) { return diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 3670f233e3..9b5ab82fc5 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -417,6 +417,110 @@ + +
+
+
+ +

+ {{ t('admin.accounts.headerOverride.hint') }} +

+
+ +
+ +
+
+

+ + {{ t('admin.accounts.headerOverride.info') }} +

+
+ +
+
+ + + +
+
+ + + +
+ +
+ +

+ {{ t('admin.accounts.headerOverride.emptyValueHint') }} +

+
+
+ @@ -2433,7 +2537,15 @@ import ModelWhitelistSelector from '@/components/account/ModelWhitelistSelector. import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' import { applyAntigravityProjectID, - applyInterceptWarmup + applyHeaderOverride, + applyInterceptWarmup, + getHeaderOverrideTemplate, + isHeaderOverridePlatform, + splitHeaderOverridesObject, + validateHeaderOverrideRows, + HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY, + HEADER_OVERRIDES_CREDENTIAL_KEY, + type HeaderOverrideRow } from '@/components/account/credentialsBuilder' import { formatDateTime, formatDateTimeLocalInput, parseDateTimeLocalInput } from '@/utils/format' import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' @@ -2564,6 +2676,30 @@ function formatPoolModeRetryStatusCodes(value: unknown): string { const customErrorCodesEnabled = ref(false) const selectedErrorCodes = ref([]) const customErrorCodeInput = ref(null) +const headerOverrideEnabled = ref(false) +const headerOverrideRows = ref([]) + +const addHeaderOverrideRow = () => { + headerOverrideRows.value.push({ name: '', value: '' }) +} + +const removeHeaderOverrideRow = (index: number) => { + headerOverrideRows.value.splice(index, 1) +} + +// 模板按钮：填入标准客户端请求头名称（值留空），跳过已存在的同名行 +const fillHeaderOverrideTemplate = () => { + const existing = new Set( + headerOverrideRows.value.map((row) => row.name.trim().toLowerCase()).filter(Boolean) + ) + const rows = headerOverrideRows.value.filter((row) => row.name.trim() || row.value.trim()) + for (const row of getHeaderOverrideTemplate(props.account?.platform || '')) { + if (!existing.has(row.name)) { + rows.push(row) + } + } + headerOverrideRows.value = rows +} const interceptWarmupRequests = ref(false) const autoPauseOnExpired = ref(false) const autoPause5hThreshold = ref(null) @@ -2580,6 +2716,7 @@ const isSyncingAntigravityUpstream = ref(false) const tempUnschedEnabled = ref(false) const tempUnschedRules = ref([]) const getModelMappingKey = createStableObjectKeyResolver('edit-model-mapping') +const getHeaderOverrideRowKey = createStableObjectKeyResolver('edit-header-override-row') const getOpenAICompactModelMappingKey = createStableObjectKeyResolver('edit-openai-compact-model-mapping') const getAntigravityModelMappingKey = createStableObjectKeyResolver('edit-antigravity-model-mapping') const getTempUnschedRuleKey = createStableObjectKeyResolver('edit-temp-unsched-rule') @@ -3186,6 +3323,10 @@ const syncFormFromAccount = (newAccount: Account | null) => { loadTempUnschedRules(credentials) + // Reset header override state (loaded below only for apikey accounts) + headerOverrideEnabled.value = false + headerOverrideRows.value = [] + // Initialize API Key fields for apikey type if (newAccount.type === 'apikey' && newAccount.credentials) { const credentials = newAccount.credentials as Record @@ -3215,6 +3356,14 @@ const syncFormFromAccount = (newAccount: Account | null) => { } else { selectedErrorCodes.value = [] } + + // Load header override (anthropic/openai apikey only) + headerOverrideEnabled.value = + isHeaderOverridePlatform(newAccount.platform) && + credentials[HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY] === true + headerOverrideRows.value = splitHeaderOverridesObject( + credentials[HEADER_OVERRIDES_CREDENTIAL_KEY] + ) } else if (newAccount.type === 'bedrock' && newAccount.credentials) { const bedrockCreds = newAccount.credentials as Record const authMode = (bedrockCreds.auth_mode as string) || 'sigv4' @@ -3850,6 +3999,18 @@ const handleSubmit = async () => { delete newCredentials.custom_error_codes } + // Add header override if enabled (anthropic/openai apikey only) + if (isHeaderOverridePlatform(props.account.platform)) { + if (headerOverrideEnabled.value) { + const headerError = validateHeaderOverrideRows(headerOverrideRows.value) + if (headerError) { + appStore.showError(t(`admin.accounts.headerOverride.${headerError}`)) + return + } + } + applyHeaderOverride(newCredentials, headerOverrideEnabled.value, headerOverrideRows.value, 'edit') + } + // Add intercept warmup requests setting applyInterceptWarmup(newCredentials, interceptWarmupRequests.value, 'edit') if (!applyTempUnschedConfig(newCredentials)) { diff --git a/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts b/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts index 31f6e3bd26..d094f5366d 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透传') + expect(wrapper.text()).toContain('3.1-Flash-Image passthrough') 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/account/__tests__/credentialsBuilder.spec.ts b/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts index 665b1732e7..c2cb093805 100644 --- a/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts +++ b/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts @@ -1,8 +1,16 @@ import { describe, it, expect } from 'vitest' import { ANTIGRAVITY_PROJECT_ID_CREDENTIAL_KEY, + HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY, + HEADER_OVERRIDES_CREDENTIAL_KEY, applyAntigravityProjectID, - applyInterceptWarmup + applyHeaderOverride, + applyInterceptWarmup, + buildHeaderOverridesObject, + getHeaderOverrideTemplate, + isHeaderOverridePlatform, + splitHeaderOverridesObject, + validateHeaderOverrideRows } from '../credentialsBuilder' describe('applyInterceptWarmup', () => { @@ -82,3 +90,202 @@ describe('applyAntigravityProjectID', () => { expect(creds[ANTIGRAVITY_PROJECT_ID_CREDENTIAL_KEY]).toBe('configured-project') }) }) + +describe('isHeaderOverridePlatform', () => { + it('only anthropic and openai are supported', () => { + expect(isHeaderOverridePlatform('anthropic')).toBe(true) + expect(isHeaderOverridePlatform('openai')).toBe(true) + expect(isHeaderOverridePlatform('gemini')).toBe(false) + expect(isHeaderOverridePlatform('grok')).toBe(false) + expect(isHeaderOverridePlatform('antigravity')).toBe(false) + expect(isHeaderOverridePlatform('')).toBe(false) + }) +}) + +describe('validateHeaderOverrideRows', () => { + it('accepts valid rows and empty placeholder rows', () => { + expect( + validateHeaderOverrideRows([ + { name: 'user-agent', value: 'my-agent/1.0' }, + { name: 'x-app', value: '' }, + { name: '', value: '' } + ]) + ).toBeNull() + }) + + it('rejects empty name with non-empty value', () => { + expect(validateHeaderOverrideRows([{ name: '', value: 'v' }])).toBe('invalidName') + }) + + it('rejects invalid header names', () => { + expect(validateHeaderOverrideRows([{ name: 'bad name', value: '' }])).toBe('invalidName') + expect(validateHeaderOverrideRows([{ name: 'bad:name', value: '' }])).toBe('invalidName') + expect(validateHeaderOverrideRows([{ name: '名称', value: '' }])).toBe('invalidName') + }) + + it('rejects blocked header names case-insensitively', () => { + expect(validateHeaderOverrideRows([{ name: 'Authorization', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'X-Api-Key', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'host', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'Content-Length', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'Content-Type', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'Cookie', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'x-goog-api-key', value: '' }])).toBe('blockedName') + }) + + it('rejects duplicate names case-insensitively', () => { + expect( + validateHeaderOverrideRows([ + { name: 'User-Agent', value: 'a' }, + { name: 'user-agent', value: 'b' } + ]) + ).toBe('duplicateName') + }) +}) + +describe('buildHeaderOverridesObject / splitHeaderOverridesObject', () => { + it('lowercases names, trims values and drops empty-name rows', () => { + expect( + buildHeaderOverridesObject([ + { name: ' User-Agent ', value: ' my-agent ' }, + { name: 'X-App', value: '' }, + { name: '', value: 'ignored' } + ]) + ).toEqual({ 'user-agent': 'my-agent', 'x-app': '' }) + }) + + it('splits an object into sorted rows and ignores non-string values', () => { + expect( + splitHeaderOverridesObject({ 'x-app': 'cli', 'user-agent': 'ua', bogus: 42 }) + ).toEqual([ + { name: 'user-agent', value: 'ua' }, + { name: 'x-app', value: 'cli' } + ]) + expect(splitHeaderOverridesObject(null)).toEqual([]) + expect(splitHeaderOverridesObject(['a'])).toEqual([]) + expect(splitHeaderOverridesObject('str')).toEqual([]) + }) + + it('roundtrips through build and split', () => { + const rows = [ + { name: 'user-agent', value: 'ua' }, + { name: 'x-app', value: 'cli' } + ] + expect(splitHeaderOverridesObject(buildHeaderOverridesObject(rows))).toEqual(rows) + }) +}) + +describe('getHeaderOverrideTemplate', () => { + it('returns Claude Code CLI headers with empty values for anthropic', () => { + const rows = getHeaderOverrideTemplate('anthropic') + expect(rows.every((r) => r.value === '')).toBe(true) + const names = rows.map((r) => r.name) + expect(names).toContain('user-agent') + expect(names).toContain('x-app') + expect(names).toContain('anthropic-beta') + expect(names).toContain('x-stainless-lang') + expect(validateHeaderOverrideRows(rows)).toBeNull() + }) + + it('returns Codex CLI headers with empty values for openai', () => { + const rows = getHeaderOverrideTemplate('openai') + expect(rows.every((r) => r.value === '')).toBe(true) + const names = rows.map((r) => r.name) + expect(names).toContain('user-agent') + expect(names).toContain('originator') + expect(names).toContain('openai-beta') + expect(validateHeaderOverrideRows(rows)).toBeNull() + }) +}) + +describe('applyHeaderOverride', () => { + it('create + enabled: writes enabled flag and overrides object', () => { + const creds: Record = { api_key: 'sk' } + applyHeaderOverride(creds, true, [{ name: 'User-Agent', value: 'ua' }], 'create') + expect(creds[HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY]).toBe(true) + expect(creds[HEADER_OVERRIDES_CREDENTIAL_KEY]).toEqual({ 'user-agent': 'ua' }) + }) + + it('create + disabled: does not add fields', () => { + const creds: Record = { api_key: 'sk' } + applyHeaderOverride(creds, false, [{ name: 'user-agent', value: 'ua' }], 'create') + expect(HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY in creds).toBe(false) + expect(HEADER_OVERRIDES_CREDENTIAL_KEY in creds).toBe(false) + }) + + it('edit + disabled: deletes existing fields', () => { + const creds: Record = { + api_key: 'sk', + [HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY]: true, + [HEADER_OVERRIDES_CREDENTIAL_KEY]: { 'user-agent': 'ua' } + } + applyHeaderOverride(creds, false, [], 'edit') + expect(HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY in creds).toBe(false) + expect(HEADER_OVERRIDES_CREDENTIAL_KEY in creds).toBe(false) + expect(creds.api_key).toBe('sk') + }) + + it('edit + enabled: replaces overrides object wholesale', () => { + const creds: Record = { + [HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY]: true, + [HEADER_OVERRIDES_CREDENTIAL_KEY]: { 'x-old': 'old' } + } + applyHeaderOverride(creds, true, [{ name: 'x-new', value: 'new' }], 'edit') + expect(creds[HEADER_OVERRIDES_CREDENTIAL_KEY]).toEqual({ 'x-new': 'new' }) + }) +}) + +describe('validateHeaderOverrideRows value/entry limits', () => { + it('rejects websocket handshake headers', () => { + expect(validateHeaderOverrideRows([{ name: 'Sec-WebSocket-Key', value: '' }])).toBe( + 'blockedName' + ) + }) + + it('rejects control characters in values', () => { + expect(validateHeaderOverrideRows([{ name: 'x-app', value: 'a\x0bb' }])).toBe('invalidValue') + }) + + it('rejects oversized values', () => { + expect(validateHeaderOverrideRows([{ name: 'x-app', value: 'a'.repeat(8193) }])).toBe( + 'invalidValue' + ) + }) + + it('measures value length in UTF-8 bytes to match backend', () => { + // 3000 个 CJK 字符 = 3000 UTF-16 code units，但 9000 UTF-8 字节 > 8192 + expect(validateHeaderOverrideRows([{ name: 'x-app', value: '测'.repeat(3000) }])).toBe( + 'invalidValue' + ) + expect(validateHeaderOverrideRows([{ name: 'x-app', value: '测'.repeat(2000) }])).toBeNull() + }) + + it('rejects too many entries', () => { + const rows = Array.from({ length: 65 }, (_, i) => ({ name: `x-h-${i}`, value: 'v' })) + expect(validateHeaderOverrideRows(rows)).toBe('tooManyEntries') + }) +}) + +describe('validateHeaderOverrideRows session isolation headers', () => { + it('rejects per-request session headers', () => { + expect(validateHeaderOverrideRows([{ name: 'session_id', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'Conversation_ID', value: '' }])).toBe('blockedName') + expect(validateHeaderOverrideRows([{ name: 'x-codex-turn-state', value: '' }])).toBe( + 'blockedName' + ) + expect(validateHeaderOverrideRows([{ name: 'X-Claude-Code-Session-Id', value: '' }])).toBe( + 'blockedName' + ) + expect(validateHeaderOverrideRows([{ name: 'x-client-request-id', value: '' }])).toBe( + 'blockedName' + ) + }) + + it('allows tab inside value', () => { + expect(validateHeaderOverrideRows([{ name: 'x-app', value: 'a\tb' }])).toBeNull() + }) + + it('rejects oversized names', () => { + expect(validateHeaderOverrideRows([{ name: 'x'.repeat(201), value: 'v' }])).toBe('invalidName') + }) +}) diff --git a/frontend/src/components/account/credentialsBuilder.ts b/frontend/src/components/account/credentialsBuilder.ts index f138976519..3cdc0bdd54 100644 --- a/frontend/src/components/account/credentialsBuilder.ts +++ b/frontend/src/components/account/credentialsBuilder.ts @@ -24,3 +24,180 @@ export function applyAntigravityProjectID( delete credentials[ANTIGRAVITY_PROJECT_ID_CREDENTIAL_KEY] } } + +// ========== 请求头覆写（仅 anthropic/openai 平台的 api_key 账号） ========== + +export const HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY = 'header_override_enabled' +export const HEADER_OVERRIDES_CREDENTIAL_KEY = 'header_overrides' + +export interface HeaderOverrideRow { + name: string + value: string +} + +/** 请求头覆写支持的平台（与后端 IsHeaderOverrideEligible 保持一致） */ +export function isHeaderOverridePlatform(platform: string): boolean { + return platform === 'anthropic' || platform === 'openai' +} + +/** 禁止覆写的请求头（与后端 headerOverrideBlockedNames 保持一致） */ +const HEADER_OVERRIDE_BLOCKED_NAMES = new Set([ + 'host', + 'content-length', + 'content-type', + 'transfer-encoding', + 'connection', + 'keep-alive', + 'proxy-authenticate', + 'proxy-authorization', + 'proxy-connection', + 'te', + 'trailer', + 'upgrade', + 'authorization', + 'x-api-key', + 'x-goog-api-key', + 'cookie', + 'accept-encoding', + 'sec-websocket-key', + 'sec-websocket-version', + 'sec-websocket-extensions', + 'sec-websocket-protocol', + 'sec-websocket-accept', + 'session_id', + 'conversation_id', + 'x-codex-turn-state', + 'x-codex-turn-metadata', + 'chatgpt-account-id', + 'x-claude-code-session-id', + 'x-client-request-id' +]) + +/** RFC 7230 token：合法的 HTTP header 名称字符集 */ +const HEADER_NAME_PATTERN = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/ + +function isValidHeaderOverrideName(name: string): boolean { + return HEADER_NAME_PATTERN.test(name) +} + +/** 模板：Claude Code CLI 标准客户端请求头（值留空由管理员填写） */ +const ANTHROPIC_HEADER_OVERRIDE_TEMPLATE = [ + 'user-agent', + 'x-app', + 'anthropic-beta', + 'anthropic-version', + 'anthropic-dangerous-direct-browser-access', + 'x-stainless-lang', + 'x-stainless-package-version', + 'x-stainless-os', + 'x-stainless-arch', + 'x-stainless-runtime', + 'x-stainless-runtime-version', + 'x-stainless-retry-count', + 'x-stainless-timeout' +] + +/** 模板：Codex CLI 标准客户端请求头（值留空由管理员填写） */ +const OPENAI_HEADER_OVERRIDE_TEMPLATE = [ + 'user-agent', + 'originator', + 'openai-beta', + 'version', + 'accept', + 'accept-language' +] + +export function getHeaderOverrideTemplate(platform: string): HeaderOverrideRow[] { + const names = + platform === 'openai' ? OPENAI_HEADER_OVERRIDE_TEMPLATE : ANTHROPIC_HEADER_OVERRIDE_TEMPLATE + return names.map((name) => ({ name, value: '' })) +} + +/** 与后端 maxHeaderOverride* 常量保持一致 */ +const HEADER_OVERRIDE_MAX_ENTRIES = 64 +const HEADER_OVERRIDE_MAX_NAME_LENGTH = 200 +const HEADER_OVERRIDE_MAX_VALUE_LENGTH = 8192 + +/** header value 不允许包含控制字符（与后端 httpguts.ValidHeaderFieldValue 对齐） */ +// eslint-disable-next-line no-control-regex +const HEADER_VALUE_INVALID_PATTERN = /[\x00-\x08\x0a-\x1f\x7f]/ + +/** 长度限制按 UTF-8 字节计（与后端 Go len() 对齐，避免多字节值前端放行后端 400） */ +const HEADER_TEXT_ENCODER = new TextEncoder() +function utf8ByteLength(value: string): number { + return HEADER_TEXT_ENCODER.encode(value).length +} + +/** + * 校验请求头覆写行，返回首个错误的 i18n key（无错误返回 null）。 + * 名称为空但值非空 → invalidName；名称非法 → invalidName； + * 禁止覆写 → blockedName；大小写不敏感重名 → duplicateName； + * 值含控制字符或超长 → invalidValue；条目过多 → tooManyEntries。 + */ +export function validateHeaderOverrideRows( + rows: HeaderOverrideRow[] +): 'invalidName' | 'blockedName' | 'duplicateName' | 'invalidValue' | 'tooManyEntries' | null { + const seen = new Set() + for (const row of rows) { + const name = row.name.trim() + const value = row.value.trim() + if (!name) { + if (value) return 'invalidName' + continue + } + if (!isValidHeaderOverrideName(name) || name.length > HEADER_OVERRIDE_MAX_NAME_LENGTH) { + return 'invalidName' + } + const lower = name.toLowerCase() + if (HEADER_OVERRIDE_BLOCKED_NAMES.has(lower)) return 'blockedName' + if (seen.has(lower)) return 'duplicateName' + if ( + HEADER_VALUE_INVALID_PATTERN.test(value) || + utf8ByteLength(value) > HEADER_OVERRIDE_MAX_VALUE_LENGTH + ) { + return 'invalidValue' + } + seen.add(lower) + } + if (seen.size > HEADER_OVERRIDE_MAX_ENTRIES) return 'tooManyEntries' + return null +} + +/** 行数组 → credentials 存储对象（名称小写化，丢弃空行） */ +export function buildHeaderOverridesObject(rows: HeaderOverrideRow[]): Record { + const result: Record = {} + for (const row of rows) { + const name = row.name.trim().toLowerCase() + if (!name) continue + result[name] = row.value.trim() + } + return result +} + +/** credentials 存储对象 → 行数组（按名称排序保证稳定展示） */ +export function splitHeaderOverridesObject(record: unknown): HeaderOverrideRow[] { + if (!record || typeof record !== 'object' || Array.isArray(record)) return [] + return Object.entries(record as Record) + .filter(([, value]) => typeof value === 'string') + .map(([name, value]) => ({ name, value: value as string })) + .sort((a, b) => a.name.localeCompare(b.name)) +} + +/** + * 将请求头覆写写入 credentials。 + * create 模式：关闭时不写入任何字段；edit 模式：关闭时删除字段（全量替换语义）。 + */ +export function applyHeaderOverride( + credentials: Record, + enabled: boolean, + rows: HeaderOverrideRow[], + mode: 'create' | 'edit' +): void { + if (enabled) { + credentials[HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY] = true + credentials[HEADER_OVERRIDES_CREDENTIAL_KEY] = buildHeaderOverridesObject(rows) + } else if (mode === 'edit') { + delete credentials[HEADER_OVERRIDE_ENABLED_CREDENTIAL_KEY] + delete credentials[HEADER_OVERRIDES_CREDENTIAL_KEY] + } +} diff --git a/frontend/src/components/admin/account/AccountTestModal.vue b/frontend/src/components/admin/account/AccountTestModal.vue index ca9d06ee1e..0a0e3dd9ae 100644 --- a/frontend/src/components/admin/account/AccountTestModal.vue +++ b/frontend/src/components/admin/account/AccountTestModal.vue @@ -55,6 +55,17 @@ /> +
+ +
-
+
+ + +

+ {{ t('payment.admin.subscriptionCnyPayPreview', { amount: subscriptionCnyPreview.amount }) }} + + {{ t('payment.admin.subscriptionCnyPayPreviewWithFee', { feeRate: subscriptionCnyPreview.feeRate, total: subscriptionCnyPreview.total }) }} + +

+
@@ -81,7 +90,9 @@ import { ref, reactive, computed, watch } from 'vue' import { useI18n } from 'vue-i18n' import { useAppStore } from '@/stores/app' import { adminPaymentAPI } from '@/api/admin/payment' +import type { AdminPaymentConfig } from '@/api/admin/payment' import { extractApiErrorMessage } from '@/utils/apiError' +import { formatPaymentAmount } from '@/components/payment/currency' import type { SubscriptionPlan } from '@/types/payment' import type { AdminGroup } from '@/types' import BaseDialog from '@/components/common/BaseDialog.vue' @@ -94,6 +105,7 @@ const props = defineProps<{ show: boolean plan: SubscriptionPlan | null groups: AdminGroup[] + paymentConfig?: AdminPaymentConfig | null }>() const emit = defineEmits<{ @@ -129,6 +141,31 @@ const selectedGroupInfo = computed(() => { return props.groups.find(g => g.id === planForm.group_id) || null }) +function roundCnyAmount(value: number): number { + return Math.round(value * 100) / 100 +} + +function ceilCnyAmount(value: number): number { + return Math.ceil(value * 100) / 100 +} + +const subscriptionCnyPreview = computed(() => { + const price = Number(planForm.price) || 0 + const rate = Number(props.paymentConfig?.subscription_usd_to_cny_rate) || 0 + if (price <= 0 || rate <= 0) return null + + const amount = roundCnyAmount(price * rate) + const feeRate = Number(props.paymentConfig?.recharge_fee_rate) || 0 + const fee = feeRate > 0 ? ceilCnyAmount((amount * feeRate) / 100) : 0 + const total = feeRate > 0 ? roundCnyAmount(amount + fee) : amount + + return { + amount: formatPaymentAmount(amount, 'CNY'), + feeRate, + total: formatPaymentAmount(total, 'CNY'), + } +}) + // Reset form when dialog opens watch(() => props.show, (visible) => { if (!visible) return diff --git a/frontend/src/views/admin/orders/__tests__/PlanEditDialog.spec.ts b/frontend/src/views/admin/orders/__tests__/PlanEditDialog.spec.ts new file mode 100644 index 0000000000..9c31e7c176 --- /dev/null +++ b/frontend/src/views/admin/orders/__tests__/PlanEditDialog.spec.ts @@ -0,0 +1,77 @@ +import { describe, expect, it, vi } from 'vitest' +import { mount } from '@vue/test-utils' +import PlanEditDialog from '../PlanEditDialog.vue' + +vi.mock('vue-i18n', () => ({ + useI18n: () => ({ + t: (key: string, params?: Record) => { + if (key === 'payment.admin.subscriptionCnyPayPreview') return `preview ${params?.amount}` + if (key === 'payment.admin.subscriptionCnyPayPreviewWithFee') return `fee ${params?.feeRate} ${params?.total}` + return key + }, + }), +})) + +vi.mock('@/stores/app', () => ({ + useAppStore: () => ({ + showError: vi.fn(), + showSuccess: vi.fn(), + }), +})) + +vi.mock('@/api/admin/payment', () => ({ + adminPaymentAPI: { + createPlan: vi.fn(), + updatePlan: vi.fn(), + }, +})) + +function mountDialog(paymentConfig: Record | null) { + return mount(PlanEditDialog, { + props: { + show: true, + plan: null, + groups: [], + paymentConfig, + }, + global: { + stubs: { + BaseDialog: { + props: ['show'], + template: '
', + }, + Select: true, + Icon: true, + GroupBadge: true, + }, + }, + }) +} + +describe('PlanEditDialog subscription CNY payment preview', () => { + it('shows CNY channel charge using the configured subscription rate and fee', async () => { + const wrapper = mountDialog({ + subscription_usd_to_cny_rate: 7.15, + recharge_fee_rate: 2.5, + }) + + await wrapper.find('input[type="number"]').setValue('9.99') + + expect(wrapper.text()).toContain('preview') + expect(wrapper.text()).toContain('¥71.43') + expect(wrapper.text()).toContain('fee 2.5') + expect(wrapper.text()).toContain('¥73.22') + }) + + it('hides the preview when the subscription rate is not configured', async () => { + const wrapper = mountDialog({ + subscription_usd_to_cny_rate: 0, + recharge_fee_rate: 2.5, + }) + + await wrapper.find('input[type="number"]').setValue('9.99') + + expect(wrapper.text()).not.toContain('preview') + expect(wrapper.text()).not.toContain('¥71.43') + }) +}) diff --git a/frontend/src/views/user/BatchImageGuideView.vue b/frontend/src/views/user/BatchImageGuideView.vue new file mode 100644 index 0000000000..70387766a8 --- /dev/null +++ b/frontend/src/views/user/BatchImageGuideView.vue @@ -0,0 +1,2799 @@ +