feat(bigquery): allow custom roles for bigquery dataset resource type (…

…#95)
goto · Oct 16, 2023 · f214eac · f214eac
1 parent 502d2bf
commit f214eac
Show file tree

Hide file tree

Showing 2 changed files with 124 additions and 38 deletions.
diff --git a/plugins/providers/bigquery/client.go b/plugins/providers/bigquery/client.go
@@ -11,7 +11,6 @@ import (
 	bqApi "google.golang.org/api/bigquery/v2"
 	"google.golang.org/api/cloudresourcemanager/v1"
 	"google.golang.org/api/iam/v1"
-	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
 )
 
@@ -365,55 +364,99 @@ func (c *bigQueryClient) CheckGrantedPermission(ctx context.Context, permissions
 	return res.Permissions, nil
 }
 
-func (c *bigQueryClient) getGrantableRolesForTables() ([]string, error) {
-	var resourceName string
-	ctx := context.Background()
-	datasetIterator := c.client.Datasets(ctx)
-	for {
-		dataset, err := datasetIterator.Next()
-		if err == iterator.Done {
-			break
-		}
-		if err != nil {
-			return nil, err
-		}
-
-		tableIterator := c.client.Dataset(dataset.DatasetID).Tables(ctx)
-		for {
-			table, err := tableIterator.Next()
-			if err == iterator.Done {
-				break
-			}
-			if err != nil {
-				return nil, err
-			}
+func (c *bigQueryClient) getGrantableRolesForDataset(ctx context.Context) ([]string, error) {
+	sampleDataset, err := c.getSampleDataset(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting a sample dataset, %w", err)
+	}
+	resourceName := fmt.Sprintf("//bigquery.googleapis.com/projects/%v/datasets/%v", sampleDataset.ProjectId, sampleDataset.DatasetId)
 
-			resourceName = fmt.Sprintf("//bigquery.googleapis.com/projects/%v/datasets/%v/tables/%v", table.ProjectID, table.DatasetID, table.TableID)
-			break
-		}
-		if resourceName != "" {
-			break
+	var grantableRoles []string
+	request := &iam.QueryGrantableRolesRequest{
+		FullResourceName: resourceName,
+	}
+	if err := c.iamService.Roles.QueryGrantableRoles(request).Pages(ctx, func(page *iam.QueryGrantableRolesResponse) error {
+		for _, role := range page.Roles {
+			grantableRoles = append(grantableRoles, role.Name)
 		}
+		return nil
+	}); err != nil {
+		return nil, err
 	}
 
-	if resourceName == "" {
-		return nil, ErrEmptyResource
+	return grantableRoles, nil
+}
+
+func (c *bigQueryClient) getGrantableRolesForTables(ctx context.Context) ([]string, error) {
+	sampleTable, err := c.getSampleTable(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting a sample table, %w", err)
 	}
 
+	resourceName := fmt.Sprintf("//bigquery.googleapis.com/projects/%v/datasets/%v/tables/%v", sampleTable.ProjectId, sampleTable.DatasetId, sampleTable.TableId)
+
+	var grantableRoles []string
 	request := &iam.QueryGrantableRolesRequest{
 		FullResourceName: resourceName,
 	}
-	response, err := c.iamService.Roles.QueryGrantableRoles(request).Do()
-	if err != nil {
+	if err := c.iamService.Roles.QueryGrantableRoles(request).Pages(ctx, func(page *iam.QueryGrantableRolesResponse) error {
+		for _, role := range page.Roles {
+			grantableRoles = append(grantableRoles, role.Name)
+		}
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	return grantableRoles, nil
+}
+
+func (c *bigQueryClient) getSampleDataset(ctx context.Context) (*bqApi.DatasetReference, error) {
+	var dataset *bqApi.DatasetReference
+	if err := c.apiClient.Datasets.List(c.projectID).Pages(ctx, func(page *bqApi.DatasetList) error {
+		for _, d := range page.Datasets {
+			dataset = d.DatasetReference
+			break
+		}
+		return nil
+	}); err != nil {
 		return nil, err
 	}
+	if dataset == nil {
+		return nil, fmt.Errorf("%w: dataset", ErrEmptyResource)
+	}
+
+	return dataset, nil
+}
 
-	var roles []string
-	for _, role := range response.Roles {
-		roles = append(roles, role.Name)
+func (c *bigQueryClient) getSampleTable(ctx context.Context) (*bqApi.TableReference, error) {
+	var table *bqApi.TableReference
+	if err := c.apiClient.Datasets.List(c.projectID).Pages(ctx, func(page *bqApi.DatasetList) error {
+		for _, d := range page.Datasets {
+			if err := c.apiClient.Tables.
+				List(c.projectID, d.DatasetReference.DatasetId).
+				Pages(ctx, func(page *bqApi.TableList) error {
+					for _, t := range page.Tables {
+						table = t.TableReference
+						break
+					}
+					return nil
+				}); err != nil {
+				return fmt.Errorf("getting a sample table, %w", err)
+			}
+			if table != nil {
+				break
+			}
+		}
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+	if table == nil {
+		return nil, fmt.Errorf("%w: table", ErrEmptyResource)
 	}
 
-	return roles, nil
+	return table, nil
 }
 
 func containsString(arr []string, v string) bool {

diff --git a/plugins/providers/bigquery/config.go b/plugins/providers/bigquery/config.go
@@ -1,6 +1,7 @@
 package bigquery
 
 import (
+	"context"
 	"encoding/base64"
 	"errors"
 	"fmt"
@@ -85,6 +86,9 @@ type Config struct {
 
 	crypto    domain.Crypto
 	validator *validator.Validate
+
+	cachedDatasetGrantableRoles []string
+	cachedTableGrantableRoles   []string
 }
 
 // NewConfig returns bigquery config struct
@@ -197,12 +201,23 @@ func (c *Config) validatePermission(value interface{}, resourceType string, clie
 		return nil, ErrInvalidPermissionConfig
 	}
 
+	ctx := context.TODO()
 	if resourceType == ResourceTypeDataset {
 		if !utils.ContainsString([]string{DatasetRoleReader, DatasetRoleWriter, DatasetRoleOwner}, permision) {
-			return nil, fmt.Errorf("%v: %v", ErrInvalidDatasetPermission, permision)
+			grantableRoles, err := c.getGrantableRolesForDataset(ctx, client)
+			if err != nil {
+				if errors.Is(err, ErrEmptyResource) {
+					return nil, fmt.Errorf("cannot verify dataset permission: %v", permision)
+				}
+				return nil, err
+			}
+
+			if !utils.ContainsString(grantableRoles, permision) {
+				return nil, fmt.Errorf("%v: %v", ErrInvalidDatasetPermission, permision)
+			}
 		}
 	} else if resourceType == ResourceTypeTable {
-		roles, err := client.getGrantableRolesForTables()
+		roles, err := c.getGrantableRolesForTables(ctx, client)
 		if err != nil {
 			if err == ErrEmptyResource {
 				return nil, ErrCannotVerifyTablePermission
@@ -220,3 +235,31 @@ func (c *Config) validatePermission(value interface{}, resourceType string, clie
 	configValue := Permission(permision)
 	return &configValue, nil
 }
+
+func (c *Config) getGrantableRolesForDataset(ctx context.Context, client *bigQueryClient) ([]string, error) {
+	if len(c.cachedDatasetGrantableRoles) > 0 {
+		return c.cachedDatasetGrantableRoles, nil
+	}
+
+	roles, err := client.getGrantableRolesForDataset(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	c.cachedDatasetGrantableRoles = roles
+	return roles, nil
+}
+
+func (c *Config) getGrantableRolesForTables(ctx context.Context, client *bigQueryClient) ([]string, error) {
+	if len(c.cachedTableGrantableRoles) > 0 {
+		return c.cachedTableGrantableRoles, nil
+	}
+
+	roles, err := client.getGrantableRolesForTables(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	c.cachedTableGrantableRoles = roles
+	return roles, nil
+}