Skip to content

Commit 7a6309f

Browse files
henrybarretohenrybarreto
committed
refactor(api,pkg): save events summary into a session associating with the seat
1 parent f1c11dc commit 7a6309f

File tree

5 files changed

+366
-5
lines changed

5 files changed

+366
-5
lines changed

‎api/store/mongo/migrations/main.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ func GenerateMigrations() []migrate.Migration {
117117
migration105,
118118
migration106,
119119
migration107,
120+
migration108,
120121
}
121122
}
122123

‎api/store/mongo/migrations/migration_108.go

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
package migrations
2+
3+
import (
4+
"context"
5+
6+
"github.com/shellhub-io/shellhub/pkg/models"
7+
log "github.com/sirupsen/logrus"
8+
migrate "github.com/xakep666/mongo-migrate"
9+
"go.mongodb.org/mongo-driver/bson"
10+
"go.mongodb.org/mongo-driver/mongo"
11+
)
12+
13+
var migration108 = migrate.Migration{
14+
Version: 108,
15+
Description: "Migrate session events to session seats structure",
16+
Up: migrate.MigrationFunc(func(ctx context.Context, db *mongo.Database) error {
17+
log.WithFields(log.Fields{"component": "migration", "version": 108, "action": "Up"}).Info("Applying migration")
18+
19+
cursor, err := db.Collection("sessions").Find(ctx, bson.M{
20+
"seats": nil,
21+
"events": bson.M{"$exists": true},
22+
})
23+
if err != nil {
24+
log.WithError(err).Error("Failed to fetch sessions")
25+
26+
return err
27+
}
28+
29+
defer cursor.Close(ctx)
30+
31+
for cursor.Next(ctx) {
32+
var session struct {
33+
UID string `bson:"uid"`
34+
Seats []models.SessionSeat `bson:"seats"`
35+
}
36+
37+
if err := cursor.Decode(&session); err != nil {
38+
log.WithError(err).Error("Failed to decode session")
39+
40+
return err
41+
}
42+
43+
eventsCursor, err := db.Collection("sessions_events").Find(ctx, bson.M{"session": session.UID})
44+
if err != nil {
45+
log.WithError(err).WithField("session", session.UID).Error("Failed to fetch events for session")
46+
47+
return err
48+
}
49+
50+
eventsBySeat := make(map[int][]string)
51+
52+
for eventsCursor.Next(ctx) {
53+
var event struct {
54+
Type string `bson:"type"`
55+
Seat int `bson:"seat"`
56+
}
57+
58+
if err := eventsCursor.Decode(&event); err != nil {
59+
log.WithError(err).Error("Failed to decode event")
60+
61+
return err
62+
}
63+
64+
if _, ok := eventsBySeat[event.Seat]; !ok {
65+
eventsBySeat[event.Seat] = []string{}
66+
}
67+
68+
eventsBySeat[event.Seat] = append(eventsBySeat[event.Seat], event.Type)
69+
}
70+
71+
eventsCursor.Close(ctx)
72+
73+
var seats []models.SessionSeat
74+
var seatIDs []int
75+
76+
eventTypes := make(map[string]bool)
77+
78+
for seatID, events := range eventsBySeat {
79+
seats = append(seats, models.SessionSeat{
80+
ID: seatID,
81+
Events: events,
82+
})
83+
84+
seatIDs = append(seatIDs, seatID)
85+
86+
for _, eventType := range events {
87+
eventTypes[eventType] = true
88+
}
89+
}
90+
91+
var types []string
92+
for eventType := range eventTypes {
93+
types = append(types, eventType)
94+
}
95+
96+
_, err = db.Collection("sessions").UpdateOne(ctx,
97+
bson.M{"uid": session.UID},
98+
bson.M{
99+
"$set": bson.M{
100+
"seats": seats,
101+
"events": models.SessionEvents{
102+
Types: types,
103+
Seats: seatIDs,
104+
},
105+
},
106+
})
107+
if err != nil {
108+
log.WithError(err).WithField("session", session.UID).Error("Failed to update session")
109+
110+
return err
111+
}
112+
}
113+
114+
log.WithFields(log.Fields{"component": "migration", "version": 108, "action": "Up"}).Info("Migration completed successfully")
115+
116+
return nil
117+
}),
118+
Down: migrate.MigrationFunc(func(ctx context.Context, db *mongo.Database) error {
119+
log.WithFields(log.Fields{"component": "migration", "version": 108, "action": "Down"}).Info("Reverting migration")
120+
121+
if _, err := db.Collection("sessions").UpdateMany(
122+
ctx,
123+
bson.M{},
124+
bson.M{
125+
"$unset": bson.M{
126+
"seats": "",
127+
},
128+
},
129+
); err != nil {
130+
log.WithError(err).Error("Failed to revert events migration")
131+
132+
return err
133+
}
134+
135+
log.WithFields(log.Fields{"component": "migration", "version": 108, "action": "Down"}).Info("Migration reverted successfully")
136+
137+
return nil
138+
}),
139+
}

‎api/store/mongo/migrations/migration_108_test.go

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
package migrations
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/shellhub-io/shellhub/pkg/models"
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
migrate "github.com/xakep666/mongo-migrate"
12+
"go.mongodb.org/mongo-driver/bson"
13+
)
14+
15+
func TestMigration108Up(t *testing.T) {
16+
ctx := context.Background()
17+
18+
cases := []struct {
19+
description string
20+
setup func() error
21+
verify func(t *testing.T)
22+
}{
23+
{
24+
description: "succeeds migrating session events to seats",
25+
setup: func() error {
26+
sessions := []bson.M{
27+
{
28+
"uid": "session-1",
29+
"events": bson.M{
30+
"types": []string{"pty-output", "window-change"},
31+
"seats": []int{0, 1},
32+
},
33+
},
34+
{
35+
"uid": "session-2",
36+
"events": bson.M{
37+
"types": []string{"pty-output"},
38+
"seats": []int{0},
39+
},
40+
},
41+
}
42+
43+
events := []bson.M{
44+
{
45+
"session": "session-1",
46+
"type": "pty-output",
47+
"seat": 0,
48+
"timestamp": time.Now(),
49+
},
50+
{
51+
"session": "session-1",
52+
"type": "window-change",
53+
"seat": 0,
54+
"timestamp": time.Now(),
55+
},
56+
{
57+
"session": "session-1",
58+
"type": "pty-output",
59+
"seat": 1,
60+
"timestamp": time.Now(),
61+
},
62+
{
63+
"session": "session-2",
64+
"type": "pty-output",
65+
"seat": 0,
66+
"timestamp": time.Now(),
67+
},
68+
}
69+
70+
if _, err := c.Database("test").Collection("sessions").InsertMany(ctx, []any{sessions[0], sessions[1]}); err != nil {
71+
return err
72+
}
73+
74+
if _, err := c.Database("test").Collection("sessions_events").InsertMany(ctx, []any{events[0], events[1], events[2], events[3]}); err != nil {
75+
return err
76+
}
77+
78+
return nil
79+
},
80+
verify: func(t *testing.T) {
81+
var session1 struct {
82+
Events models.SessionEvents `bson:"events"`
83+
Seats []models.SessionSeat `bson:"seats"`
84+
}
85+
86+
err := c.Database("test").Collection("sessions").FindOne(ctx, bson.M{"uid": "session-1"}).Decode(&session1)
87+
assert.NoError(t, err)
88+
89+
assert.ElementsMatch(t, []string{"pty-output", "window-change"}, session1.Events.Types)
90+
assert.ElementsMatch(t, []int{0, 1}, session1.Events.Seats)
91+
92+
assert.Len(t, session1.Seats, 2)
93+
for _, seat := range session1.Seats {
94+
switch seat.ID {
95+
case 0:
96+
assert.ElementsMatch(t, []string{"pty-output", "window-change"}, seat.Events)
97+
case 1:
98+
assert.ElementsMatch(t, []string{"pty-output"}, seat.Events)
99+
}
100+
}
101+
102+
var session2 struct {
103+
Events models.SessionEvents `bson:"events"`
104+
Seats []models.SessionSeat `bson:"seats"`
105+
}
106+
107+
err = c.Database("test").Collection("sessions").FindOne(ctx, bson.M{"uid": "session-2"}).Decode(&session2)
108+
assert.NoError(t, err)
109+
110+
assert.ElementsMatch(t, []string{"pty-output"}, session2.Events.Types)
111+
assert.ElementsMatch(t, []int{0}, session2.Events.Seats)
112+
113+
assert.Len(t, session2.Seats, 1)
114+
for _, seat := range session2.Seats {
115+
switch seat.ID {
116+
case 0:
117+
assert.ElementsMatch(t, []string{"pty-output"}, seat.Events)
118+
}
119+
}
120+
},
121+
},
122+
}
123+
124+
for _, tc := range cases {
125+
t.Run(tc.description, func(t *testing.T) {
126+
t.Cleanup(func() { assert.NoError(t, srv.Reset()) })
127+
128+
require.NoError(t, tc.setup())
129+
migrates := migrate.NewMigrate(c.Database("test"), GenerateMigrations()[107])
130+
require.NoError(t, migrates.Up(ctx, migrate.AllAvailable))
131+
132+
tc.verify(t)
133+
})
134+
}
135+
}
136+
137+
func TestMigration108Down(t *testing.T) {
138+
ctx := context.Background()
139+
140+
cases := []struct {
141+
description string
142+
setup func() error
143+
verify func(t *testing.T)
144+
}{
145+
{
146+
description: "succeeds removing events field while keeping seats structure",
147+
setup: func() error {
148+
sessions := []bson.M{
149+
{
150+
"uid": "session-1",
151+
"events": bson.M{
152+
"types": []string{"pty-output", "window-change"},
153+
"seats": []int{0, 1},
154+
},
155+
"seats": []bson.M{
156+
{
157+
"id": 0,
158+
"events": []string{"pty-output", "window-change"},
159+
},
160+
{
161+
"id": 1,
162+
"events": []string{"pty-output"},
163+
},
164+
},
165+
},
166+
}
167+
168+
if _, err := c.Database("test").Collection("sessions").InsertMany(ctx, []any{sessions[0]}); err != nil {
169+
return err
170+
}
171+
172+
return nil
173+
},
174+
verify: func(t *testing.T) {
175+
var session struct {
176+
UID string `bson:"uid"`
177+
Events models.SessionEvents `bson:"events"`
178+
Seats []models.SessionSeat `bson:"seats"`
179+
}
180+
181+
err := c.Database("test").Collection("sessions").FindOne(ctx, bson.M{"uid": "session-1"}).Decode(&session)
182+
assert.NoError(t, err)
183+
184+
assert.ElementsMatch(t, []string{"pty-output", "window-change"}, session.Events.Types)
185+
assert.ElementsMatch(t, []int{0, 1}, session.Events.Seats)
186+
187+
assert.Len(t, session.Seats, 0, "Seats should be removed")
188+
},
189+
},
190+
}
191+
192+
for _, tc := range cases {
193+
t.Run(tc.description, func(t *testing.T) {
194+
t.Cleanup(func() { assert.NoError(t, srv.Reset()) })
195+
196+
require.NoError(t, tc.setup())
197+
migrates := migrate.NewMigrate(c.Database("test"), GenerateMigrations()[107])
198+
require.NoError(t, migrates.Up(ctx, migrate.AllAvailable))
199+
require.NoError(t, migrates.Down(ctx, migrate.AllAvailable))
200+
201+
tc.verify(t)
202+
})
203+
}
204+
}

‎api/store/mongo/session.go

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ func (s *Store) SessionCreate(ctx context.Context, session models.Session) (*mod
217217
session.StartedAt = clock.Now()
218218
session.LastSeen = session.StartedAt
219219
session.Recorded = false
220+
session.Seats = []models.SessionSeat{}
220221

221222
device, err := s.DeviceResolve(ctx, store.DeviceUIDResolver, string(session.DeviceUID))
222223
if err != nil {
@@ -329,15 +330,28 @@ func (s *Store) SessionEvent(ctx context.Context, uid models.UID, event *models.
329330

330331
if _, err := session.WithTransaction(ctx, func(ctx mongo.SessionContext) (any, error) {
331332
if _, err := s.db.Collection("sessions").UpdateOne(ctx,
332-
bson.M{"uid": uid},
333+
bson.M{"uid": uid, "seats.id": bson.M{"$ne": event.Seat}},
334+
bson.M{
335+
"$push": bson.M{
336+
"seats": bson.M{
337+
"id": event.Seat,
338+
"events": bson.A{},
339+
},
340+
},
341+
},
342+
); err != nil {
343+
return nil, FromMongoError(err)
344+
}
345+
346+
if _, err := s.db.Collection("sessions").UpdateOne(ctx,
347+
bson.M{"uid": uid, "seats.id": event.Seat},
333348
bson.M{
334349
"$addToSet": bson.M{
335-
"events.types": event.Type,
336-
"events.seats": event.Seat,
350+
"seats.$.events": event.Type,
337351
},
338352
},
339353
); err != nil {
340-
return nil, err
354+
return nil, FromMongoError(err)
341355
}
342356

343357
if _, err := s.db.Collection("sessions_events").InsertOne(ctx, event); err != nil {

0 commit comments

Comments
 (0)