Skip to content

Commit 02da338

Browse files
lucasmrodgetvictor
lucasmrod
authored and
getvictor
committed
Additional changes to happy path and cleanup cron job (#17757)
#17441 & #17442
1 parent 90c8160 commit 02da338

12 files changed

+553
-117
lines changed

‎cmd/fleet/calendar_cron.go

+214-77
Large diffs are not rendered by default.

‎cmd/fleet/calendar_cron_test.go

+42-15
Original file line numberDiff line numberDiff line change
@@ -12,34 +12,61 @@ func TestGetPreferredCalendarEventDate(t *testing.T) {
1212
return time.Date(year, month, day, 0, 0, 0, 0, time.UTC)
1313
}
1414
for _, tc := range []struct {
15-
name string
16-
year int
17-
month time.Month
18-
days int
15+
name string
16+
year int
17+
month time.Month
18+
daysStart int
19+
daysEnd int
20+
webhookFiredThisMonth bool
1921

2022
expected time.Time
2123
}{
2224
{
23-
year: 2024,
24-
month: 3,
25-
days: 31,
26-
name: "March 2024",
25+
name: "March 2024 (webhook hasn't fired)",
26+
year: 2024,
27+
month: 3,
28+
daysStart: 1,
29+
daysEnd: 31,
30+
webhookFiredThisMonth: false,
31+
32+
expected: date(2024, 3, 19),
33+
},
34+
{
35+
name: "March 2024 (webhook has fired, days before 3rd Tuesday)",
36+
year: 2024,
37+
month: 3,
38+
daysStart: 1,
39+
daysEnd: 18,
40+
webhookFiredThisMonth: true,
41+
2742
expected: date(2024, 3, 19),
2843
},
2944
{
30-
year: 2024,
31-
month: 4,
32-
days: 30,
33-
name: "April 2024",
45+
name: "March 2024 (webhook has fired, days after 3rd Tuesday)",
46+
year: 2024,
47+
month: 3,
48+
daysStart: 20,
49+
daysEnd: 30,
50+
webhookFiredThisMonth: true,
51+
52+
expected: date(2024, 4, 16),
53+
},
54+
{
55+
name: "April 2024 (webhook hasn't fired)",
56+
year: 2024,
57+
month: 4,
58+
daysEnd: 30,
59+
webhookFiredThisMonth: false,
60+
3461
expected: date(2024, 4, 16),
3562
},
3663
} {
3764
t.Run(tc.name, func(t *testing.T) {
38-
for day := 1; day <= tc.days; day++ {
39-
actual := getPreferredCalendarEventDate(tc.year, tc.month, day)
65+
for day := tc.daysStart; day <= tc.daysEnd; day++ {
66+
actual := getPreferredCalendarEventDate(tc.year, tc.month, day, tc.webhookFiredThisMonth)
4067
require.NotEqual(t, actual.Weekday(), time.Saturday)
4168
require.NotEqual(t, actual.Weekday(), time.Sunday)
42-
if day <= tc.expected.Day() {
69+
if day <= tc.expected.Day() || tc.webhookFiredThisMonth {
4370
require.Equal(t, tc.expected, actual)
4471
} else {
4572
today := date(tc.year, tc.month, day)

‎server/datastore/mysql/calendar_events.go

+78-14
Original file line numberDiff line numberDiff line change
@@ -11,23 +11,29 @@ import (
1111
"github.com/jmoiron/sqlx"
1212
)
1313

14-
func (ds *Datastore) NewCalendarEvent(
14+
func (ds *Datastore) CreateOrUpdateCalendarEvent(
1515
ctx context.Context,
1616
email string,
1717
startTime time.Time,
1818
endTime time.Time,
1919
data []byte,
2020
hostID uint,
21+
webhookStatus fleet.CalendarWebhookStatus,
2122
) (*fleet.CalendarEvent, error) {
22-
var calendarEvent *fleet.CalendarEvent
23+
var id int64
2324
if err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
2425
const calendarEventsQuery = `
2526
INSERT INTO calendar_events (
2627
email,
2728
start_time,
2829
end_time,
2930
event
30-
) VALUES (?, ?, ?, ?);
31+
) VALUES (?, ?, ?, ?)
32+
ON DUPLICATE KEY UPDATE
33+
start_time = VALUES(start_time),
34+
end_time = VALUES(end_time),
35+
event = VALUES(event),
36+
updated_at = CURRENT_TIMESTAMP;
3137
`
3238
result, err := tx.ExecContext(
3339
ctx,
@@ -41,28 +47,31 @@ func (ds *Datastore) NewCalendarEvent(
4147
return ctxerr.Wrap(ctx, err, "insert calendar event")
4248
}
4349

44-
id, _ := result.LastInsertId()
45-
calendarEvent = &fleet.CalendarEvent{
46-
ID: uint(id),
47-
Email: email,
48-
StartTime: startTime,
49-
EndTime: endTime,
50-
Data: data,
50+
if insertOnDuplicateDidInsert(result) {
51+
id, _ = result.LastInsertId()
52+
} else {
53+
stmt := `SELECT id FROM calendar_events WHERE email = ?`
54+
if err := sqlx.GetContext(ctx, tx, &id, stmt, email); err != nil {
55+
return ctxerr.Wrap(ctx, err, "query mdm solution id")
56+
}
5157
}
5258

5359
const hostCalendarEventsQuery = `
5460
INSERT INTO host_calendar_events (
5561
host_id,
5662
calendar_event_id,
5763
webhook_status
58-
) VALUES (?, ?, ?);
64+
) VALUES (?, ?, ?)
65+
ON DUPLICATE KEY UPDATE
66+
webhook_status = VALUES(webhook_status),
67+
calendar_event_id = VALUES(calendar_event_id);
5968
`
6069
result, err = tx.ExecContext(
6170
ctx,
6271
hostCalendarEventsQuery,
6372
hostID,
64-
calendarEvent.ID,
65-
fleet.CalendarWebhookStatusPending,
73+
id,
74+
webhookStatus,
6675
)
6776
if err != nil {
6877
return ctxerr.Wrap(ctx, err, "insert host calendar event")
@@ -71,9 +80,29 @@ func (ds *Datastore) NewCalendarEvent(
7180
}); err != nil {
7281
return nil, ctxerr.Wrap(ctx, err)
7382
}
83+
84+
calendarEvent, err := getCalendarEventByID(ctx, ds.writer(ctx), uint(id))
85+
if err != nil {
86+
return nil, ctxerr.Wrap(ctx, err, "get created calendar event by id")
87+
}
7488
return calendarEvent, nil
7589
}
7690

91+
func getCalendarEventByID(ctx context.Context, q sqlx.QueryerContext, id uint) (*fleet.CalendarEvent, error) {
92+
const calendarEventsQuery = `
93+
SELECT * FROM calendar_events WHERE id = ?;
94+
`
95+
var calendarEvent fleet.CalendarEvent
96+
err := sqlx.GetContext(ctx, q, &calendarEvent, calendarEventsQuery, id)
97+
if err != nil {
98+
if err == sql.ErrNoRows {
99+
return nil, ctxerr.Wrap(ctx, notFound("CalendarEvent").WithID(id))
100+
}
101+
return nil, ctxerr.Wrap(ctx, err, "get calendar event")
102+
}
103+
return &calendarEvent, nil
104+
}
105+
77106
func (ds *Datastore) GetCalendarEvent(ctx context.Context, email string) (*fleet.CalendarEvent, error) {
78107
const calendarEventsQuery = `
79108
SELECT * FROM calendar_events WHERE email = ?;
@@ -94,7 +123,8 @@ func (ds *Datastore) UpdateCalendarEvent(ctx context.Context, calendarEventID ui
94123
UPDATE calendar_events SET
95124
start_time = ?,
96125
end_time = ?,
97-
event = ?
126+
event = ?,
127+
updated_at = CURRENT_TIMESTAMP
98128
WHERE id = ?;
99129
`
100130
if _, err := ds.writer(ctx).ExecContext(ctx, calendarEventsQuery, startTime, endTime, data, calendarEventID); err != nil {
@@ -148,3 +178,37 @@ func (ds *Datastore) UpdateHostCalendarWebhookStatus(ctx context.Context, hostID
148178
}
149179
return nil
150180
}
181+
182+
func (ds *Datastore) ListCalendarEvents(ctx context.Context, teamID *uint) ([]*fleet.CalendarEvent, error) {
183+
calendarEventsQuery := `
184+
SELECT ce.* FROM calendar_events ce
185+
`
186+
187+
var args []interface{}
188+
if teamID != nil {
189+
// TODO(lucas): Should we add a team_id column to calendar_events?
190+
calendarEventsQuery += ` JOIN host_calendar_events hce ON ce.id=hce.calendar_event_id
191+
JOIN hosts h ON h.id=hce.host_id WHERE h.team_id = ?`
192+
args = append(args, *teamID)
193+
}
194+
195+
var calendarEvents []*fleet.CalendarEvent
196+
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &calendarEvents, calendarEventsQuery, args...); err != nil {
197+
if err == sql.ErrNoRows {
198+
return nil, nil
199+
}
200+
return nil, ctxerr.Wrap(ctx, err, "get all calendar events")
201+
}
202+
return calendarEvents, nil
203+
}
204+
205+
func (ds *Datastore) ListOutOfDateCalendarEvents(ctx context.Context, t time.Time) ([]*fleet.CalendarEvent, error) {
206+
calendarEventsQuery := `
207+
SELECT ce.* FROM calendar_events ce WHERE updated_at < ?
208+
`
209+
var calendarEvents []*fleet.CalendarEvent
210+
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &calendarEvents, calendarEventsQuery, t); err != nil {
211+
return nil, ctxerr.Wrap(ctx, err, "get all calendar events")
212+
}
213+
return calendarEvents, nil
214+
}

‎server/datastore/mysql/calendar_events_test.go

+123-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,128 @@
11
package mysql
22

3-
import "testing"
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/fleetdm/fleet/v4/server/fleet"
9+
"github.com/fleetdm/fleet/v4/server/ptr"
10+
"github.com/stretchr/testify/require"
11+
)
412

513
func TestCalendarEvents(t *testing.T) {
14+
ds := CreateMySQLDS(t)
15+
16+
cases := []struct {
17+
name string
18+
fn func(t *testing.T, ds *Datastore)
19+
}{
20+
{"UpdateCalendarEvent", testUpdateCalendarEvent},
21+
{"CreateOrUpdateCalendarEvent", testCreateOrUpdateCalendarEvent},
22+
}
23+
for _, c := range cases {
24+
t.Run(c.name, func(t *testing.T) {
25+
defer TruncateTables(t, ds)
26+
27+
c.fn(t, ds)
28+
})
29+
}
30+
}
31+
32+
func testUpdateCalendarEvent(t *testing.T, ds *Datastore) {
33+
ctx := context.Background()
34+
35+
host, err := ds.NewHost(context.Background(), &fleet.Host{
36+
DetailUpdatedAt: time.Now(),
37+
LabelUpdatedAt: time.Now(),
38+
PolicyUpdatedAt: time.Now(),
39+
SeenTime: time.Now(),
40+
NodeKey: ptr.String("1"),
41+
UUID: "1",
42+
Hostname: "foo.local",
43+
PrimaryIP: "192.168.1.1",
44+
PrimaryMac: "30-65-EC-6F-C4-58",
45+
})
46+
require.NoError(t, err)
47+
err = ds.ReplaceHostDeviceMapping(ctx, host.ID, []*fleet.HostDeviceMapping{
48+
{
49+
HostID: host.ID,
50+
Email: "[email protected]",
51+
Source: "google_chrome_profiles",
52+
},
53+
}, "google_chrome_profiles")
54+
require.NoError(t, err)
55+
56+
startTime1 := time.Now()
57+
endTime1 := startTime1.Add(30 * time.Minute)
58+
calendarEvent, err := ds.CreateOrUpdateCalendarEvent(ctx, "[email protected]", startTime1, endTime1, []byte(`{}`), host.ID, fleet.CalendarWebhookStatusNone)
59+
require.NoError(t, err)
60+
61+
time.Sleep(1 * time.Second)
62+
63+
err = ds.UpdateCalendarEvent(ctx, calendarEvent.ID, startTime1, endTime1, []byte(`{}`))
64+
require.NoError(t, err)
65+
66+
calendarEvent2, err := ds.GetCalendarEvent(ctx, "[email protected]")
67+
require.NoError(t, err)
68+
require.NotEqual(t, *calendarEvent, *calendarEvent2)
69+
calendarEvent.UpdatedAt = calendarEvent2.UpdatedAt
70+
require.Equal(t, *calendarEvent, *calendarEvent2)
71+
72+
// TODO(lucas): Add more tests here.
73+
}
74+
75+
func testCreateOrUpdateCalendarEvent(t *testing.T, ds *Datastore) {
76+
ctx := context.Background()
77+
78+
host, err := ds.NewHost(context.Background(), &fleet.Host{
79+
DetailUpdatedAt: time.Now(),
80+
LabelUpdatedAt: time.Now(),
81+
PolicyUpdatedAt: time.Now(),
82+
SeenTime: time.Now(),
83+
NodeKey: ptr.String("1"),
84+
UUID: "1",
85+
Hostname: "foo.local",
86+
PrimaryIP: "192.168.1.1",
87+
PrimaryMac: "30-65-EC-6F-C4-58",
88+
})
89+
require.NoError(t, err)
90+
err = ds.ReplaceHostDeviceMapping(ctx, host.ID, []*fleet.HostDeviceMapping{
91+
{
92+
HostID: host.ID,
93+
Email: "[email protected]",
94+
Source: "google_chrome_profiles",
95+
},
96+
}, "google_chrome_profiles")
97+
require.NoError(t, err)
98+
99+
startTime1 := time.Now()
100+
endTime1 := startTime1.Add(30 * time.Minute)
101+
calendarEvent, err := ds.CreateOrUpdateCalendarEvent(ctx, "[email protected]", startTime1, endTime1, []byte(`{}`), host.ID, fleet.CalendarWebhookStatusNone)
102+
require.NoError(t, err)
103+
104+
time.Sleep(1 * time.Second)
105+
106+
calendarEvent2, err := ds.CreateOrUpdateCalendarEvent(ctx, "[email protected]", startTime1, endTime1, []byte(`{}`), host.ID, fleet.CalendarWebhookStatusNone)
107+
require.NoError(t, err)
108+
require.Greater(t, calendarEvent2.UpdatedAt, calendarEvent.UpdatedAt)
109+
calendarEvent.UpdatedAt = calendarEvent2.UpdatedAt
110+
require.Equal(t, *calendarEvent, *calendarEvent2)
111+
112+
time.Sleep(1 * time.Second)
113+
114+
startTime2 := startTime1.Add(1 * time.Hour)
115+
endTime2 := startTime1.Add(30 * time.Minute)
116+
calendarEvent3, err := ds.CreateOrUpdateCalendarEvent(ctx, "[email protected]", startTime2, endTime2, []byte(`{"foo": "bar"}`), host.ID, fleet.CalendarWebhookStatusPending)
117+
require.NoError(t, err)
118+
require.Greater(t, calendarEvent3.UpdatedAt, calendarEvent2.UpdatedAt)
119+
require.WithinDuration(t, startTime2, calendarEvent3.StartTime, 1*time.Second)
120+
require.WithinDuration(t, endTime2, calendarEvent3.EndTime, 1*time.Second)
121+
require.Equal(t, string(calendarEvent3.Data), `{"foo": "bar"}`)
122+
123+
calendarEvent3b, err := ds.GetCalendarEvent(ctx, "[email protected]")
124+
require.NoError(t, err)
125+
require.Equal(t, calendarEvent3, calendarEvent3b)
126+
127+
// TODO(lucas): Add more tests here.
6128
}

‎server/datastore/mysql/migrations/tables/20240314085226_AddCalendarEventTables.go

+3-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ func Up_20240314085226(tx *sql.Tx) error {
2121
event JSON NOT NULL,
2222
2323
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
24-
updated_at TIMESTAMP NOT NULL NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
24+
updated_at TIMESTAMP NOT NULL NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
25+
26+
UNIQUE KEY idx_one_calendar_event_per_email (email)
2527
);
2628
`); err != nil {
2729
return fmt.Errorf("create calendar_events table: %w", err)

‎server/datastore/mysql/policies.go

+1
Original file line numberDiff line numberDiff line change
@@ -1172,6 +1172,7 @@ func (ds *Datastore) GetCalendarPolicies(ctx context.Context, teamID uint) ([]fl
11721172
}
11731173

11741174
// TODO(lucas): Must be tested at scale.
1175+
// TODO(lucas): Filter out hosts with team_id == NULL
11751176
func (ds *Datastore) GetHostsPolicyMemberships(ctx context.Context, domain string, policyIDs []uint) ([]fleet.HostPolicyMembershipData, error) {
11761177
query := `
11771178
SELECT

‎server/datastore/mysql/schema.sql

+2-1
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@ CREATE TABLE `calendar_events` (
5252
`event` json NOT NULL,
5353
`created_at` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
5454
`updated_at` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
55-
PRIMARY KEY (`id`)
55+
PRIMARY KEY (`id`),
56+
UNIQUE KEY `idx_one_calendar_event_per_email` (`email`)
5657
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
5758
/*!40101 SET character_set_client = @saved_cs_client */;
5859
/*!40101 SET @saved_cs_client = @@character_set_client */;

‎server/fleet/calendar_events.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ type CalendarEvent struct {
1515
type CalendarWebhookStatus int
1616

1717
const (
18-
CalendarWebhookStatusPending CalendarWebhookStatus = iota
18+
CalendarWebhookStatusNone CalendarWebhookStatus = iota
19+
CalendarWebhookStatusPending
1920
CalendarWebhookStatusSent
2021
)
2122

0 commit comments

Comments
 (0)