feat(kv): Add kv.Service org & bucket ID validation hooks.

Adds the ability to customize validation of organization and bucket
IDs in the `kv.Service`.

kv/bucket.go

@@ -3,6 +3,7 @@ package kv
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strings"
 	"time"
@@ -402,7 +403,7 @@ func (s *Service) CreateBucket(ctx context.Context, b *influxdb.Bucket) error {
 	})
 }
 
-func (s *Service) createBucket(ctx context.Context, tx Tx, b *influxdb.Bucket) error {
+func (s *Service) createBucket(ctx context.Context, tx Tx, b *influxdb.Bucket) (err error) {
 	if b.OrgID.Valid() {
 		span, ctx := tracing.StartSpanFromContext(ctx)
 		defer span.Finish()
@@ -423,7 +424,9 @@ func (s *Service) createBucket(ctx context.Context, tx Tx, b *influxdb.Bucket) e
 	if b.IsSystem() {
 		b.ID = influxdb.ID(b.Type)
 	} else {
-		b.ID = s.IDGenerator.ID()
+		if b.ID, err = s.generateBucketID(); err != nil {
+			return err
+		}
 	}
 
 	b.CreatedAt = s.Now()
@@ -445,6 +448,16 @@ func (s *Service) createBucket(ctx context.Context, tx Tx, b *influxdb.Bucket) e
 	return nil
 }
 
+func (s *Service) generateBucketID() (influxdb.ID, error) {
+	for i := 0; i < MaxIDGenerationN; i++ {
+		id := s.IDGenerator.ID()
+		if s.IsValidOrgBucketID == nil || s.IsValidOrgBucketID(id) {
+			return id, nil
+		}
+	}
+	return 0, errors.New("unable to generate valid bucket id")
+}
+
 // PutBucket will put a bucket without setting an ID.
 func (s *Service) PutBucket(ctx context.Context, b *influxdb.Bucket) error {
 	return s.kv.Update(ctx, func(tx Tx) error {

kv/bucket_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/influxdata/influxdb"
+	"github.com/influxdata/influxdb/inmem"
 	"github.com/influxdata/influxdb/kv"
 	influxdbtesting "github.com/influxdata/influxdb/testing"
 )
@@ -78,3 +79,16 @@ func initBucketService(s kv.Store, f influxdbtesting.BucketFields, t *testing.T)
 		}
 	}
 }
+
+func TestService_CreateBucket(t *testing.T) {
+	t.Run("InvalidBucketID", func(t *testing.T) {
+		svc := kv.NewService(inmem.NewKVStore())
+		if err := svc.PutOrganization(context.Background(), &influxdb.Organization{ID: 123, Name: "ORG"}); err != nil {
+			t.Fatal(err)
+		}
+		svc.IsValidOrgBucketID = func(id influxdb.ID) bool { return false }
+		if err := svc.CreateBucket(context.Background(), &influxdb.Bucket{OrgID: 123, Name: "BUCKET"}); err == nil || err.Error() != `unable to generate valid bucket id` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+}

kv/org.go

@@ -3,6 +3,7 @@ package kv
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strings"
 	"time"
@@ -14,6 +15,9 @@ import (
 	"github.com/influxdata/influxdb/kit/tracing"
 )
 
+// MaxIDGenerationN is the maximum number of times an ID generation is done before failing.
+const MaxIDGenerationN = 100
+
 var (
 	organizationBucket = []byte("organizationsv1")
 	organizationIndex  = []byte("organizationindexv1")
@@ -262,12 +266,14 @@ func (s *Service) addOrgOwner(ctx context.Context, tx Tx, orgID influxdb.ID) err
 	return s.addResourceOwner(ctx, tx, influxdb.OrgsResourceType, orgID)
 }
 
-func (s *Service) createOrganization(ctx context.Context, tx Tx, o *influxdb.Organization) error {
+func (s *Service) createOrganization(ctx context.Context, tx Tx, o *influxdb.Organization) (err error) {
 	if err := s.validOrganizationName(ctx, tx, o); err != nil {
 		return err
 	}
 
-	o.ID = s.IDGenerator.ID()
+	if o.ID, err = s.generateOrgID(); err != nil {
+		return err
+	}
 	o.CreatedAt = s.Now()
 	o.UpdatedAt = s.Now()
 	if err := s.appendOrganizationEventToLog(ctx, tx, o.ID, organizationCreatedEvent); err != nil {
@@ -291,6 +297,16 @@ func (s *Service) createOrganization(ctx context.Context, tx Tx, o *influxdb.Org
 	return nil
 }
 
+func (s *Service) generateOrgID() (influxdb.ID, error) {
+	for i := 0; i < MaxIDGenerationN; i++ {
+		id := s.IDGenerator.ID()
+		if s.IsValidOrgBucketID == nil || s.IsValidOrgBucketID(id) {
+			return id, nil
+		}
+	}
+	return 0, errors.New("unable to generate valid org id")
+}
+
 // PutOrganization will put a organization without setting an ID.
 func (s *Service) PutOrganization(ctx context.Context, o *influxdb.Organization) error {
 	var err error

kv/org_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/influxdata/influxdb"
+	"github.com/influxdata/influxdb/inmem"
 	"github.com/influxdata/influxdb/kv"
 	influxdbtesting "github.com/influxdata/influxdb/testing"
 )
@@ -70,3 +71,13 @@ func initOrganizationService(s kv.Store, f influxdbtesting.OrganizationFields, t
 		}
 	}
 }
+
+func TestService_CreateOrganization(t *testing.T) {
+	t.Run("InvalidOrgID", func(t *testing.T) {
+		svc := kv.NewService(inmem.NewKVStore())
+		svc.IsValidOrgBucketID = func(id influxdb.ID) bool { return false }
+		if err := svc.CreateOrganization(context.Background(), &influxdb.Organization{Name: "ORG"}); err == nil || err.Error() != `unable to generate valid org id` {
+			t.Fatalf("unexpected error: %#v", err)
+		}
+	})
+}

kv/service.go

@@ -28,6 +28,9 @@ type Service struct {
 	TokenGenerator influxdb.TokenGenerator
 	influxdb.TimeGenerator
 	Hash Crypt
+
+	// Organization & bucket specific validation.
+	IsValidOrgBucketID func(influxdb.ID) bool
 }
 
 // NewService returns an instance of a Service.