-
Notifications
You must be signed in to change notification settings - Fork 100
/
resource_watermark.go
125 lines (104 loc) · 3.16 KB
/
resource_watermark.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
package queries
import (
"context"
"fmt"
"io"
"time"
runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
"github.com/rilldata/rill/runtime"
"github.com/rilldata/rill/runtime/drivers"
)
type ResourceWatermark struct {
ResourceKind string `json:"resource_kind,omitempty"`
ResourceName string `json:"resource_name,omitempty"`
Result *time.Time `json:"-"`
}
var _ runtime.Query = &ResourceWatermark{}
func (q *ResourceWatermark) Key() string {
return fmt.Sprintf("ResourceWatermark:%s/%s", q.ResourceKind, q.ResourceName)
}
func (q *ResourceWatermark) Deps() []*runtimev1.ResourceName {
return []*runtimev1.ResourceName{
{Kind: q.ResourceKind, Name: q.ResourceName},
}
}
func (q *ResourceWatermark) MarshalResult() *runtime.QueryResult {
return &runtime.QueryResult{
Value: q.Result,
Bytes: 24,
}
}
func (q *ResourceWatermark) UnmarshalResult(v any) error {
res, ok := v.(*time.Time)
if !ok {
return fmt.Errorf("ResourceWatermark: mismatched unmarshal input")
}
q.Result = res
return nil
}
func (q *ResourceWatermark) Resolve(ctx context.Context, rt *runtime.Runtime, instanceID string, priority int) error {
ctrl, err := rt.Controller(ctx, instanceID)
if err != nil {
return err
}
rs, err := ctrl.Get(ctx, &runtimev1.ResourceName{Kind: q.ResourceKind, Name: q.ResourceName}, false)
if err != nil {
return err
}
switch q.ResourceKind {
case runtime.ResourceKindMetricsView:
return q.resolveMetricsView(ctx, rt, instanceID, priority, rs)
default:
// For resources without watermark support, q.Result will be nil.
return nil
}
}
func (q *ResourceWatermark) Export(ctx context.Context, rt *runtime.Runtime, instanceID string, w io.Writer, opts *runtime.ExportOptions) error {
return ErrExportNotSupported
}
func (q *ResourceWatermark) resolveMetricsView(ctx context.Context, rt *runtime.Runtime, instanceID string, priority int, rs *runtimev1.Resource) error {
mv := rs.GetMetricsView()
if mv == nil {
return fmt.Errorf("internal: resource %q is not a metrics view", rs.Meta.Name.Name)
}
spec := mv.State.ValidSpec
if spec == nil {
return fmt.Errorf("metrics view %q is not valid", rs.Meta.Name.Name)
}
sql := ""
if spec.WatermarkExpression != "" {
sql = fmt.Sprintf("SELECT %s FROM %s", spec.WatermarkExpression, safeName(spec.Table))
} else if spec.TimeDimension != "" {
sql = fmt.Sprintf("SELECT MAX(%s) FROM %s", safeName(spec.TimeDimension), safeName(spec.Table))
} else {
// No watermark available
return nil
}
olap, release, err := rt.OLAP(ctx, instanceID, spec.Connector)
if err != nil {
return err
}
defer release()
res, err := olap.Execute(ctx, &drivers.Statement{
Query: sql,
Priority: priority,
ExecutionTimeout: defaultExecutionTimeout,
})
if err != nil {
return err
}
defer res.Close()
var t time.Time
for res.Next() {
if err := res.Scan(&t); err != nil {
return err
}
}
if !t.IsZero() {
// Hacky workaround for the following issue: the watermark is used as the *exclusive* upper bound for time ranges.
// We need to add a small delta to ensure the row with the watermark is included in the result.
t = t.Add(time.Second)
q.Result = &t
}
return nil
}