1// Requires golang.org/x/tools/cmd/goyacc and modernc.org/golex
2//
3//go:generate goyacc -o datemath.y.go datemath.y
4//go:generate golex -o datemath.l.go datemath.l
5
6/*
7Package datemath provides an expression language for relative dates based on Elasticsearch's date math.
8
9This package is useful for letting end-users describe dates in a simple format similar to Grafana and Kibana and for
10persisting them as relative dates.
11
12The expression starts with an anchor date, which can either be "now", or an ISO8601 date string ending with ||. This
13anchor date can optionally be followed by one or more date math expressions, for example:
14
15	now+1h	Add one hour
16	now-1d	Subtract one day
17	now/d	Round down to the nearest day
18
19The supported time units are:
20	y Years
21	M Months
22	w Weeks
23	d Days
24	b Business Days (excludes Saturday and Sunday by default, use WithBusinessDayFunc to override)
25	h Hours
26	H Hours
27	m Minutes
28	s Seconds
29
30Compatibility with Elasticsearch datemath
31
32This package aims to be a superset of Elasticsearch's expressions. That is, any datemath expression that is valid for
33Elasticsearch should evaluate in the same way here.
34
35Currently the package does not support expressions outside of those also considered valid by Elasticsearch, but this may
36change in the future to include additional functionality.
37*/
38package datemath
39
40import (
41	"fmt"
42	"strconv"
43	"strings"
44	"time"
45)
46
47func init() {
48	// have goyacc parser return more verbose syntax error messages
49	yyErrorVerbose = true
50}
51
52var missingTimeZone = time.FixedZone("MISSING", 0)
53
54type timeUnit rune
55
56const (
57	timeUnitYear        = timeUnit('y')
58	timeUnitMonth       = timeUnit('M')
59	timeUnitWeek        = timeUnit('w')
60	timeUnitDay         = timeUnit('d')
61	timeUnitBusinessDay = timeUnit('b')
62	timeUnitHour        = timeUnit('h')
63	timeUnitMinute      = timeUnit('m')
64	timeUnitSecond      = timeUnit('s')
65)
66
67func (u timeUnit) String() string {
68	return string(u)
69}
70
71// Expression represents a parsed datemath expression
72type Expression struct {
73	input string
74
75	mathExpression
76}
77
78type mathExpression struct {
79	anchorDateExpression anchorDateExpression
80	adjustments          []timeAdjuster
81}
82
83func newMathExpression(anchorDateExpression anchorDateExpression, adjustments []timeAdjuster) mathExpression {
84	return mathExpression{
85		anchorDateExpression: anchorDateExpression,
86		adjustments:          adjustments,
87	}
88}
89
90// MarshalJSON implements the json.Marshaler interface
91//
92// It serializes as the string expression the Expression was created with
93func (e Expression) MarshalJSON() ([]byte, error) {
94	return []byte(strconv.Quote(e.String())), nil
95}
96
97// UnmarshalJSON implements the json.Unmarshaler interface
98//
99// Parses the datemath expression from a JSON string
100func (e *Expression) UnmarshalJSON(data []byte) error {
101	s, err := strconv.Unquote(string(data))
102	if err != nil {
103		return err
104	}
105
106	expression, err := Parse(s)
107	if err != nil {
108		return nil
109	}
110
111	*e = expression
112	return nil
113}
114
115// String returns a the string used to create the expression
116func (e Expression) String() string {
117	return e.input
118}
119
120// Options represesent configurable behavior for interpreting the datemath expression
121type Options struct {
122	// Use this this time as "now"
123	// Default is `time.Now()`
124	Now time.Time
125
126	// Use this location if there is no timezone in the expression
127	// Defaults to time.UTC
128	Location *time.Location
129
130	// Use this weekday as the start of the week
131	// Defaults to time.Monday
132	StartOfWeek time.Weekday
133
134	// Rounding to period should be done to the end of the period
135	// Defaults to false
136	RoundUp bool
137
138	BusinessDayFunc func(time.Time) bool
139}
140
141// WithNow use the given time as "now"
142func WithNow(now time.Time) func(*Options) {
143	return func(o *Options) {
144		o.Now = now
145	}
146}
147
148// WithStartOfWeek uses the given weekday as the start of the week
149func WithStartOfWeek(day time.Weekday) func(*Options) {
150	return func(o *Options) {
151		o.StartOfWeek = day
152	}
153}
154
155// WithLocation uses the given location as the timezone of the date if unspecified
156func WithLocation(l *time.Location) func(*Options) {
157	return func(o *Options) {
158		o.Location = l
159	}
160}
161
162// WithRoundUp sets the rounding of time to the end of the period instead of the beginning
163func WithRoundUp(b bool) func(*Options) {
164	return func(o *Options) {
165		o.RoundUp = b
166	}
167}
168
169// WithBusinessDayFunc use the given fn to check if a day is a business day
170func WithBusinessDayFunc(fn func(time.Time) bool) func(*Options) {
171	return func(o *Options) {
172		o.BusinessDayFunc = fn
173	}
174}
175
176func isNotWeekend(t time.Time) bool {
177	return t.Weekday() != time.Saturday && t.Weekday() != time.Sunday
178}
179
180// Time evaluate the expression with the given options to get the time it represents
181func (e Expression) Time(opts ...func(*Options)) time.Time {
182	options := Options{
183		Now:         time.Now(),
184		Location:    time.UTC,
185		StartOfWeek: time.Monday,
186	}
187	for _, opt := range opts {
188		opt(&options)
189	}
190
191	t := e.anchorDateExpression(options)
192	for _, adjustment := range e.adjustments {
193		t = adjustment(t, options)
194	}
195	return t
196}
197
198// Parse parses the datemath expression which can later be evaluated
199func Parse(s string) (Expression, error) {
200	lex := newLexer([]byte(s))
201	lexWrapper := newLexerWrapper(lex)
202
203	yyParse(lexWrapper)
204
205	if len(lex.errors) > 0 {
206		return Expression{}, fmt.Errorf(strings.Join(lex.errors, "\n"))
207	}
208
209	return Expression{input: s, mathExpression: lexWrapper.expression}, nil
210}
211
212// MustParse is the same as Parse() but panic's on error
213func MustParse(s string) Expression {
214	e, err := Parse(s)
215	if err != nil {
216		panic(err)
217	}
218	return e
219}
220
221// ParseAndEvaluate is a convience wrapper to parse and return the time that the expression represents
222func ParseAndEvaluate(s string, opts ...func(*Options)) (time.Time, error) {
223	expression, err := Parse(s)
224	if err != nil {
225		return time.Time{}, err
226	}
227
228	return expression.Time(opts...), nil
229}
230
231type anchorDateExpression func(opts Options) time.Time
232
233func anchorDateNow(opts Options) time.Time {
234	return opts.Now.In(opts.Location)
235}
236
237func anchorDate(t time.Time) func(opts Options) time.Time {
238	return func(opts Options) time.Time {
239		location := t.Location()
240		if location == missingTimeZone {
241			location = opts.Location
242		}
243
244		return time.Date(t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second(), t.Nanosecond(), location)
245	}
246}
247
248type timeAdjuster func(time.Time, Options) time.Time
249
250func addUnits(factor int, u timeUnit) func(time.Time, Options) time.Time {
251	return func(t time.Time, options Options) time.Time {
252		switch u {
253		case timeUnitYear:
254			return t.AddDate(factor, 0, 0)
255		case timeUnitMonth:
256			return t.AddDate(0, factor, 0)
257		case timeUnitWeek:
258			return t.AddDate(0, 0, 7*factor)
259		case timeUnitDay:
260			return t.AddDate(0, 0, factor)
261		case timeUnitBusinessDay:
262
263			fn := options.BusinessDayFunc
264			if fn == nil {
265				fn = isNotWeekend
266			}
267
268			increment := 1
269			if factor < 0 {
270				increment = -1
271			}
272
273			for i := factor; i != 0; i -= increment {
274				t = t.AddDate(0, 0, increment)
275				for !fn(t) {
276					t = t.AddDate(0, 0, increment)
277				}
278			}
279
280			return t
281
282		case timeUnitHour:
283			return t.Add(time.Duration(factor) * time.Hour)
284		case timeUnitMinute:
285			return t.Add(time.Duration(factor) * time.Minute)
286		case timeUnitSecond:
287			return t.Add(time.Duration(factor) * time.Second)
288		default:
289			panic(fmt.Sprintf("unknown time unit: %s", u))
290		}
291	}
292}
293
294func truncateUnits(u timeUnit) func(time.Time, Options) time.Time {
295	var roundDown = func(t time.Time, options Options) time.Time {
296		switch u {
297		case timeUnitYear:
298			return time.Date(t.Year(), 1, 1, 0, 0, 0, 0, t.Location())
299		case timeUnitMonth:
300			return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, t.Location())
301		case timeUnitWeek:
302			diff := int(t.Weekday() - options.StartOfWeek)
303			if diff < 0 {
304				return time.Date(t.Year(), t.Month(), t.Day()+diff-1, 0, 0, 0, 0, t.Location())
305			}
306			return time.Date(t.Year(), t.Month(), t.Day()-diff, 0, 0, 0, 0, t.Location())
307		case timeUnitDay:
308			return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, t.Location())
309		case timeUnitHour:
310			return t.Truncate(time.Hour)
311		case timeUnitMinute:
312			return t.Truncate(time.Minute)
313		case timeUnitSecond:
314			return t.Truncate(time.Second)
315		default:
316			panic(fmt.Sprintf("unknown time unit: %s", u))
317		}
318	}
319
320	return func(t time.Time, options Options) time.Time {
321		if options.RoundUp {
322			return addUnits(1, u)(roundDown(t, options), options).Add(-time.Millisecond)
323		}
324		return roundDown(t, options)
325	}
326}
327
328func daysIn(m time.Month, year int) int {
329	return time.Date(year, m+1, 0, 0, 0, 0, 0, time.UTC).Day()
330}
331
332// lexerWrapper wraps the golex generated wrapper to store the parsed expression for later and provide needed data to
333// the parser
334type lexerWrapper struct {
335	lex yyLexer
336
337	expression mathExpression
338}
339
340func newLexerWrapper(lex yyLexer) *lexerWrapper {
341	return &lexerWrapper{
342		lex: lex,
343	}
344}
345
346func (l *lexerWrapper) Lex(lval *yySymType) int {
347	return l.lex.Lex(lval)
348}
349
350func (l *lexerWrapper) Error(s string) {
351	l.lex.Error(s)
352}
353