-
Notifications
You must be signed in to change notification settings - Fork 418
/
resource_monitor.go
278 lines (232 loc) · 9.46 KB
/
resource_monitor.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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
package snowflake
import (
"database/sql"
"errors"
"fmt"
"log"
"strings"
"github.com/jmoiron/sqlx"
)
// ResourceMonitorBuilder extends the generic builder to provide support for triggers.
type ResourceMonitorBuilder struct {
Builder
}
// ResourceMonitor returns a pointer to a ResourceMonitorBuilder that abstracts the DDL operations for a resource monitor.
//
// Supported DDL operations are:
// - CREATE RESOURCE MONITOR
// - ALTER RESOURCE MONITOR
// - DROP RESOURCE MONITOR
// - SHOW RESOURCE MONITOR
//
// [Snowflake Reference](https://docs.snowflake.net/manuals/user-guide/resource-monitors.html#ddl-for-resource-monitors)
func NewResourceMonitorBuilder(name string) *ResourceMonitorBuilder {
return &ResourceMonitorBuilder{
Builder{
entityType: ResourceMonitorType,
name: name,
},
}
}
// @TODO support for a ResourceMonitorAlterBuilder so that we can alter triggers
// ResourceMonitorCreateBuilder extends the generic create builder to provide support for triggers.
type ResourceMonitorCreateBuilder struct {
CreateBuilder
// triggers consist of the type (DO SUSPEND | SUSPEND_IMMEDIATE | NOTIFY) and
// the threshold (a percentage value)
triggers []trigger
}
type ResourceMonitorAlterBuilder struct {
AlterPropertiesBuilder
// triggers consist of the type (DO SUSPEND | SUSPEND_IMMEDIATE | NOTIFY) and
// the threshold (a percentage value)
triggers []trigger
}
type action string
type trigger struct {
action action
threshold int
}
const (
// SuspendTrigger suspends all assigned warehouses while allowing currently running queries to complete.
SuspendTrigger action = "SUSPEND"
// SuspendImmediatelyTrigger suspends all assigned warehouses immediately and cancel any currently running queries or statements using the warehouses.
SuspendImmediatelyTrigger action = "SUSPEND_IMMEDIATE"
// NotifyTrigger sends an alert (to all users who have enabled notifications for themselves), but do not take any other action.
NotifyTrigger action = "NOTIFY"
)
// Create returns a pointer to a ResourceMonitorCreateBuilder.
func (rb *ResourceMonitorBuilder) Create() *ResourceMonitorCreateBuilder {
return &ResourceMonitorCreateBuilder{
CreateBuilder{
name: rb.name,
entityType: rb.entityType,
stringProperties: make(map[string]string),
boolProperties: make(map[string]bool),
intProperties: make(map[string]int),
floatProperties: make(map[string]float64),
stringListProperties: make(map[string][]string),
},
make([]trigger, 0),
}
}
// NotifyAt adds a notify trigger at the specified percentage threshold.
func (rcb *ResourceMonitorCreateBuilder) NotifyAt(pct int) *ResourceMonitorCreateBuilder {
rcb.triggers = append(rcb.triggers, trigger{NotifyTrigger, pct})
return rcb
}
// SuspendAt adds a suspend trigger at the specified percentage threshold.
func (rcb *ResourceMonitorCreateBuilder) SuspendAt(pct int) *ResourceMonitorCreateBuilder {
rcb.triggers = append(rcb.triggers, trigger{SuspendTrigger, pct})
return rcb
}
// SuspendImmediatelyAt adds a suspend immediately trigger at the specified percentage threshold.
func (rcb *ResourceMonitorCreateBuilder) SuspendImmediatelyAt(pct int) *ResourceMonitorCreateBuilder {
rcb.triggers = append(rcb.triggers, trigger{SuspendImmediatelyTrigger, pct})
return rcb
}
// Statement returns the SQL statement needed to actually create the resource.
func (rcb *ResourceMonitorCreateBuilder) Statement() string {
var sb strings.Builder
sb.WriteString(fmt.Sprintf(`CREATE %v "%v"`, rcb.entityType, rcb.name))
for k, v := range rcb.stringProperties {
sb.WriteString(fmt.Sprintf(` %v='%v'`, strings.ToUpper(k), EscapeString(v)))
}
for k, v := range rcb.intProperties {
sb.WriteString(fmt.Sprintf(` %v=%d`, strings.ToUpper(k), v))
}
for k, v := range rcb.floatProperties {
sb.WriteString(fmt.Sprintf(` %v=%.2f`, strings.ToUpper(k), v))
}
for k, v := range rcb.stringListProperties {
sb.WriteString(fmt.Sprintf(" %s=%s", strings.ToUpper(k), formatStringList(v)))
}
if len(rcb.triggers) > 0 {
sb.WriteString(" TRIGGERS")
}
for _, trig := range rcb.triggers {
sb.WriteString(fmt.Sprintf(` ON %d PERCENT DO %v`, trig.threshold, trig.action))
}
return sb.String()
}
// SetOnAccount returns the SQL query that will set the resource monitor globally on your Snowflake account.
func (rcb *ResourceMonitorCreateBuilder) SetOnAccount() string {
return fmt.Sprintf(`ALTER ACCOUNT SET RESOURCE_MONITOR = "%v"`, rcb.name)
}
// SetOnWarehouse returns the SQL query that will set the resource monitor on the specified warehouse.
func (rcb *ResourceMonitorCreateBuilder) SetOnWarehouse(warehouse string) string {
return fmt.Sprintf(`ALTER WAREHOUSE "%v" SET RESOURCE_MONITOR = "%v"`, warehouse, rcb.name)
}
// Alter returns a pointer to a ResourceMonitorAlterBuilder.
func (rb *ResourceMonitorBuilder) Alter() *ResourceMonitorAlterBuilder {
return &ResourceMonitorAlterBuilder{
AlterPropertiesBuilder{
name: rb.name,
entityType: rb.entityType,
stringProperties: make(map[string]string),
boolProperties: make(map[string]bool),
intProperties: make(map[string]int),
floatProperties: make(map[string]float64),
stringListProperties: make(map[string][]string),
},
make([]trigger, 0),
}
}
// NotifyAt adds a notify trigger at the specified percentage threshold.
func (rcb *ResourceMonitorAlterBuilder) NotifyAt(pct int) *ResourceMonitorAlterBuilder {
rcb.triggers = append(rcb.triggers, trigger{NotifyTrigger, pct})
return rcb
}
// SuspendAt adds a suspend trigger at the specified percentage threshold.
func (rcb *ResourceMonitorAlterBuilder) SuspendAt(pct int) *ResourceMonitorAlterBuilder {
rcb.triggers = append(rcb.triggers, trigger{SuspendTrigger, pct})
return rcb
}
// SuspendImmediatelyAt adds a suspend immediately trigger at the specified percentage threshold.
func (rcb *ResourceMonitorAlterBuilder) SuspendImmediatelyAt(pct int) *ResourceMonitorAlterBuilder {
rcb.triggers = append(rcb.triggers, trigger{SuspendImmediatelyTrigger, pct})
return rcb
}
// Statement returns the SQL statement needed to actually alter the resource.
func (rcb *ResourceMonitorAlterBuilder) Statement() string {
var sb strings.Builder
sb.WriteString(fmt.Sprintf(`ALTER %v "%v" SET`, rcb.entityType, rcb.name))
for k, v := range rcb.stringProperties {
sb.WriteString(fmt.Sprintf(` %v='%v'`, strings.ToUpper(k), EscapeString(v)))
}
for k, v := range rcb.intProperties {
sb.WriteString(fmt.Sprintf(` %v=%d`, strings.ToUpper(k), v))
}
for k, v := range rcb.floatProperties {
sb.WriteString(fmt.Sprintf(` %v=%.2f`, strings.ToUpper(k), v))
}
for k, v := range rcb.stringListProperties {
sb.WriteString(fmt.Sprintf(" %s=%s", strings.ToUpper(k), formatStringList(v)))
}
// If the only change is the trigges, then we do not need to add the SET keyword
if strings.HasSuffix(sb.String(), "SET") {
sb.Reset()
sb.WriteString(fmt.Sprintf(`ALTER %v "%v"`, rcb.entityType, rcb.name))
}
if len(rcb.triggers) > 0 {
sb.WriteString(" TRIGGERS")
}
for _, trig := range rcb.triggers {
sb.WriteString(fmt.Sprintf(` ON %d PERCENT DO %v`, trig.threshold, trig.action))
}
return sb.String()
}
// SetOnAccount returns the SQL query that will set the resource monitor globally on your Snowflake account.
func (rcb *ResourceMonitorAlterBuilder) SetOnAccount() string {
return fmt.Sprintf(`ALTER ACCOUNT SET RESOURCE_MONITOR = "%v"`, rcb.name)
}
func (rcb *ResourceMonitorAlterBuilder) UnsetOnAccount() string {
return `ALTER ACCOUNT SET RESOURCE_MONITOR = NULL`
}
// SetOnWarehouse returns the SQL query that will set the resource monitor on the specified warehouse.
func (rcb *ResourceMonitorAlterBuilder) SetOnWarehouse(warehouse string) string {
return fmt.Sprintf(`ALTER WAREHOUSE "%v" SET RESOURCE_MONITOR = "%v"`, warehouse, rcb.name)
}
// UnsetOnWarehouse returns the SQL query that will unset the resource monitor on the specified warehouse.
func (rcb *ResourceMonitorAlterBuilder) UnsetOnWarehouse(warehouse string) string {
return fmt.Sprintf(`ALTER WAREHOUSE "%v" SET RESOURCE_MONITOR = NULL`, warehouse)
}
type ResourceMonitor struct {
Name sql.NullString `db:"name"`
CreditQuota sql.NullString `db:"credit_quota"`
UsedCredits sql.NullString `db:"used_credits"`
RemainingCredits sql.NullString `db:"remaining_credits"`
Level sql.NullString `db:"level"`
Frequency sql.NullString `db:"frequency"`
StartTime sql.NullString `db:"start_time"`
EndTime sql.NullString `db:"end_time"`
NotifyAt sql.NullString `db:"notify_at"`
SuspendAt sql.NullString `db:"suspend_at"`
SuspendImmediatelyAt sql.NullString `db:"suspend_immediately_at"`
CreatedOn sql.NullString `db:"created_on"`
Owner sql.NullString `db:"owner"`
Comment sql.NullString `db:"comment"`
NotifyUsers sql.NullString `db:"notify_users"`
}
func ScanResourceMonitor(row *sqlx.Row) (*ResourceMonitor, error) {
rm := &ResourceMonitor{}
err := row.StructScan(rm)
return rm, err
}
func ListResourceMonitors(db *sql.DB) ([]ResourceMonitor, error) {
stmt := "SHOW RESOURCE MONITORS"
rows, err := Query(db, stmt)
if err != nil {
return nil, err
}
defer rows.Close()
dbs := []ResourceMonitor{}
if err := sqlx.StructScan(rows, &dbs); err != nil {
if errors.Is(err, sql.ErrNoRows) {
log.Println("[DEBUG] no resource monitors found")
return nil, nil
}
return nil, fmt.Errorf("unable to scan row for %s err = %w", stmt, err)
}
return dbs, nil
}