this repo has no description
wire through ctx for all database (gorm) operations (#1043)
This involved both adding `.WithContext(ctx)` to any database calls, and
wiring through `ctx` in a couple functions which should have had it
already.
The one place I didn't wire this through is gorm auto database
migrations in the setup command. It felt less necessary during process
initialization, though I can add if requested.
+38
-34
+3
-2
cmd/relay/handlers.go
+3
-2
cmd/relay/handlers.go
···
57
57
return c.JSON(http.StatusBadRequest, xrpc.XRPCError{ErrStr: "HostNotFound", Message: fmt.Sprintf("host server unreachable: %s", err)})
58
58
}
59
59
60
-
return s.relay.SubscribeToHost(hostname, noSSL, false)
60
+
return s.relay.SubscribeToHost(ctx, hostname, noSSL, false)
61
61
}
62
62
63
63
func (s *Service) handleComAtprotoSyncListHosts(c echo.Context, cursor int64, limit int) (*comatproto.SyncListHosts_Output, error) {
···
240
240
}
241
241
242
242
func (svc *Service) HandleHealthCheck(c echo.Context) error {
243
-
if err := svc.relay.Healthcheck(); err != nil {
243
+
ctx := c.Request().Context()
244
+
if err := svc.relay.Healthcheck(ctx); err != nil {
244
245
svc.logger.Error("healthcheck can't connect to database", "err", err)
245
246
return c.JSON(http.StatusInternalServerError, HealthStatus{Status: "error", Message: "can't connect to database"})
246
247
} else {
+2
-1
cmd/relay/main.go
+2
-1
cmd/relay/main.go
···
205
205
}
206
206
207
207
func runRelay(cctx *cli.Context) error {
208
+
ctx := cctx.Context
208
209
logger := configLogger(cctx, os.Stdout)
209
210
210
211
// Trap SIGINT to trigger a shutdown.
···
297
298
}
298
299
299
300
// restart any existing subscriptions as worker goroutines
300
-
if err := r.ResubscribeAllHosts(); err != nil {
301
+
if err := r.ResubscribeAllHosts(ctx); err != nil {
301
302
return err
302
303
}
303
304
+10
-10
cmd/relay/relay/account.go
+10
-10
cmd/relay/relay/account.go
···
25
25
}
26
26
27
27
var acc models.Account
28
-
if err := r.db.Where("did = ?", did).First(&acc).Error; err != nil {
28
+
if err := r.db.WithContext(ctx).Where("did = ?", did).First(&acc).Error; err != nil {
29
29
if errors.Is(err, gorm.ErrRecordNotFound) {
30
30
return nil, ErrAccountNotFound
31
31
}
···
44
44
45
45
func (r *Relay) GetAccountRepo(ctx context.Context, uid uint64) (*models.AccountRepo, error) {
46
46
var repo models.AccountRepo
47
-
if err := r.db.First(&repo, uid).Error; err != nil {
47
+
if err := r.db.WithContext(ctx).First(&repo, uid).Error; err != nil {
48
48
if errors.Is(err, gorm.ErrRecordNotFound) {
49
49
return nil, ErrAccountRepoNotFound
50
50
}
···
102
102
}
103
103
104
104
// create Account row and increment host count in the same transaction
105
-
err = r.db.Transaction(func(tx *gorm.DB) error {
105
+
err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
106
106
if err := tx.Model(&models.Host{}).Where("id = ?", hostID).Update("account_count", gorm.Expr("account_count + 1")).Error; err != nil {
107
107
return fmt.Errorf("failed to increment account count for host (%s): %w", hostname, err)
108
108
}
···
157
157
// TODO: check new upstream status here (using r.HostChecker). In particular, a moved account might go from takendown to active
158
158
159
159
// create Account row and increment host count in the same transaction
160
-
err = r.db.Transaction(func(tx *gorm.DB) error {
160
+
err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
161
161
// decrement old host count
162
162
if err := tx.Model(&models.Host{}).Where("id = ?", acc.HostID).Update("account_count", gorm.Expr("account_count - 1")).Error; err != nil {
163
163
return fmt.Errorf("failed to decrement account count for former host (%d): %w", acc.HostID, err)
···
187
187
// The DID and UID are both required, and *must* match; it is assumed that calling code has already done an account lookup.
188
188
func (r *Relay) UpdateAccountUpstreamStatus(ctx context.Context, did syntax.DID, uid uint64, status models.AccountStatus) error {
189
189
190
-
if err := r.db.Model(models.Account{}).Where("uid = ?", uid).Update("upstream_status", status).Error; err != nil {
190
+
if err := r.db.WithContext(ctx).Model(models.Account{}).Where("uid = ?", uid).Update("upstream_status", status).Error; err != nil {
191
191
return err
192
192
}
193
193
···
206
206
return err
207
207
}
208
208
209
-
if err := r.db.Model(models.Account{}).Where("uid = ?", acc.UID).Update("status", status).Error; err != nil {
209
+
if err := r.db.WithContext(ctx).Model(models.Account{}).Where("uid = ?", acc.UID).Update("status", status).Error; err != nil {
210
210
return err
211
211
}
212
212
···
239
239
func (r *Relay) ListAccounts(ctx context.Context, cursor int64, limit int) ([]*models.Account, error) {
240
240
241
241
accounts := []*models.Account{}
242
-
if err := r.db.Model(&models.Account{}).Where("uid > ? AND status = 'active' AND upstream_status = 'active'", cursor).Order("uid").Limit(limit).Find(&accounts).Error; err != nil {
242
+
if err := r.db.WithContext(ctx).Model(&models.Account{}).Where("uid > ? AND status = 'active' AND upstream_status = 'active'", cursor).Order("uid").Limit(limit).Find(&accounts).Error; err != nil {
243
243
return nil, err
244
244
}
245
245
return accounts, nil
···
248
248
func (r *Relay) ListAccountTakedowns(ctx context.Context, cursor int64, limit int) ([]*models.Account, error) {
249
249
250
250
accounts := []*models.Account{}
251
-
if err := r.db.Model(&models.Account{}).Where("uid > ? AND status = ?", cursor, models.AccountStatusTakendown).Order("uid").Limit(limit).Find(&accounts).Error; err != nil {
251
+
if err := r.db.WithContext(ctx).Model(&models.Account{}).Where("uid > ? AND status = ?", cursor, models.AccountStatusTakendown).Order("uid").Limit(limit).Find(&accounts).Error; err != nil {
252
252
return nil, err
253
253
}
254
254
return accounts, nil
255
255
}
256
256
257
-
func (r *Relay) UpsertAccountRepo(uid uint64, rev syntax.TID, commitCID, commitDataCID string) error {
258
-
return r.db.Exec("INSERT INTO account_repo (uid, rev, commit_cid, commit_data_cid) VALUES (?, ?, ?, ?) ON CONFLICT (uid) DO UPDATE SET rev = EXCLUDED.rev, commit_cid = EXCLUDED.commit_cid, commit_data_cid = EXCLUDED.commit_data_cid", uid, rev, commitCID, commitDataCID).Error
257
+
func (r *Relay) UpsertAccountRepo(ctx context.Context, uid uint64, rev syntax.TID, commitCID, commitDataCID string) error {
258
+
return r.db.WithContext(ctx).Exec("INSERT INTO account_repo (uid, rev, commit_cid, commit_data_cid) VALUES (?, ?, ?, ?) ON CONFLICT (uid) DO UPDATE SET rev = EXCLUDED.rev, commit_cid = EXCLUDED.commit_cid, commit_data_cid = EXCLUDED.commit_data_cid", uid, rev, commitCID, commitDataCID).Error
259
259
}
260
260
261
261
// This implements the `diskpersist.UidSource` interface
+6
-5
cmd/relay/relay/crawl.go
+6
-5
cmd/relay/relay/crawl.go
···
2
2
3
3
import (
4
4
"fmt"
5
+
"context"
5
6
6
7
"github.com/bluesky-social/indigo/cmd/relay/relay/models"
7
8
)
8
9
9
-
func (r *Relay) SubscribeToHost(hostname string, noSSL, adminForce bool) error {
10
+
func (r *Relay) SubscribeToHost(ctx context.Context, hostname string, noSSL, adminForce bool) error {
10
11
11
12
// if we already have an active subscription, exit early
12
13
if r.Slurper.CheckIfSubscribed(hostname) {
···
16
17
// fetch host info from database. this query will not error if host does not yet exist
17
18
newHost := false
18
19
var host models.Host
19
-
if err := r.db.Find(&host, "hostname = ?", hostname).Error; err != nil {
20
+
if err := r.db.WithContext(ctx).Find(&host, "hostname = ?", hostname).Error; err != nil {
20
21
return err
21
22
}
22
23
···
43
44
AccountLimit: accountLimit,
44
45
}
45
46
46
-
if err := r.db.Create(&host).Error; err != nil {
47
+
if err := r.db.WithContext(ctx).Create(&host).Error; err != nil {
47
48
return err
48
49
}
49
50
···
56
57
}
57
58
58
59
// This function expects to be run when starting up, to re-connect to known active hosts
59
-
func (r *Relay) ResubscribeAllHosts() error {
60
+
func (r *Relay) ResubscribeAllHosts(ctx context.Context) error {
60
61
61
62
var all []models.Host
62
-
if err := r.db.Find(&all, "status = ?", "active").Error; err != nil {
63
+
if err := r.db.WithContext(ctx).Find(&all, "status = ?", "active").Error; err != nil {
63
64
return err
64
65
}
65
66
+4
-4
cmd/relay/relay/domain_ban.go
+4
-4
cmd/relay/relay/domain_ban.go
···
45
45
46
46
func (r *Relay) findDomainBan(ctx context.Context, domain string) (bool, error) {
47
47
var ban models.DomainBan
48
-
if err := r.db.Model(&models.DomainBan{}).Where("domain = ?", domain).First(&ban).Error; err != nil {
48
+
if err := r.db.WithContext(ctx).Model(&models.DomainBan{}).Where("domain = ?", domain).First(&ban).Error; err != nil {
49
49
if errors.Is(err, gorm.ErrRecordNotFound) {
50
50
return false, nil
51
51
}
···
56
56
57
57
func (r *Relay) CreateDomainBan(ctx context.Context, domain string) error {
58
58
domainBan := models.DomainBan{Domain: domain}
59
-
return r.db.Create(&domainBan).Error
59
+
return r.db.WithContext(ctx).Create(&domainBan).Error
60
60
}
61
61
62
62
func (r *Relay) RemoveDomainBan(ctx context.Context, domain string) error {
63
-
return r.db.Delete(&models.DomainBan{}, "domain = ?", domain).Error
63
+
return r.db.WithContext(ctx).Delete(&models.DomainBan{}, "domain = ?", domain).Error
64
64
}
65
65
66
66
// returns all domain bans
67
67
func (r *Relay) ListDomainBans(ctx context.Context) ([]models.DomainBan, error) {
68
68
bans := []models.DomainBan{}
69
-
if err := r.db.Model(&models.DomainBan{}).Find(&bans).Error; err != nil {
69
+
if err := r.db.WithContext(ctx).Model(&models.DomainBan{}).Find(&bans).Error; err != nil {
70
70
return nil, err
71
71
}
72
72
return bans, nil
+7
-7
cmd/relay/relay/host.go
+7
-7
cmd/relay/relay/host.go
···
18
18
defer span.End()
19
19
20
20
var host models.Host
21
-
if err := r.db.Model(models.Host{}).Where("hostname = ?", hostname).First(&host).Error; err != nil {
21
+
if err := r.db.WithContext(ctx).Model(models.Host{}).Where("hostname = ?", hostname).First(&host).Error; err != nil {
22
22
if errors.Is(err, gorm.ErrRecordNotFound) {
23
23
return nil, ErrHostNotFound
24
24
}
···
38
38
defer span.End()
39
39
40
40
var host models.Host
41
-
if err := r.db.Find(&host, hostID).Error; err != nil {
41
+
if err := r.db.WithContext(ctx).Find(&host, hostID).Error; err != nil {
42
42
if errors.Is(err, gorm.ErrRecordNotFound) {
43
43
return nil, ErrHostNotFound
44
44
}
···
58
58
// filters for accounts which have seen at least one event
59
59
// TODO: also filter based on status?
60
60
hosts := []*models.Host{}
61
-
if err := r.db.Model(&models.Host{}).Where("id > ? AND last_seq > 0", cursor).Order("id").Limit(limit).Find(&hosts).Error; err != nil {
61
+
if err := r.db.WithContext(ctx).Model(&models.Host{}).Where("id > ? AND last_seq > 0", cursor).Order("id").Limit(limit).Find(&hosts).Error; err != nil {
62
62
return nil, err
63
63
}
64
64
return hosts, nil
65
65
}
66
66
67
67
func (r *Relay) UpdateHostStatus(ctx context.Context, hostID uint64, status models.HostStatus) error {
68
-
return r.db.Model(models.Host{}).Where("id = ?", hostID).Update("status", status).Error
68
+
return r.db.WithContext(ctx).Model(models.Host{}).Where("id = ?", hostID).Update("status", status).Error
69
69
}
70
70
71
71
func (r *Relay) UpdateHostAccountLimit(ctx context.Context, hostID uint64, accountLimit int64) error {
···
82
82
delta := accountLimit - host.AccountLimit
83
83
r.Logger.Info("updating host account limit", "host", host.Hostname, "accountLimit", accountLimit, "previousAccountLimit", host.AccountLimit)
84
84
85
-
if err := r.db.Model(models.Host{}).Where("id = ?", hostID).Update("account_limit", accountLimit).Error; err != nil {
85
+
if err := r.db.WithContext(ctx).Model(models.Host{}).Where("id = ?", hostID).Update("account_limit", accountLimit).Error; err != nil {
86
86
return err
87
87
}
88
88
···
91
91
// if limit increased: potentially mark some "host-throttled" accounts as "active" (ordered by UID ascending)
92
92
// fetch accounts and update individually. this ensures that account cache is cleared for each (as well as any future code around account status changes)
93
93
var accounts []models.Account
94
-
if err := r.db.Model(models.Account{}).Where("status = ? AND upstream_status = ? AND host_id = ?", models.AccountStatusHostThrottled, models.AccountStatusActive, host.ID).Order("uid ASC").Limit(int(delta)).Find(&accounts).Error; err != nil {
94
+
if err := r.db.WithContext(ctx).Model(models.Account{}).Where("status = ? AND upstream_status = ? AND host_id = ?", models.AccountStatusHostThrottled, models.AccountStatusActive, host.ID).Order("uid ASC").Limit(int(delta)).Find(&accounts).Error; err != nil {
95
95
return err
96
96
}
97
97
r.Logger.Info("marking host-throttled accounts as active", "count", len(accounts), "delta", delta, "accountLimit", accountLimit, "host", host.Hostname)
···
118
118
//
119
119
// Note that in some situations this may have partial success.
120
120
func (r *Relay) PersistHostCursors(ctx context.Context, cursors *[]HostCursor) error {
121
-
tx := r.db.WithContext(ctx).Begin()
121
+
tx := r.db.WithContext(ctx).WithContext(ctx).Begin()
122
122
for _, cur := range *cursors {
123
123
if cur.LastSeq <= 0 {
124
124
continue
+2
-2
cmd/relay/relay/ingest.go
+2
-2
cmd/relay/relay/ingest.go
···
137
137
return err
138
138
}
139
139
140
-
err = r.UpsertAccountRepo(acc.UID, syntax.TID(newRepo.Rev), newRepo.CommitCID, newRepo.CommitDataCID)
140
+
err = r.UpsertAccountRepo(ctx, acc.UID, syntax.TID(newRepo.Rev), newRepo.CommitCID, newRepo.CommitDataCID)
141
141
if err != nil {
142
142
return fmt.Errorf("failed to upsert account repo (%s): %w", acc.DID, err)
143
143
}
···
182
182
return err
183
183
}
184
184
185
-
err = r.UpsertAccountRepo(acc.UID, syntax.TID(newRepo.Rev), newRepo.CommitCID, newRepo.CommitDataCID)
185
+
err = r.UpsertAccountRepo(ctx, acc.UID, syntax.TID(newRepo.Rev), newRepo.CommitCID, newRepo.CommitDataCID)
186
186
if err != nil {
187
187
return fmt.Errorf("failed to upsert account repo (%s): %w", acc.DID, err)
188
188
}
+3
-2
cmd/relay/relay/relay.go
+3
-2
cmd/relay/relay/relay.go
···
3
3
import (
4
4
"log/slog"
5
5
"sync"
6
+
"context"
6
7
7
8
"github.com/bluesky-social/indigo/atproto/identity"
8
9
"github.com/bluesky-social/indigo/cmd/relay/relay/models"
···
125
126
}
126
127
127
128
// simple check of connection to database
128
-
func (r *Relay) Healthcheck() error {
129
-
return r.db.Exec("SELECT 1").Error
129
+
func (r *Relay) Healthcheck(ctx context.Context) error {
130
+
return r.db.WithContext(ctx).Exec("SELECT 1").Error
130
131
}
+1
-1
cmd/relay/testing/runner.go
+1
-1
cmd/relay/testing/runner.go
···
132
132
133
133
sr := MustSimpleRelay(&dir, tmpd, s.Lenient)
134
134
135
-
err = sr.Relay.SubscribeToHost(fmt.Sprintf("localhost:%d", hostPort), true, true)
135
+
err = sr.Relay.SubscribeToHost(ctx, fmt.Sprintf("localhost:%d", hostPort), true, true)
136
136
if err != nil {
137
137
return err
138
138
}