Changeset 0.22.0.1 (#94)

server/store/postgres/cap_override.go

@@ -0,0 +1,85 @@
+package postgres
+
+import (
+	"context"
+	"database/sql"
+
+	"git.gobha.me/xcaliber/chat-switchboard/models"
+)
+
+// CapOverrideStore implements store.CapabilityOverrideStore for Postgres.
+type CapOverrideStore struct {
+	db *sql.DB
+}
+
+func NewCapOverrideStore(db *sql.DB) *CapOverrideStore {
+	return &CapOverrideStore{db: db}
+}
+
+func (s *CapOverrideStore) Set(ctx context.Context, o *models.CapabilityOverride) error {
+	_, err := s.db.ExecContext(ctx, `
+		INSERT INTO capability_overrides (provider_config_id, model_id, field, value, set_by)
+		VALUES ($1, $2, $3, $4, $5)
+		ON CONFLICT (provider_config_id, model_id, field) DO UPDATE SET
+			value = EXCLUDED.value,
+			set_by = EXCLUDED.set_by,
+			created_at = now()
+	`, o.ProviderConfigID, o.ModelID, o.Field, o.Value, o.SetBy)
+	return err
+}
+
+func (s *CapOverrideStore) Delete(ctx context.Context, id string) error {
+	_, err := s.db.ExecContext(ctx, `DELETE FROM capability_overrides WHERE id = $1`, id)
+	return err
+}
+
+func (s *CapOverrideStore) ListForModel(ctx context.Context, modelID string) ([]models.CapabilityOverride, error) {
+	return s.query(ctx, `
+		SELECT id, provider_config_id, model_id, field, value, set_by, created_at
+		FROM capability_overrides
+		WHERE model_id = $1
+		ORDER BY provider_config_id NULLS LAST, field
+	`, modelID)
+}
+
+func (s *CapOverrideStore) ListForProviderModel(ctx context.Context, providerConfigID, modelID string) ([]models.CapabilityOverride, error) {
+	return s.query(ctx, `
+		SELECT id, provider_config_id, model_id, field, value, set_by, created_at
+		FROM capability_overrides
+		WHERE model_id = $1 AND (provider_config_id = $2 OR provider_config_id IS NULL)
+		ORDER BY provider_config_id NULLS LAST, field
+	`, modelID, providerConfigID)
+}
+
+func (s *CapOverrideStore) ListAll(ctx context.Context) ([]models.CapabilityOverride, error) {
+	return s.query(ctx, `
+		SELECT id, provider_config_id, model_id, field, value, set_by, created_at
+		FROM capability_overrides
+		ORDER BY model_id, provider_config_id NULLS LAST, field
+	`)
+}
+
+func (s *CapOverrideStore) DeleteForProvider(ctx context.Context, providerConfigID string) error {
+	_, err := s.db.ExecContext(ctx, `
+		DELETE FROM capability_overrides WHERE provider_config_id = $1
+	`, providerConfigID)
+	return err
+}
+
+func (s *CapOverrideStore) query(ctx context.Context, q string, args ...interface{}) ([]models.CapabilityOverride, error) {
+	rows, err := s.db.QueryContext(ctx, q, args...)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var result []models.CapabilityOverride
+	for rows.Next() {
+		var o models.CapabilityOverride
+		if err := rows.Scan(&o.ID, &o.ProviderConfigID, &o.ModelID, &o.Field, &o.Value, &o.SetBy, &o.CreatedAt); err != nil {
+			return nil, err
+		}
+		result = append(result, o)
+	}
+	return result, rows.Err()
+}