Support owner role type column

Snowflake-Labs · Apr 23, 2024 · 9dc6423 · 9dc6423
1 parent f0a4349
commit 9dc6423
Show file tree

Hide file tree

Showing 10 changed files with 114 additions and 54 deletions.
diff --git a/pkg/sdk/alerts.go b/pkg/sdk/alerts.go
@@ -2,6 +2,7 @@ package sdk
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 	"time"
@@ -217,35 +218,37 @@ func (v *Alert) ObjectType() ObjectType {
 }
 
 type Alert struct {
-	CreatedOn    time.Time
-	Name         string
-	DatabaseName string
-	SchemaName   string
-	Owner        string
-	Comment      *string
-	Warehouse    string
-	Schedule     string
-	State        AlertState
-	Condition    string
-	Action       string
+	CreatedOn     time.Time
+	Name          string
+	DatabaseName  string
+	SchemaName    string
+	Owner         string
+	Comment       *string
+	Warehouse     string
+	Schedule      string
+	State         AlertState
+	Condition     string
+	Action        string
+	OwnerRoleType string
 }
 
 type alertDBRow struct {
-	CreatedOn    time.Time `db:"created_on"`
-	Name         string    `db:"name"`
-	DatabaseName string    `db:"database_name"`
-	SchemaName   string    `db:"schema_name"`
-	Owner        string    `db:"owner"`
-	Comment      *string   `db:"comment"`
-	Warehouse    string    `db:"warehouse"`
-	Schedule     string    `db:"schedule"`
-	State        string    `db:"state"` // suspended, started
-	Condition    string    `db:"condition"`
-	Action       string    `db:"action"`
+	CreatedOn     time.Time      `db:"created_on"`
+	Name          string         `db:"name"`
+	DatabaseName  string         `db:"database_name"`
+	SchemaName    string         `db:"schema_name"`
+	Owner         string         `db:"owner"`
+	Comment       *string        `db:"comment"`
+	Warehouse     string         `db:"warehouse"`
+	Schedule      string         `db:"schedule"`
+	State         string         `db:"state"` // suspended, started
+	Condition     string         `db:"condition"`
+	Action        string         `db:"action"`
+	OwnerRoleType sql.NullString `db:"owner_role_type"`
 }
 
 func (row alertDBRow) convert() *Alert {
-	return &Alert{
+	alert := &Alert{
 		CreatedOn:    row.CreatedOn,
 		Name:         row.Name,
 		DatabaseName: row.DatabaseName,
@@ -258,6 +261,11 @@ func (row alertDBRow) convert() *Alert {
 		Condition:    row.Condition,
 		Action:       row.Action,
 	}
+	if row.OwnerRoleType.Valid {
+		alert.OwnerRoleType = row.OwnerRoleType.String
+	}
+
+	return alert
 }
 
 func (opts *ShowAlertOptions) validate() error {

diff --git a/pkg/sdk/databases.go b/pkg/sdk/databases.go
@@ -57,6 +57,7 @@ type Database struct {
 	DroppedOn     time.Time
 	Transient     bool
 	Kind          string
+	OwnerRoleType string
 }
 
 func (v *Database) ID() AccountObjectIdentifier {
@@ -80,6 +81,7 @@ type databaseRow struct {
 	ResourceGroup sql.NullString `db:"resource_group"`
 	DroppedOn     sql.NullTime   `db:"dropped_on"`
 	Kind          sql.NullString `db:"kind"`
+	OwnerRoleType sql.NullString `db:"owner_role_type"`
 }
 
 func (row databaseRow) convert() *Database {
@@ -129,6 +131,9 @@ func (row databaseRow) convert() *Database {
 	if row.Kind.Valid {
 		database.Kind = row.Kind.String
 	}
+	if row.OwnerRoleType.Valid {
+		database.OwnerRoleType = row.OwnerRoleType.String
+	}
 	return database
 }
 

diff --git a/pkg/sdk/dynamic_table.go b/pkg/sdk/dynamic_table.go
@@ -124,6 +124,7 @@ type DynamicTable struct {
 	IsClone             bool
 	IsReplica           bool
 	DataTimestamp       time.Time
+	OwnerRoleType       string
 }
 
 func (dt *DynamicTable) ID() SchemaObjectIdentifier {
@@ -152,6 +153,7 @@ type dynamicTableRow struct {
 	IsClone             bool           `db:"is_clone"`
 	IsReplica           bool           `db:"is_replica"`
 	DataTimestamp       sql.NullTime   `db:"data_timestamp"`
+	OwnerRoleType       sql.NullString `db:"owner_role_type"`
 }
 
 func (dtr dynamicTableRow) convert() *DynamicTable {
@@ -184,6 +186,10 @@ func (dtr dynamicTableRow) convert() *DynamicTable {
 	if dtr.LastSuspendedOn.Valid {
 		dt.LastSuspendedOn = dtr.LastSuspendedOn.Time
 	}
+	if dtr.OwnerRoleType.Valid {
+		dt.OwnerRoleType = dtr.OwnerRoleType.String
+	}
+
 	return dt
 }
 

diff --git a/pkg/sdk/testint/alerts_integration_test.go b/pkg/sdk/testint/alerts_integration_test.go
@@ -424,4 +424,14 @@ func TestInt_AlertsShowByID(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, id2, e2.ID())
 	})
+
+	t.Run("show by id: check fields", func(t *testing.T) {
+		name := random.AlphaN(4)
+		id1 := sdk.NewSchemaObjectIdentifier(databaseTest.Name, schemaTest.Name, name)
+		createAlertHandle(t, id1)
+
+		alert, err := client.Alerts.ShowByID(ctx, id1)
+		require.NoError(t, err)
+		assert.Equal(t, "ROLE", alert.OwnerRoleType)
+	})
 }
diff --git a/pkg/sdk/testint/databases_integration_test.go b/pkg/sdk/testint/databases_integration_test.go
@@ -330,6 +330,7 @@ func TestInt_DatabasesShow(t *testing.T) {
 		}
 		assert.Contains(t, databaseIDs, databaseTest.ID())
 		assert.Contains(t, databaseIDs, databaseTest2.ID())
+		assert.Equal(t, "ROLE", databases[0].OwnerRoleType)
 	})
 
 	t.Run("with terse", func(t *testing.T) {

diff --git a/pkg/sdk/testint/dynamic_table_integration_test.go b/pkg/sdk/testint/dynamic_table_integration_test.go
@@ -46,6 +46,7 @@ func TestInt_DynamicTableCreateAndDrop(t *testing.T) {
 		require.Equal(t, name.Name(), dynamicTableById.Name)
 		require.Equal(t, testWarehouse(t).ID().Name(), dynamicTableById.Warehouse)
 		require.Equal(t, *targetLag.MaximumDuration, dynamicTableById.TargetLag)
+		assert.Equal(t, "ROLE", dynamicTableById.OwnerRoleType)
 	})
 
 	t.Run("test complete with target lag", func(t *testing.T) {

diff --git a/pkg/sdk/testint/network_rule_gen_integration_test.go b/pkg/sdk/testint/network_rule_gen_integration_test.go
@@ -182,4 +182,14 @@ func TestInt_NetworkRulesShowByID(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, id2, e2.ID())
 	})
+
+	t.Run("show by id: check fields", func(t *testing.T) {
+		name := random.AlphaN(4)
+		id1 := sdk.NewSchemaObjectIdentifier(databaseTest.Name, schemaTest.Name, name)
+		createNetworkRuleHandle(t, id1)
+
+		nr, err := client.NetworkRules.ShowByID(ctx, id1)
+		require.NoError(t, err)
+		assert.Equal(t, "ROLE", nr.OwnerRoleType)
+	})
 }
diff --git a/pkg/sdk/testint/streamlits_integration_test.go b/pkg/sdk/testint/streamlits_integration_test.go
@@ -340,4 +340,17 @@ func TestInt_StreamlitsShowByID(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, id2, e2.ID())
 	})
+
+	t.Run("show by id: check fields", func(t *testing.T) {
+		name := random.AlphaN(4)
+		stage, cleanupStage := createStage(t, client, sdk.NewSchemaObjectIdentifier(TestDatabaseName, TestSchemaName, random.AlphaN(4)))
+		t.Cleanup(cleanupStage)
+		manifest := "manifest.yml"
+		id1 := sdk.NewSchemaObjectIdentifier(databaseTest.Name, schemaTest.Name, name)
+		createStreamlitHandle(t, id1, stage, manifest)
+
+		sl, err := client.Streamlits.ShowByID(ctx, id1)
+		require.NoError(t, err)
+		assert.Equal(t, "ROLE", sl.OwnerRoleType)
+	})
 }
diff --git a/pkg/sdk/testint/warehouses_integration_test.go b/pkg/sdk/testint/warehouses_integration_test.go
@@ -39,6 +39,7 @@ func TestInt_WarehousesShow(t *testing.T) {
 		assert.Equal(t, 1, len(warehouses))
 		assert.Equal(t, warehouse.Name, warehouses[0].Name)
 		assert.Equal(t, sdk.WarehouseSizeSmall, warehouses[0].Size)
+		assert.Equal(t, "ROLE", warehouses[0].OwnerRoleType)
 	})
 
 	t.Run("when searching a non-existent password policy", func(t *testing.T) {

diff --git a/pkg/sdk/warehouses.go b/pkg/sdk/warehouses.go
@@ -378,40 +378,42 @@ type Warehouse struct {
 	QueryAccelerationMaxScaleFactor int
 	ResourceMonitor                 string
 	ScalingPolicy                   ScalingPolicy
+	OwnerRoleType                   string
 }
 
 type warehouseDBRow struct {
-	Name                            string        `db:"name"`
-	State                           string        `db:"state"`
-	Type                            string        `db:"type"`
-	Size                            string        `db:"size"`
-	MinClusterCount                 int           `db:"min_cluster_count"`
-	MaxClusterCount                 int           `db:"max_cluster_count"`
-	StartedClusters                 int           `db:"started_clusters"`
-	Running                         int           `db:"running"`
-	Queued                          int           `db:"queued"`
-	IsDefault                       string        `db:"is_default"`
-	IsCurrent                       string        `db:"is_current"`
-	AutoSuspend                     sql.NullInt64 `db:"auto_suspend"`
-	AutoResume                      bool          `db:"auto_resume"`
-	Available                       string        `db:"available"`
-	Provisioning                    string        `db:"provisioning"`
-	Quiescing                       string        `db:"quiescing"`
-	Other                           string        `db:"other"`
-	CreatedOn                       time.Time     `db:"created_on"`
-	ResumedOn                       time.Time     `db:"resumed_on"`
-	UpdatedOn                       time.Time     `db:"updated_on"`
-	Owner                           string        `db:"owner"`
-	Comment                         string        `db:"comment"`
-	EnableQueryAcceleration         bool          `db:"enable_query_acceleration"`
-	QueryAccelerationMaxScaleFactor int           `db:"query_acceleration_max_scale_factor"`
-	ResourceMonitor                 string        `db:"resource_monitor"`
-	Actives                         string        `db:"actives"`
-	Pendings                        string        `db:"pendings"`
-	Failed                          string        `db:"failed"`
-	Suspended                       string        `db:"suspended"`
-	UUID                            string        `db:"uuid"`
-	ScalingPolicy                   string        `db:"scaling_policy"`
+	Name                            string         `db:"name"`
+	State                           string         `db:"state"`
+	Type                            string         `db:"type"`
+	Size                            string         `db:"size"`
+	MinClusterCount                 int            `db:"min_cluster_count"`
+	MaxClusterCount                 int            `db:"max_cluster_count"`
+	StartedClusters                 int            `db:"started_clusters"`
+	Running                         int            `db:"running"`
+	Queued                          int            `db:"queued"`
+	IsDefault                       string         `db:"is_default"`
+	IsCurrent                       string         `db:"is_current"`
+	AutoSuspend                     sql.NullInt64  `db:"auto_suspend"`
+	AutoResume                      bool           `db:"auto_resume"`
+	Available                       string         `db:"available"`
+	Provisioning                    string         `db:"provisioning"`
+	Quiescing                       string         `db:"quiescing"`
+	Other                           string         `db:"other"`
+	CreatedOn                       time.Time      `db:"created_on"`
+	ResumedOn                       time.Time      `db:"resumed_on"`
+	UpdatedOn                       time.Time      `db:"updated_on"`
+	Owner                           string         `db:"owner"`
+	Comment                         string         `db:"comment"`
+	EnableQueryAcceleration         bool           `db:"enable_query_acceleration"`
+	QueryAccelerationMaxScaleFactor int            `db:"query_acceleration_max_scale_factor"`
+	ResourceMonitor                 string         `db:"resource_monitor"`
+	Actives                         string         `db:"actives"`
+	Pendings                        string         `db:"pendings"`
+	Failed                          string         `db:"failed"`
+	Suspended                       string         `db:"suspended"`
+	UUID                            string         `db:"uuid"`
+	ScalingPolicy                   string         `db:"scaling_policy"`
+	OwnerRoleType                   sql.NullString `db:"owner_role_type"`
 }
 
 func (row warehouseDBRow) convert() *Warehouse {
@@ -453,6 +455,9 @@ func (row warehouseDBRow) convert() *Warehouse {
 	if row.AutoSuspend.Valid {
 		wh.AutoSuspend = int(row.AutoSuspend.Int64)
 	}
+	if row.OwnerRoleType.Valid {
+		wh.OwnerRoleType = row.OwnerRoleType.String
+	}
 	return wh
 }