@@ -5,44 +5,58 @@ import (
5
5
"errors"
6
6
"time"
7
7
8
+ "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs"
8
9
"github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs/types"
9
10
)
10
11
11
12
// QueryAndTailLogGroup queries the log group from the give start time and initiates a Live Tail session.
12
13
// This function also handles the case where the log group does not exist yet.
13
14
// The caller should call `Close()` on the returned EventStream when done.
14
- func QueryAndTailLogGroup (ctx context.Context , lgi LogGroupInput , start time.Time ) (EventStream , error ) {
15
+ func QueryAndTailLogGroup (ctx context.Context , lgi LogGroupInput , start , end time.Time ) (EventStream , error ) {
15
16
ctx , cancel := context .WithCancel (ctx )
16
17
17
18
es := & eventStream {
18
19
cancel : cancel ,
19
20
ch : make (chan types.StartLiveTailResponseStream ),
20
21
}
21
22
22
- // First call TailLogGroup once to check if the log group exists or we have another error
23
- eventStream , err := TailLogGroup (ctx , lgi )
24
- if err != nil {
25
- var resourceNotFound * types.ResourceNotFoundException
26
- if ! errors .As (err , & resourceNotFound ) {
27
- return nil , err
23
+ doTail := ! end .IsZero ()
24
+
25
+ var tailStream cloudwatchlogs.StartLiveTailResponseStreamReader
26
+ if doTail {
27
+ // First call TailLogGroup once to check if the log group exists or we have another error
28
+ var err error
29
+ tailStream , err = TailLogGroup (ctx , lgi )
30
+ if err != nil {
31
+ var resourceNotFound * types.ResourceNotFoundException
32
+ if ! errors .As (err , & resourceNotFound ) {
33
+ return nil , err
34
+ }
35
+ // Doesn't exist yet, continue to poll for it
28
36
}
29
37
}
30
38
31
39
// Start goroutine to wait for the log group to be created and then tail it
32
40
go func () {
33
41
defer close (es .ch )
34
42
35
- if eventStream == nil {
43
+ if doTail {
36
44
// If the log group does not exist yet, poll until it does
37
- eventStream , err = pollTailLogGroup (ctx , lgi )
38
- if err != nil {
39
- es .err = err
40
- return
45
+ if tailStream == nil {
46
+ var err error
47
+ tailStream , err = pollTailLogGroup (ctx , lgi )
48
+ if err != nil {
49
+ es .err = err
50
+ return
51
+ }
41
52
}
53
+ defer tailStream .Close ()
42
54
}
43
- defer eventStream .Close ()
44
55
45
56
if ! start .IsZero () {
57
+ if end .IsZero () {
58
+ end = time .Now ()
59
+ }
46
60
// Query the logs between the start time and now
47
61
if err := QueryLogGroup (ctx , lgi , start , time .Now (), func (events []LogEvent ) error {
48
62
es .ch <- & types.StartLiveTailResponseStreamMemberSessionUpdate {
@@ -55,14 +69,17 @@ func QueryAndTailLogGroup(ctx context.Context, lgi LogGroupInput, start time.Tim
55
69
}
56
70
}
57
71
58
- es .pipeEvents (ctx , eventStream )
72
+ if doTail {
73
+ // Pipe the events from the tail stream to the internal channel
74
+ es .pipeEvents (ctx , tailStream )
75
+ }
59
76
}()
60
77
61
78
return es , nil
62
79
}
63
80
64
81
// pollTailLogGroup polls the log group and starts the Live Tail session once it's available
65
- func pollTailLogGroup (ctx context.Context , lgi LogGroupInput ) (EventStream , error ) {
82
+ func pollTailLogGroup (ctx context.Context , lgi LogGroupInput ) (cloudwatchlogs. StartLiveTailResponseStreamReader , error ) {
66
83
ticker := time .NewTicker (time .Second )
67
84
defer ticker .Stop ()
68
85
0 commit comments