feat: Add missing database role operations to SDK (#2014)

* Add missing operations to database role (WIP) * Add options and validations to missing database role operations (WIP) * Add dto opts mapping to missing database role operations (WIP) * Add unit tests to missing database role operations (WIP) * Add unit tests of validations for missing database role operations (WIP) * Add integration tests for missing database role operations (WIP) * Rename to parent role * Fix granting syntax * Add assertions for granted... fields * Fix grant/revoke database role to/from share test
Snowflake-Labs · Aug 17, 2023 · d2ea67d · d2ea67d
1 parent f5efc09
commit d2ea67d
Show file tree

Hide file tree

Showing 7 changed files with 514 additions and 18 deletions.
diff --git a/pkg/sdk/database_role.go b/pkg/sdk/database_role.go
@@ -13,6 +13,11 @@ type DatabaseRoles interface {
 	Drop(ctx context.Context, request *DropDatabaseRoleRequest) error
 	Show(ctx context.Context, request *ShowDatabaseRoleRequest) ([]DatabaseRole, error)
 	ShowByID(ctx context.Context, id DatabaseObjectIdentifier) (*DatabaseRole, error)
+
+	Grant(ctx context.Context, request *GrantDatabaseRoleRequest) error
+	Revoke(ctx context.Context, request *RevokeDatabaseRoleRequest) error
+	GrantToShare(ctx context.Context, request *GrantDatabaseRoleToShareRequest) error
+	RevokeFromShare(ctx context.Context, request *RevokeDatabaseRoleFromShareRequest) error
 }
 
 // createDatabaseRoleOptions is based on https://docs.snowflake.com/en/sql-reference/sql/create-database-role.
@@ -77,9 +82,9 @@ type databaseRoleDBRow struct {
 	IsDefault              sql.NullString `db:"is_default"`
 	IsCurrent              sql.NullString `db:"is_current"`
 	IsInherited            sql.NullString `db:"is_inherited"`
-	GrantedToRoles         sql.NullString `db:"granted_to_roles"`
-	GrantedToDatabaseRoles sql.NullString `db:"granted_to_database_roles"`
-	GrantedDatabaseRoles   sql.NullString `db:"granted_database_roles"`
+	GrantedToRoles         int            `db:"granted_to_roles"`
+	GrantedToDatabaseRoles int            `db:"granted_to_database_roles"`
+	GrantedDatabaseRoles   int            `db:"granted_database_roles"`
 	Owner                  string         `db:"owner"`
 	Comment                sql.NullString `db:"comment"`
 	OwnerRoleType          sql.NullString `db:"owner_role_type"`
@@ -93,19 +98,22 @@ type DatabaseRole struct {
 	IsDefault              bool
 	IsCurrent              bool
 	IsInherited            bool
-	GrantedToRoles         string
-	GrantedToDatabaseRoles string
-	GrantedDatabaseRoles   string
+	GrantedToRoles         int
+	GrantedToDatabaseRoles int
+	GrantedDatabaseRoles   int
 	Owner                  string
 	Comment                string
 	OwnerRoleType          string
 }
 
 func (row databaseRoleDBRow) convert() *DatabaseRole {
 	databaseRole := DatabaseRole{
-		CreatedOn: row.CreatedOn,
-		Name:      row.Name,
-		Owner:     row.Owner,
+		CreatedOn:              row.CreatedOn,
+		Name:                   row.Name,
+		Owner:                  row.Owner,
+		GrantedToRoles:         row.GrantedToRoles,
+		GrantedToDatabaseRoles: row.GrantedToDatabaseRoles,
+		GrantedDatabaseRoles:   row.GrantedDatabaseRoles,
 	}
 	if row.IsDefault.Valid {
 		databaseRole.IsDefault = row.IsDefault.String == "Y"
@@ -116,15 +124,6 @@ func (row databaseRoleDBRow) convert() *DatabaseRole {
 	if row.IsInherited.Valid {
 		databaseRole.IsInherited = row.IsInherited.String == "Y"
 	}
-	if row.GrantedToRoles.Valid {
-		databaseRole.GrantedToRoles = row.GrantedToRoles.String
-	}
-	if row.GrantedToDatabaseRoles.Valid {
-		databaseRole.GrantedToDatabaseRoles = row.GrantedToDatabaseRoles.String
-	}
-	if row.GrantedDatabaseRoles.Valid {
-		databaseRole.GrantedDatabaseRoles = row.GrantedDatabaseRoles.String
-	}
 	if row.Comment.Valid {
 		databaseRole.Comment = row.Comment.String
 	}
@@ -133,3 +132,45 @@ func (row databaseRoleDBRow) convert() *DatabaseRole {
 	}
 	return &databaseRole
 }
+
+// grantDatabaseRoleOptions is based on https://docs.snowflake.com/en/sql-reference/sql/grant-database-role.
+type grantDatabaseRoleOptions struct {
+	grant        bool                            `ddl:"static" sql:"GRANT"`
+	databaseRole bool                            `ddl:"static" sql:"DATABASE ROLE"`
+	name         DatabaseObjectIdentifier        `ddl:"identifier"`
+	to           bool                            `ddl:"static" sql:"TO"`
+	ParentRole   grantOrRevokeDatabaseRoleObject `ddl:"-"`
+}
+
+// revokeDatabaseRoleOptions is based on https://docs.snowflake.com/en/sql-reference/sql/revoke-database-role.
+type revokeDatabaseRoleOptions struct {
+	revoke       bool                            `ddl:"static" sql:"REVOKE"`
+	databaseRole bool                            `ddl:"static" sql:"DATABASE ROLE"`
+	name         DatabaseObjectIdentifier        `ddl:"identifier"`
+	from         bool                            `ddl:"static" sql:"FROM"`
+	ParentRole   grantOrRevokeDatabaseRoleObject `ddl:"-"`
+}
+
+type grantOrRevokeDatabaseRoleObject struct {
+	// One of
+	DatabaseRoleName *DatabaseObjectIdentifier `ddl:"identifier" sql:"DATABASE ROLE"`
+	AccountRoleName  *AccountObjectIdentifier  `ddl:"identifier" sql:"ROLE"`
+}
+
+// grantDatabaseRoleToShareOptions is based on https://docs.snowflake.com/en/sql-reference/sql/grant-database-role-share.
+type grantDatabaseRoleToShareOptions struct {
+	grant        bool                     `ddl:"static" sql:"GRANT"`
+	databaseRole bool                     `ddl:"static" sql:"DATABASE ROLE"`
+	name         DatabaseObjectIdentifier `ddl:"identifier"`
+	toShare      bool                     `ddl:"static" sql:"TO SHARE"`
+	Share        AccountObjectIdentifier  `ddl:"identifier"`
+}
+
+// revokeDatabaseRoleFromShareOptions is based on https://docs.snowflake.com/en/sql-reference/sql/grant-database-role-share.
+type revokeDatabaseRoleFromShareOptions struct {
+	revoke       bool                     `ddl:"static" sql:"REVOKE"`
+	databaseRole bool                     `ddl:"static" sql:"DATABASE ROLE"`
+	name         DatabaseObjectIdentifier `ddl:"identifier"`
+	fromShare    bool                     `ddl:"static" sql:"FROM SHARE"`
+	Share        AccountObjectIdentifier  `ddl:"identifier"`
+}
diff --git a/pkg/sdk/database_role_dto.go b/pkg/sdk/database_role_dto.go
@@ -1,5 +1,7 @@
 package sdk
 
+//go:generate go run ./dto-builder-generator/main.go
+
 var (
 	_ optionsProvider[createDatabaseRoleOptions] = new(CreateDatabaseRoleRequest)
 	_ optionsProvider[alterDatabaseRoleOptions]  = new(AlterDatabaseRoleRequest)
@@ -43,3 +45,29 @@ type ShowDatabaseRoleRequest struct {
 	like     *Like
 	database AccountObjectIdentifier // required
 }
+
+type GrantDatabaseRoleRequest struct {
+	name DatabaseObjectIdentifier // required
+
+	// One of
+	databaseRole *DatabaseObjectIdentifier
+	accountRole  *AccountObjectIdentifier
+}
+
+type RevokeDatabaseRoleRequest struct {
+	name DatabaseObjectIdentifier // required
+
+	// One of
+	databaseRole *DatabaseObjectIdentifier
+	accountRole  *AccountObjectIdentifier
+}
+
+type GrantDatabaseRoleToShareRequest struct {
+	name  DatabaseObjectIdentifier // required
+	share AccountObjectIdentifier  // required
+}
+
+type RevokeDatabaseRoleFromShareRequest struct {
+	name  DatabaseObjectIdentifier // required
+	share AccountObjectIdentifier  // required
+}
diff --git a/pkg/sdk/database_role_dto_builders.go b/pkg/sdk/database_role_dto_builders.go
@@ -98,3 +98,63 @@ func (s *ShowDatabaseRoleRequest) WithLike(pattern string) *ShowDatabaseRoleRequ
 	}
 	return s
 }
+
+func NewGrantDatabaseRoleRequest(
+	name DatabaseObjectIdentifier,
+) *GrantDatabaseRoleRequest {
+	s := GrantDatabaseRoleRequest{}
+	s.name = name
+	return &s
+}
+
+func (s *GrantDatabaseRoleRequest) WithDatabaseRole(databaseRole DatabaseObjectIdentifier) *GrantDatabaseRoleRequest {
+	s.accountRole = nil
+	s.databaseRole = &databaseRole
+	return s
+}
+
+func (s *GrantDatabaseRoleRequest) WithAccountRole(accountRole AccountObjectIdentifier) *GrantDatabaseRoleRequest {
+	s.databaseRole = nil
+	s.accountRole = &accountRole
+	return s
+}
+
+func NewRevokeDatabaseRoleRequest(
+	name DatabaseObjectIdentifier,
+) *RevokeDatabaseRoleRequest {
+	s := RevokeDatabaseRoleRequest{}
+	s.name = name
+	return &s
+}
+
+func (s *RevokeDatabaseRoleRequest) WithDatabaseRole(databaseRole DatabaseObjectIdentifier) *RevokeDatabaseRoleRequest {
+	s.accountRole = nil
+	s.databaseRole = &databaseRole
+	return s
+}
+
+func (s *RevokeDatabaseRoleRequest) WithAccountRole(accountRole AccountObjectIdentifier) *RevokeDatabaseRoleRequest {
+	s.databaseRole = nil
+	s.accountRole = &accountRole
+	return s
+}
+
+func NewGrantDatabaseRoleToShareRequest(
+	name DatabaseObjectIdentifier,
+	share AccountObjectIdentifier,
+) *GrantDatabaseRoleToShareRequest {
+	s := GrantDatabaseRoleToShareRequest{}
+	s.name = name
+	s.share = share
+	return &s
+}
+
+func NewRevokeDatabaseRoleFromShareRequest(
+	name DatabaseObjectIdentifier,
+	share AccountObjectIdentifier,
+) *RevokeDatabaseRoleFromShareRequest {
+	s := RevokeDatabaseRoleFromShareRequest{}
+	s.name = name
+	s.share = share
+	return &s
+}
diff --git a/pkg/sdk/database_role_impl.go b/pkg/sdk/database_role_impl.go
@@ -45,6 +45,26 @@ func (v *databaseRoles) ShowByID(ctx context.Context, id DatabaseObjectIdentifie
 	return findOne(databaseRoles, func(r DatabaseRole) bool { return r.Name == id.Name() })
 }
 
+func (v *databaseRoles) Grant(ctx context.Context, request *GrantDatabaseRoleRequest) error {
+	opts := request.toOpts()
+	return validateAndExec(v.client, ctx, opts)
+}
+
+func (v *databaseRoles) Revoke(ctx context.Context, request *RevokeDatabaseRoleRequest) error {
+	opts := request.toOpts()
+	return validateAndExec(v.client, ctx, opts)
+}
+
+func (v *databaseRoles) GrantToShare(ctx context.Context, request *GrantDatabaseRoleToShareRequest) error {
+	opts := request.toOpts()
+	return validateAndExec(v.client, ctx, opts)
+}
+
+func (v *databaseRoles) RevokeFromShare(ctx context.Context, request *RevokeDatabaseRoleFromShareRequest) error {
+	opts := request.toOpts()
+	return validateAndExec(v.client, ctx, opts)
+}
+
 func (s *CreateDatabaseRoleRequest) toOpts() *createDatabaseRoleOptions {
 	return &createDatabaseRoleOptions{
 		OrReplace:   Bool(s.orReplace),
@@ -84,3 +104,51 @@ func (s *ShowDatabaseRoleRequest) toOpts() *showDatabaseRoleOptions {
 		Database: s.database,
 	}
 }
+
+func (s *GrantDatabaseRoleRequest) toOpts() *grantDatabaseRoleOptions {
+	opts := grantDatabaseRoleOptions{
+		name: s.name,
+	}
+
+	grantToRole := grantOrRevokeDatabaseRoleObject{}
+	if s.databaseRole != nil {
+		grantToRole.DatabaseRoleName = s.databaseRole
+	}
+	if s.accountRole != nil {
+		grantToRole.AccountRoleName = s.accountRole
+	}
+	opts.ParentRole = grantToRole
+
+	return &opts
+}
+
+func (s *RevokeDatabaseRoleRequest) toOpts() *revokeDatabaseRoleOptions {
+	opts := revokeDatabaseRoleOptions{
+		name: s.name,
+	}
+
+	revokeFromRole := grantOrRevokeDatabaseRoleObject{}
+	if s.databaseRole != nil {
+		revokeFromRole.DatabaseRoleName = s.databaseRole
+	}
+	if s.accountRole != nil {
+		revokeFromRole.AccountRoleName = s.accountRole
+	}
+	opts.ParentRole = revokeFromRole
+
+	return &opts
+}
+
+func (s *GrantDatabaseRoleToShareRequest) toOpts() *grantDatabaseRoleToShareOptions {
+	return &grantDatabaseRoleToShareOptions{
+		name:  s.name,
+		Share: s.share,
+	}
+}
+
+func (s *RevokeDatabaseRoleFromShareRequest) toOpts() *revokeDatabaseRoleFromShareOptions {
+	return &revokeDatabaseRoleFromShareOptions{
+		name:  s.name,
+		Share: s.share,
+	}
+}
diff --git a/pkg/sdk/database_role_integration_test.go b/pkg/sdk/database_role_integration_test.go
@@ -21,6 +21,9 @@ func TestInt_DatabaseRoles(t *testing.T) {
 		assert.Equal(t, expectedName, databaseRole.Name)
 		assert.Equal(t, "ACCOUNTADMIN", databaseRole.Owner)
 		assert.Equal(t, expectedComment, databaseRole.Comment)
+		assert.Equal(t, 0, databaseRole.GrantedToRoles)
+		assert.Equal(t, 0, databaseRole.GrantedToDatabaseRoles)
+		assert.Equal(t, 0, databaseRole.GrantedDatabaseRoles)
 	}
 
 	cleanupDatabaseRoleProvider := func(id DatabaseObjectIdentifier) func() {
@@ -203,4 +206,72 @@ func TestInt_DatabaseRoles(t *testing.T) {
 		require.NoError(t, err)
 		assert.Equal(t, 0, len(returnedDatabaseRoles))
 	})
+
+	t.Run("grant and revoke database_role: to database role", func(t *testing.T) {
+		role1 := createDatabaseRole(t)
+		id1 := NewDatabaseObjectIdentifier(database.Name, role1.Name)
+		role2 := createDatabaseRole(t)
+		id2 := NewDatabaseObjectIdentifier(database.Name, role2.Name)
+
+		grantRequest := NewGrantDatabaseRoleRequest(id1).WithDatabaseRole(id2)
+		err := client.DatabaseRoles.Grant(ctx, grantRequest)
+		require.NoError(t, err)
+
+		extractedRole, err := client.DatabaseRoles.ShowByID(ctx, id1)
+		require.NoError(t, err)
+		assert.Equal(t, 0, extractedRole.GrantedToRoles)
+		assert.Equal(t, 1, extractedRole.GrantedToDatabaseRoles)
+		assert.Equal(t, 0, extractedRole.GrantedDatabaseRoles)
+
+		extractedRole, err = client.DatabaseRoles.ShowByID(ctx, id2)
+		require.NoError(t, err)
+		assert.Equal(t, 0, extractedRole.GrantedToRoles)
+		assert.Equal(t, 0, extractedRole.GrantedToDatabaseRoles)
+		assert.Equal(t, 1, extractedRole.GrantedDatabaseRoles)
+
+		revokeRequest := NewRevokeDatabaseRoleRequest(id1).WithDatabaseRole(id2)
+		err = client.DatabaseRoles.Revoke(ctx, revokeRequest)
+		require.NoError(t, err)
+	})
+
+	t.Run("grant and revoke database_role: to account role", func(t *testing.T) {
+		role := createDatabaseRole(t)
+		roleId := NewDatabaseObjectIdentifier(database.Name, role.Name)
+
+		accountRole, accountRoleCleanup := createRole(t, client)
+		t.Cleanup(accountRoleCleanup)
+
+		grantRequest := NewGrantDatabaseRoleRequest(roleId).WithAccountRole(accountRole.ID())
+		err := client.DatabaseRoles.Grant(ctx, grantRequest)
+		require.NoError(t, err)
+
+		extractedRole, err := client.DatabaseRoles.ShowByID(ctx, roleId)
+		require.NoError(t, err)
+		assert.Equal(t, 1, extractedRole.GrantedToRoles)
+		assert.Equal(t, 0, extractedRole.GrantedToDatabaseRoles)
+		assert.Equal(t, 0, extractedRole.GrantedDatabaseRoles)
+
+		revokeRequest := NewRevokeDatabaseRoleRequest(roleId).WithAccountRole(accountRole.ID())
+		err = client.DatabaseRoles.Revoke(ctx, revokeRequest)
+		require.NoError(t, err)
+	})
+
+	t.Run("grant and revoke database_role: to share", func(t *testing.T) {
+		role := createDatabaseRole(t)
+		roleId := NewDatabaseObjectIdentifier(database.Name, role.Name)
+
+		share, shareCleanup := createShare(t, client)
+		t.Cleanup(shareCleanup)
+
+		err := client.Grants.GrantPrivilegeToShare(ctx, ObjectPrivilegeUsage, &GrantPrivilegeToShareOn{Database: database.ID()}, share.ID())
+		require.NoError(t, err)
+
+		grantRequest := NewGrantDatabaseRoleToShareRequest(roleId, share.ID())
+		err = client.DatabaseRoles.GrantToShare(ctx, grantRequest)
+		require.NoError(t, err)
+
+		revokeRequest := NewRevokeDatabaseRoleFromShareRequest(roleId, share.ID())
+		err = client.DatabaseRoles.RevokeFromShare(ctx, revokeRequest)
+		require.NoError(t, err)
+	})
 }