| package sql |
|
|
| import ( |
| "context" |
| "database/sql" |
| "errors" |
| "fmt" |
| "strings" |
| "sync" |
| "time" |
|
|
| "ccLoad/internal/model" |
| ) |
|
|
| |
|
|
| |
| func (s *SQLStore) ListConfigs(ctx context.Context) ([]*model.Config, error) { |
| |
| |
| |
| query := ` |
| SELECT c.id, c.name, c.url, c.priority, c.rpm_limit, c.channel_type, c.protocol_transform_mode, c.enabled, |
| c.scheduled_check_enabled, c.scheduled_check_model, |
| c.cooldown_until, c.cooldown_duration_ms, c.daily_cost_limit, c.cost_multiplier, c.custom_request_rules, |
| SUM(CASE WHEN k.id IS NOT NULL AND k.disabled = 0 THEN 1 ELSE 0 END) as key_count, |
| c.created_at, c.updated_at |
| FROM channels c |
| LEFT JOIN api_keys k ON c.id = k.channel_id |
| GROUP BY c.id |
| ORDER BY c.priority DESC, c.id ASC |
| ` |
| rows, err := s.db.QueryContext(ctx, query) |
| if err != nil { |
| return nil, err |
| } |
| defer func() { _ = rows.Close() }() |
|
|
| |
| scanner := NewConfigScanner() |
| configs, err := scanner.ScanConfigs(rows) |
| if err != nil { |
| return nil, err |
| } |
|
|
| if err := s.loadConfigsAuxConcurrent(ctx, configs); err != nil { |
| return nil, err |
| } |
|
|
| return configs, nil |
| } |
|
|
| |
| func (s *SQLStore) GetConfig(ctx context.Context, id int64) (*model.Config, error) { |
| |
| |
| query := ` |
| SELECT c.id, c.name, c.url, c.priority, c.rpm_limit, c.channel_type, c.protocol_transform_mode, c.enabled, |
| c.scheduled_check_enabled, c.scheduled_check_model, |
| c.cooldown_until, c.cooldown_duration_ms, c.daily_cost_limit, c.cost_multiplier, c.custom_request_rules, |
| SUM(CASE WHEN k.id IS NOT NULL AND k.disabled = 0 THEN 1 ELSE 0 END) as key_count, |
| c.created_at, c.updated_at |
| FROM channels c |
| LEFT JOIN api_keys k ON c.id = k.channel_id |
| WHERE c.id = ? |
| GROUP BY c.id |
| ` |
| row := s.db.QueryRowContext(ctx, query, id) |
|
|
| |
| scanner := NewConfigScanner() |
| config, err := scanner.ScanConfig(row) |
| if err != nil { |
| if errors.Is(err, sql.ErrNoRows) { |
| return nil, errors.New("not found") |
| } |
| return nil, err |
| } |
|
|
| if err := s.loadConfigsAuxConcurrent(ctx, []*model.Config{config}); err != nil { |
| return nil, err |
| } |
|
|
| return config, nil |
| } |
|
|
| |
| func (s *SQLStore) GetEnabledChannelsByModel(ctx context.Context, modelName string) ([]*model.Config, error) { |
| var query string |
| var args []any |
|
|
| if modelName == "*" { |
| |
| |
| query = ` |
| SELECT c.id, c.name, c.url, c.priority, c.rpm_limit, |
| c.channel_type, c.protocol_transform_mode, c.enabled, c.scheduled_check_enabled, c.scheduled_check_model, |
| c.cooldown_until, c.cooldown_duration_ms, c.daily_cost_limit, c.cost_multiplier, c.custom_request_rules, |
| SUM(CASE WHEN k.id IS NOT NULL AND k.disabled = 0 THEN 1 ELSE 0 END) as key_count, |
| c.created_at, c.updated_at |
| FROM channels c |
| LEFT JOIN api_keys k ON c.id = k.channel_id |
| WHERE c.enabled = 1 |
| GROUP BY c.id |
| ORDER BY c.priority DESC, c.id ASC |
| ` |
| } else { |
| |
| query = ` |
| SELECT c.id, c.name, c.url, c.priority, c.rpm_limit, |
| c.channel_type, c.protocol_transform_mode, c.enabled, c.scheduled_check_enabled, c.scheduled_check_model, |
| c.cooldown_until, c.cooldown_duration_ms, c.daily_cost_limit, c.cost_multiplier, c.custom_request_rules, |
| SUM(CASE WHEN k.id IS NOT NULL AND k.disabled = 0 THEN 1 ELSE 0 END) as key_count, |
| c.created_at, c.updated_at |
| FROM channels c |
| INNER JOIN channel_models cm ON c.id = cm.channel_id |
| LEFT JOIN api_keys k ON c.id = k.channel_id |
| WHERE c.enabled = 1 |
| AND cm.model = ? |
| GROUP BY c.id |
| ORDER BY c.priority DESC, c.id ASC |
| ` |
| args = []any{modelName} |
| } |
|
|
| rows, err := s.db.QueryContext(ctx, query, args...) |
| if err != nil { |
| return nil, err |
| } |
| defer func() { _ = rows.Close() }() |
|
|
| scanner := NewConfigScanner() |
| configs, err := scanner.ScanConfigs(rows) |
| if err != nil { |
| return nil, err |
| } |
|
|
| |
| if err := s.loadConfigsAuxConcurrent(ctx, configs); err != nil { |
| return nil, err |
| } |
|
|
| return configs, nil |
| } |
|
|
| |
| func (s *SQLStore) GetEnabledChannelsByType(ctx context.Context, channelType string) ([]*model.Config, error) { |
| |
| query := ` |
| SELECT c.id, c.name, c.url, c.priority, c.rpm_limit, |
| c.channel_type, c.protocol_transform_mode, c.enabled, c.scheduled_check_enabled, c.scheduled_check_model, |
| c.cooldown_until, c.cooldown_duration_ms, c.daily_cost_limit, c.cost_multiplier, c.custom_request_rules, |
| SUM(CASE WHEN k.id IS NOT NULL AND k.disabled = 0 THEN 1 ELSE 0 END) as key_count, |
| c.created_at, c.updated_at |
| FROM channels c |
| LEFT JOIN api_keys k ON c.id = k.channel_id |
| WHERE c.enabled = 1 |
| AND c.channel_type = ? |
| GROUP BY c.id |
| ORDER BY c.priority DESC, c.id ASC |
| ` |
|
|
| rows, err := s.db.QueryContext(ctx, query, channelType) |
| if err != nil { |
| return nil, err |
| } |
| defer func() { _ = rows.Close() }() |
|
|
| scanner := NewConfigScanner() |
| configs, err := scanner.ScanConfigs(rows) |
| if err != nil { |
| return nil, err |
| } |
|
|
| |
| if err := s.loadConfigsAuxConcurrent(ctx, configs); err != nil { |
| return nil, err |
| } |
|
|
| return configs, nil |
| } |
|
|
| |
| func (s *SQLStore) GetEnabledChannelsByModelAndProtocol(ctx context.Context, modelName string, protocol string) ([]*model.Config, error) { |
| protocol = strings.TrimSpace(strings.ToLower(protocol)) |
| if protocol == "" { |
| return s.GetEnabledChannelsByModel(ctx, modelName) |
| } |
|
|
| args := []any{protocol, protocol} |
| query := ` |
| SELECT c.id, c.name, c.url, c.priority, c.rpm_limit, |
| c.channel_type, c.protocol_transform_mode, c.enabled, c.scheduled_check_enabled, c.scheduled_check_model, |
| c.cooldown_until, c.cooldown_duration_ms, c.daily_cost_limit, c.cost_multiplier, c.custom_request_rules, |
| SUM(CASE WHEN k.id IS NOT NULL AND k.disabled = 0 THEN 1 ELSE 0 END) as key_count, |
| c.created_at, c.updated_at |
| FROM channels c |
| LEFT JOIN api_keys k ON c.id = k.channel_id |
| WHERE c.enabled = 1 |
| AND ( |
| c.channel_type = ? |
| OR EXISTS ( |
| SELECT 1 |
| FROM channel_protocol_transforms cpt |
| WHERE cpt.channel_id = c.id AND cpt.protocol = ? |
| ) |
| ) |
| ` |
|
|
| if modelName != "*" { |
| query += ` |
| AND EXISTS ( |
| SELECT 1 |
| FROM channel_models cm |
| WHERE cm.channel_id = c.id AND cm.model = ? |
| ) |
| ` |
| args = append(args, modelName) |
| } |
|
|
| query += ` |
| GROUP BY c.id |
| ORDER BY c.priority DESC, c.id ASC |
| ` |
|
|
| rows, err := s.db.QueryContext(ctx, query, args...) |
| if err != nil { |
| return nil, err |
| } |
| defer func() { _ = rows.Close() }() |
|
|
| scanner := NewConfigScanner() |
| configs, err := scanner.ScanConfigs(rows) |
| if err != nil { |
| return nil, err |
| } |
|
|
| if err := s.loadConfigsAuxConcurrent(ctx, configs); err != nil { |
| return nil, err |
| } |
|
|
| configs = filterConfigsByProtocol(configs, protocol) |
| return configs, nil |
| } |
|
|
| |
| func (s *SQLStore) GetEnabledChannelsByExposedProtocol(ctx context.Context, protocol string) ([]*model.Config, error) { |
| protocol = strings.TrimSpace(strings.ToLower(protocol)) |
| if protocol == "" { |
| return []*model.Config{}, nil |
| } |
| return s.GetEnabledChannelsByModelAndProtocol(ctx, "*", protocol) |
| } |
|
|
| |
| func (s *SQLStore) CreateConfig(ctx context.Context, c *model.Config) (*model.Config, error) { |
| nowUnix := timeToUnix(time.Now()) |
|
|
| |
| channelType := c.GetChannelType() |
| protocolTransformMode := c.GetProtocolTransformMode() |
| customRules, err := marshalCustomRequestRules(c.CustomRequestRules) |
| if err != nil { |
| return nil, err |
| } |
|
|
| id := c.ID |
| err = s.WithTransaction(ctx, func(tx *sql.Tx) error { |
| if id == 0 { |
| |
| res, err := tx.ExecContext(ctx, ` |
| INSERT INTO channels(name, url, priority, rpm_limit, channel_type, protocol_transform_mode, enabled, scheduled_check_enabled, scheduled_check_model, daily_cost_limit, cost_multiplier, custom_request_rules, created_at, updated_at) |
| VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| `, c.Name, c.URL, c.Priority, c.RPMLimit, channelType, protocolTransformMode, |
| boolToInt(c.Enabled), boolToInt(c.ScheduledCheckEnabled), c.ScheduledCheckModel, c.DailyCostLimit, normalizeCostMultiplier(c.CostMultiplier), customRules, nowUnix, nowUnix) |
| if err != nil { |
| return err |
| } |
|
|
| id, err = res.LastInsertId() |
| if err != nil { |
| return fmt.Errorf("get last insert id: %w", err) |
| } |
| } else { |
| |
| if s.IsSQLite() { |
| _, err := tx.ExecContext(ctx, ` |
| INSERT INTO channels(id, name, url, priority, rpm_limit, channel_type, protocol_transform_mode, enabled, scheduled_check_enabled, scheduled_check_model, daily_cost_limit, cost_multiplier, custom_request_rules, created_at, updated_at) |
| VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| `, id, c.Name, c.URL, c.Priority, c.RPMLimit, channelType, protocolTransformMode, |
| boolToInt(c.Enabled), boolToInt(c.ScheduledCheckEnabled), c.ScheduledCheckModel, c.DailyCostLimit, normalizeCostMultiplier(c.CostMultiplier), customRules, nowUnix, nowUnix) |
| if err != nil { |
| return err |
| } |
| } else { |
| _, err := tx.ExecContext(ctx, ` |
| INSERT INTO channels(id, name, url, priority, rpm_limit, channel_type, protocol_transform_mode, enabled, scheduled_check_enabled, scheduled_check_model, daily_cost_limit, cost_multiplier, custom_request_rules, created_at, updated_at) |
| VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| ON DUPLICATE KEY UPDATE |
| name = VALUES(name), |
| url = VALUES(url), |
| priority = VALUES(priority), |
| rpm_limit = VALUES(rpm_limit), |
| channel_type = VALUES(channel_type), |
| protocol_transform_mode = VALUES(protocol_transform_mode), |
| enabled = VALUES(enabled), |
| scheduled_check_enabled = VALUES(scheduled_check_enabled), |
| scheduled_check_model = VALUES(scheduled_check_model), |
| daily_cost_limit = VALUES(daily_cost_limit), |
| cost_multiplier = VALUES(cost_multiplier), |
| custom_request_rules = VALUES(custom_request_rules), |
| updated_at = VALUES(updated_at) |
| `, id, c.Name, c.URL, c.Priority, c.RPMLimit, channelType, protocolTransformMode, |
| boolToInt(c.Enabled), boolToInt(c.ScheduledCheckEnabled), c.ScheduledCheckModel, c.DailyCostLimit, normalizeCostMultiplier(c.CostMultiplier), customRules, nowUnix, nowUnix) |
| if err != nil { |
| return err |
| } |
| } |
| } |
|
|
| |
| if err := s.saveModelEntriesTx(ctx, tx, id, c.ModelEntries); err != nil { |
| return fmt.Errorf("save model entries: %w", err) |
| } |
| if err := s.saveProtocolTransformsTx(ctx, tx, id, c.GetProtocolTransforms()); err != nil { |
| return fmt.Errorf("save protocol transforms: %w", err) |
| } |
|
|
| return nil |
| }) |
| if err != nil { |
| return nil, err |
| } |
|
|
| |
| config, err := s.GetConfig(ctx, id) |
| if err != nil { |
| return nil, err |
| } |
|
|
| return config, nil |
| } |
|
|
| |
| func (s *SQLStore) UpdateConfig(ctx context.Context, id int64, upd *model.Config) (*model.Config, error) { |
| if upd == nil { |
| return nil, errors.New("update payload cannot be nil") |
| } |
|
|
| |
| if _, err := s.GetConfig(ctx, id); err != nil { |
| return nil, err |
| } |
|
|
| name := strings.TrimSpace(upd.Name) |
| url := strings.TrimSpace(upd.URL) |
|
|
| |
| channelType := upd.GetChannelType() |
| protocolTransformMode := upd.GetProtocolTransformMode() |
| customRules, err := marshalCustomRequestRules(upd.CustomRequestRules) |
| if err != nil { |
| return nil, err |
| } |
| updatedAtUnix := timeToUnix(time.Now()) |
|
|
| err = s.WithTransaction(ctx, func(tx *sql.Tx) error { |
| |
| _, err := tx.ExecContext(ctx, ` |
| UPDATE channels |
| SET name=?, url=?, priority=?, rpm_limit=?, channel_type=?, protocol_transform_mode=?, enabled=?, scheduled_check_enabled=?, scheduled_check_model=?, daily_cost_limit=?, cost_multiplier=?, custom_request_rules=?, updated_at=? |
| WHERE id=? |
| `, name, url, upd.Priority, upd.RPMLimit, channelType, protocolTransformMode, |
| boolToInt(upd.Enabled), boolToInt(upd.ScheduledCheckEnabled), upd.ScheduledCheckModel, upd.DailyCostLimit, normalizeCostMultiplier(upd.CostMultiplier), customRules, updatedAtUnix, id) |
| if err != nil { |
| return err |
| } |
|
|
| |
| if err := s.saveModelEntriesTx(ctx, tx, id, upd.ModelEntries); err != nil { |
| return fmt.Errorf("save model entries: %w", err) |
| } |
| if err := s.saveProtocolTransformsTx(ctx, tx, id, upd.GetProtocolTransforms()); err != nil { |
| return fmt.Errorf("save protocol transforms: %w", err) |
| } |
|
|
| return nil |
| }) |
| if err != nil { |
| return nil, err |
| } |
|
|
| |
| config, err := s.GetConfig(ctx, id) |
| if err != nil { |
| return nil, err |
| } |
|
|
| return config, nil |
| } |
|
|
| |
| |
| |
| func (s *SQLStore) UpdateChannelEnabled(ctx context.Context, id int64, enabled bool) (*model.Config, error) { |
| updatedAtUnix := timeToUnix(time.Now()) |
| result, err := s.db.ExecContext(ctx, ` |
| UPDATE channels |
| SET enabled = ?, updated_at = ? |
| WHERE id = ? |
| `, boolToInt(enabled), updatedAtUnix, id) |
| if err != nil { |
| return nil, fmt.Errorf("update channel enabled: %w", err) |
| } |
|
|
| rowsAffected, err := result.RowsAffected() |
| if err == nil && rowsAffected == 0 { |
| cfg, getErr := s.GetConfig(ctx, id) |
| if getErr != nil { |
| return nil, getErr |
| } |
| return cfg, nil |
| } |
|
|
| config, err := s.GetConfig(ctx, id) |
| if err != nil { |
| return nil, err |
| } |
| return config, nil |
| } |
|
|
| |
| func (s *SQLStore) DeleteConfig(ctx context.Context, id int64) error { |
| |
| if _, err := s.GetConfig(ctx, id); err != nil { |
| if strings.Contains(err.Error(), "not found") { |
| return nil |
| } |
| return err |
| } |
|
|
| |
| err := s.WithTransaction(ctx, func(tx *sql.Tx) error { |
| if _, err := tx.ExecContext(ctx, `DELETE FROM api_keys WHERE channel_id = ?`, id); err != nil { |
| return fmt.Errorf("delete channel api keys: %w", err) |
| } |
| if _, err := tx.ExecContext(ctx, `DELETE FROM channel_models WHERE channel_id = ?`, id); err != nil { |
| return fmt.Errorf("delete channel models: %w", err) |
| } |
| if _, err := tx.ExecContext(ctx, `DELETE FROM logs WHERE channel_id = ?`, id); err != nil { |
| return fmt.Errorf("delete channel logs: %w", err) |
| } |
| if _, err := tx.ExecContext(ctx, `DELETE FROM channels WHERE id = ?`, id); err != nil { |
| return fmt.Errorf("delete channel: %w", err) |
| } |
| return nil |
| }) |
| if err != nil { |
| return err |
| } |
|
|
| return nil |
| } |
|
|
| |
| |
| func (s *SQLStore) BatchUpdatePriority(ctx context.Context, updates []struct { |
| ID int64 |
| Priority int |
| }) (int64, error) { |
| if len(updates) == 0 { |
| return 0, nil |
| } |
|
|
| updatedAtUnix := timeToUnix(time.Now()) |
|
|
| |
| var caseBuilder strings.Builder |
| |
| args := make([]any, 0, len(updates)*2+1+len(updates)) |
|
|
| caseBuilder.WriteString("UPDATE channels SET priority = CASE id ") |
| for _, update := range updates { |
| caseBuilder.WriteString("WHEN ? THEN ? ") |
| args = append(args, update.ID, update.Priority) |
| } |
| caseBuilder.WriteString("END, updated_at = ? WHERE id IN (") |
| args = append(args, updatedAtUnix) |
|
|
| for i, update := range updates { |
| if i > 0 { |
| caseBuilder.WriteString(",") |
| } |
| caseBuilder.WriteString("?") |
| args = append(args, update.ID) |
| } |
| caseBuilder.WriteString(")") |
|
|
| |
| result, err := s.db.ExecContext(ctx, caseBuilder.String(), args...) |
| if err != nil { |
| return 0, fmt.Errorf("batch update priority: %w", err) |
| } |
|
|
| rowsAffected, _ := result.RowsAffected() |
|
|
| return rowsAffected, nil |
| } |
|
|
| |
|
|
| |
| |
| |
| |
| |
| func (s *SQLStore) loadModelEntriesForConfigs(ctx context.Context, configs []*model.Config) error { |
| if len(configs) == 0 { |
| return nil |
| } |
|
|
| |
| channelIDs := make([]any, len(configs)) |
| placeholders := make([]string, len(configs)) |
| idToConfig := make(map[int64]*model.Config) |
| for i, cfg := range configs { |
| channelIDs[i] = cfg.ID |
| placeholders[i] = "?" |
| idToConfig[cfg.ID] = cfg |
| cfg.ModelEntries = nil |
| } |
|
|
| |
| query := fmt.Sprintf( |
| `SELECT channel_id, model, redirect_model FROM channel_models WHERE channel_id IN (%s) ORDER BY channel_id, created_at ASC, model ASC`, |
| strings.Join(placeholders, ","), |
| ) |
|
|
| rows, err := s.db.QueryContext(ctx, query, channelIDs...) |
| if err != nil { |
| return fmt.Errorf("query model entries: %w", err) |
| } |
| defer func() { _ = rows.Close() }() |
|
|
| for rows.Next() { |
| var channelID int64 |
| var entry model.ModelEntry |
| if err := rows.Scan(&channelID, &entry.Model, &entry.RedirectModel); err != nil { |
| return fmt.Errorf("scan model entry: %w", err) |
| } |
| if cfg, ok := idToConfig[channelID]; ok { |
| cfg.ModelEntries = append(cfg.ModelEntries, entry) |
| } |
| } |
|
|
| return rows.Err() |
| } |
|
|
| func (s *SQLStore) loadProtocolTransformsForConfigs(ctx context.Context, configs []*model.Config) error { |
| if len(configs) == 0 { |
| return nil |
| } |
|
|
| channelIDs := make([]any, len(configs)) |
| placeholders := make([]string, len(configs)) |
| idToConfig := make(map[int64]*model.Config, len(configs)) |
| for i, cfg := range configs { |
| channelIDs[i] = cfg.ID |
| placeholders[i] = "?" |
| idToConfig[cfg.ID] = cfg |
| cfg.ProtocolTransforms = nil |
| } |
|
|
| query := fmt.Sprintf( |
| `SELECT channel_id, protocol FROM channel_protocol_transforms WHERE channel_id IN (%s) ORDER BY channel_id, protocol ASC`, |
| strings.Join(placeholders, ","), |
| ) |
| rows, err := s.db.QueryContext(ctx, query, channelIDs...) |
| if err != nil { |
| return fmt.Errorf("query protocol transforms: %w", err) |
| } |
| defer func() { _ = rows.Close() }() |
|
|
| for rows.Next() { |
| var channelID int64 |
| var protocol string |
| if err := rows.Scan(&channelID, &protocol); err != nil { |
| return fmt.Errorf("scan protocol transform: %w", err) |
| } |
| if cfg, ok := idToConfig[channelID]; ok { |
| cfg.ProtocolTransforms = append(cfg.ProtocolTransforms, protocol) |
| } |
| } |
|
|
| if err := rows.Err(); err != nil { |
| return err |
| } |
| for _, cfg := range configs { |
| normalizeLoadedProtocolTransforms(cfg) |
| } |
| return nil |
| } |
|
|
| |
| |
| func (s *SQLStore) loadConfigsAuxConcurrent(ctx context.Context, configs []*model.Config) error { |
| if len(configs) == 0 { |
| return nil |
| } |
| var ( |
| wg sync.WaitGroup |
| modelErr error |
| protocolErr error |
| ) |
| wg.Add(2) |
| go func() { |
| defer wg.Done() |
| modelErr = s.loadModelEntriesForConfigs(ctx, configs) |
| }() |
| go func() { |
| defer wg.Done() |
| protocolErr = s.loadProtocolTransformsForConfigs(ctx, configs) |
| }() |
| wg.Wait() |
| if modelErr != nil { |
| return modelErr |
| } |
| return protocolErr |
| } |
|
|
| func normalizeLoadedProtocolTransforms(cfg *model.Config) { |
| if cfg == nil { |
| return |
| } |
| cfg.ProtocolTransforms = cfg.GetProtocolTransforms() |
| } |
|
|
| func filterConfigsByProtocol(configs []*model.Config, protocol string) []*model.Config { |
| if protocol == "" { |
| return configs |
| } |
| filtered := make([]*model.Config, 0, len(configs)) |
| for _, cfg := range configs { |
| if cfg != nil && cfg.SupportsProtocol(protocol) { |
| filtered = append(filtered, cfg) |
| } |
| } |
| return filtered |
| } |
|
|
| |
| func (s *SQLStore) saveModelEntriesTx(ctx context.Context, tx *sql.Tx, channelID int64, entries []model.ModelEntry) error { |
| return s.saveModelEntriesImpl(ctx, tx, channelID, entries) |
| } |
|
|
| func (s *SQLStore) saveProtocolTransformsTx(ctx context.Context, tx *sql.Tx, channelID int64, transforms []string) error { |
| if _, err := tx.ExecContext(ctx, `DELETE FROM channel_protocol_transforms WHERE channel_id = ?`, channelID); err != nil { |
| return fmt.Errorf("delete old protocol transforms: %w", err) |
| } |
| if len(transforms) == 0 { |
| return nil |
| } |
|
|
| var b strings.Builder |
| b.WriteString(`INSERT INTO channel_protocol_transforms (channel_id, protocol) VALUES `) |
| args := make([]any, 0, len(transforms)*2) |
| for i, protocol := range transforms { |
| if i > 0 { |
| b.WriteByte(',') |
| } |
| b.WriteString("(?, ?)") |
| args = append(args, channelID, protocol) |
| } |
| if _, err := tx.ExecContext(ctx, b.String(), args...); err != nil { |
| return fmt.Errorf("save protocol transforms: %w", err) |
| } |
| return nil |
| } |
|
|
| |
| type dbExecutor interface { |
| ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) |
| } |
|
|
| |
| |
| func (s *SQLStore) saveModelEntriesImpl(ctx context.Context, exec dbExecutor, channelID int64, entries []model.ModelEntry) error { |
| |
| if _, err := exec.ExecContext(ctx, `DELETE FROM channel_models WHERE channel_id = ?`, channelID); err != nil { |
| return fmt.Errorf("delete old model entries: %w", err) |
| } |
|
|
| if len(entries) == 0 { |
| return nil |
| } |
|
|
| |
| |
| const batchSize = 200 |
| baseCreatedAt := time.Now().UnixMilli() |
|
|
| for offset := 0; offset < len(entries); offset += batchSize { |
| end := min(offset+batchSize, len(entries)) |
| chunk := entries[offset:end] |
|
|
| var b strings.Builder |
| b.WriteString(`INSERT INTO channel_models (channel_id, model, redirect_model, created_at) VALUES `) |
| args := make([]any, 0, len(chunk)*4) |
| for i, entry := range chunk { |
| if i > 0 { |
| b.WriteByte(',') |
| } |
| b.WriteString("(?, ?, ?, ?)") |
| args = append(args, channelID, entry.Model, entry.RedirectModel, baseCreatedAt+int64(offset+i)) |
| } |
| if _, err := exec.ExecContext(ctx, b.String(), args...); err != nil { |
| return fmt.Errorf("save model entries (offset %d): %w", offset, err) |
| } |
| } |
|
|
| return nil |
| } |
|
|