diff --git a/pkg/snowflake/generic.go b/pkg/snowflake/generic.go index 44246c8f2d..887f829bcc 100644 --- a/pkg/snowflake/generic.go +++ b/pkg/snowflake/generic.go @@ -9,12 +9,13 @@ import ( type EntityType string const ( - DatabaseType EntityType = "DATABASE" - ManagedAccountType EntityType = "MANAGED ACCOUNT" - RoleType EntityType = "ROLE" - ShareType EntityType = "SHARE" - UserType EntityType = "USER" - WarehouseType EntityType = "WAREHOUSE" + DatabaseType EntityType = "DATABASE" + ManagedAccountType EntityType = "MANAGED ACCOUNT" + ResourceMonitorType EntityType = "RESOURCE MONITOR" + RoleType EntityType = "ROLE" + ShareType EntityType = "SHARE" + UserType EntityType = "USER" + WarehouseType EntityType = "WAREHOUSE" ) type Builder struct { diff --git a/pkg/snowflake/resource_monitor.go b/pkg/snowflake/resource_monitor.go new file mode 100644 index 0000000000..33f2c3b6db --- /dev/null +++ b/pkg/snowflake/resource_monitor.go @@ -0,0 +1,108 @@ +package snowflake + +import ( + "fmt" + "strings" +) + +// 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 ResourceMonitor(name string) *ResourceMonitorBuilder { + return &ResourceMonitorBuilder{ + Builder{ + entityType: ResourceMonitorType, + name: name, + }, + } +} + +// 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 trigger struct { + action string + threshold int +} + +const ( + // SuspendTrigger suspends all assigned warehouses while allowing currently running queries to complete. + SuspendTrigger = "SUSPEND" + // SuspendImmediatelyTrigger suspends all assigned warehouses immediately and cancel any currently running queries or statements using the warehouses. + SuspendImmediatelyTrigger = "SUSPEND_IMMEDIATE" + // NotifyTrigger sends an alert (to all users who have enabled notifications for themselves), but do not take any other action. + NotifyTrigger = "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), + }, + 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)) + } + + 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() +} diff --git a/pkg/snowflake/resource_monitor_test.go b/pkg/snowflake/resource_monitor_test.go new file mode 100644 index 0000000000..487ccad1dc --- /dev/null +++ b/pkg/snowflake/resource_monitor_test.go @@ -0,0 +1,35 @@ +package snowflake_test + +import ( + "testing" + + "github.com/chanzuckerberg/terraform-provider-snowflake/pkg/snowflake" + "github.com/stretchr/testify/assert" +) + +func TestResourceMonitor(t *testing.T) { + a := assert.New(t) + rm := snowflake.ResourceMonitor("resource_monitor") + a.NotNil(rm) + + q := rm.Show() + a.Equal(`SHOW RESOURCE MONITORS LIKE 'resource_monitor'`, q) + + q = rm.Create().Statement() + a.Equal(`CREATE RESOURCE MONITOR "resource_monitor"`, q) + + q = rm.Drop() + a.Equal(`DROP RESOURCE MONITOR "resource_monitor"`, q) + + ab := rm.Alter() + ab.SetInt("credit_quota", 66) + q = ab.Statement() + a.Equal(`ALTER RESOURCE MONITOR "resource_monitor" SET CREDIT_QUOTA=66`, q) + + cb := snowflake.ResourceMonitor("resource_monitor").Create() + cb.NotifyAt(80).NotifyAt(90).SuspendAt(95).SuspendImmediatelyAt(100) + cb.SetString("frequency", "YEARLY") + cb.SetInt("credit_quota", 666) + q = cb.Statement() + a.Equal(`CREATE RESOURCE MONITOR "resource_monitor" FREQUENCY='YEARLY' CREDIT_QUOTA=666 TRIGGERS ON 80 PERCENT DO NOTIFY ON 90 PERCENT DO NOTIFY ON 95 PERCENT DO SUSPEND ON 100 PERCENT DO SUSPEND_IMMEDIATE`, q) +}