package controlplane import ( "context" "errors" "fmt" "strings" "time" "github.com/jackc/pgx/v5" ) const maxLowBalanceThresholdMicros int64 = 1_000_000_000_000_000 func (s *Store) GetTenantPreferences(ctx context.Context, tenantID string, defaultThresholdMicros int64) (TenantPreferences, error) { if defaultThresholdMicros <= 0 { defaultThresholdMicros = 5_000_000 } result := TenantPreferences{TenantID: strings.TrimSpace(tenantID), LowBalanceEnabled: true, LowBalanceThresholdMicros: defaultThresholdMicros} if result.TenantID == "" { return result, nil } var defaultModel, fallbackModel *string var updatedAt time.Time err := s.db.QueryRow(ctx, ` SELECT default_model, fallback_model, low_balance_enabled, low_balance_threshold_micros, updated_at FROM tenant_preferences WHERE tenant_id=$1`, result.TenantID).Scan( &defaultModel, &fallbackModel, &result.LowBalanceEnabled, &result.LowBalanceThresholdMicros, &updatedAt) if errors.Is(err, pgx.ErrNoRows) { return result, nil } if err != nil { return TenantPreferences{}, fmt.Errorf("query tenant preferences: %w", err) } if defaultModel != nil { result.DefaultModel = *defaultModel } if fallbackModel != nil { result.FallbackModel = *fallbackModel } result.UpdatedAt = &updatedAt return result, nil } func (s *Store) SetDeveloperPreferences(ctx context.Context, input SetDeveloperPreferencesInput) (TenantPreferences, error) { input.TenantID = strings.TrimSpace(input.TenantID) input.DefaultModel = strings.TrimSpace(input.DefaultModel) input.FallbackModel = strings.TrimSpace(input.FallbackModel) if input.TenantID == "" { return TenantPreferences{}, errors.New("tenant_id is required") } if input.DefaultModel != "" && input.DefaultModel == input.FallbackModel { return TenantPreferences{}, errors.New("default_model and fallback_model must be different") } available, err := s.ListDeveloperModels(ctx, input.TenantID) if err != nil { return TenantPreferences{}, err } allowed := make(map[string]struct{}, len(available)) for _, model := range available { allowed[model.PublicID] = struct{}{} } for field, model := range map[string]string{"default_model": input.DefaultModel, "fallback_model": input.FallbackModel} { if model != "" { if _, ok := allowed[model]; !ok { return TenantPreferences{}, fmt.Errorf("%s is not available to this tenant", field) } } } tx, err := s.db.Begin(ctx) if err != nil { return TenantPreferences{}, err } defer tx.Rollback(ctx) if err := tx.QueryRow(ctx, `SELECT id FROM tenants WHERE id=$1`, input.TenantID).Scan(new(string)); err != nil { if errors.Is(err, pgx.ErrNoRows) { return TenantPreferences{}, ErrNotFound } return TenantPreferences{}, err } var result TenantPreferences var defaultModel, fallbackModel *string var updatedAt time.Time if err := tx.QueryRow(ctx, ` INSERT INTO tenant_preferences (tenant_id, default_model, fallback_model) VALUES ($1, NULLIF($2,''), NULLIF($3,'')) ON CONFLICT (tenant_id) DO UPDATE SET default_model=EXCLUDED.default_model, fallback_model=EXCLUDED.fallback_model, updated_at=now() RETURNING tenant_id::text, default_model, fallback_model, low_balance_enabled, low_balance_threshold_micros, updated_at`, input.TenantID, input.DefaultModel, input.FallbackModel).Scan( &result.TenantID, &defaultModel, &fallbackModel, &result.LowBalanceEnabled, &result.LowBalanceThresholdMicros, &updatedAt); err != nil { return TenantPreferences{}, fmt.Errorf("save developer preferences: %w", err) } if defaultModel != nil { result.DefaultModel = *defaultModel } if fallbackModel != nil { result.FallbackModel = *fallbackModel } result.UpdatedAt = &updatedAt if err := tx.Commit(ctx); err != nil { return TenantPreferences{}, err } return result, nil } func (s *Store) SetBillingPreferences(ctx context.Context, input SetBillingPreferencesInput, defaultThresholdMicros int64) (TenantPreferences, error) { input.TenantID = strings.TrimSpace(input.TenantID) if input.TenantID == "" { return TenantPreferences{}, errors.New("tenant_id is required") } if defaultThresholdMicros <= 0 { defaultThresholdMicros = 5_000_000 } if input.LowBalanceThresholdMicros != nil && (*input.LowBalanceThresholdMicros < 0 || *input.LowBalanceThresholdMicros > maxLowBalanceThresholdMicros) { return TenantPreferences{}, errors.New("low balance threshold is outside the supported range") } tx, err := s.db.Begin(ctx) if err != nil { return TenantPreferences{}, err } defer tx.Rollback(ctx) if err := tx.QueryRow(ctx, `SELECT id FROM tenants WHERE id=$1`, input.TenantID).Scan(new(string)); err != nil { if errors.Is(err, pgx.ErrNoRows) { return TenantPreferences{}, ErrNotFound } return TenantPreferences{}, err } var result TenantPreferences var defaultModel, fallbackModel *string var updatedAt time.Time if err := tx.QueryRow(ctx, ` INSERT INTO tenant_preferences (tenant_id, low_balance_enabled, low_balance_threshold_micros) VALUES ($1, COALESCE($2::boolean, TRUE), COALESCE($3::bigint, $4::bigint)) ON CONFLICT (tenant_id) DO UPDATE SET low_balance_enabled=COALESCE($2::boolean, tenant_preferences.low_balance_enabled), low_balance_threshold_micros=COALESCE($3::bigint, tenant_preferences.low_balance_threshold_micros), updated_at=now() RETURNING tenant_id::text, default_model, fallback_model, low_balance_enabled, low_balance_threshold_micros, updated_at`, input.TenantID, input.LowBalanceEnabled, input.LowBalanceThresholdMicros, defaultThresholdMicros).Scan( &result.TenantID, &defaultModel, &fallbackModel, &result.LowBalanceEnabled, &result.LowBalanceThresholdMicros, &updatedAt); err != nil { return TenantPreferences{}, fmt.Errorf("save billing preferences: %w", err) } if defaultModel != nil { result.DefaultModel = *defaultModel } if fallbackModel != nil { result.FallbackModel = *fallbackModel } result.UpdatedAt = &updatedAt if err := tx.Commit(ctx); err != nil { return TenantPreferences{}, err } return result, nil }