94 lines
2.3 KiB
Go
94 lines
2.3 KiB
Go
package tenant
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/influxdata/influxdb/v2"
|
|
"github.com/influxdata/influxdb/v2/kit/tracing"
|
|
"github.com/influxdata/influxdb/v2/kv"
|
|
"github.com/influxdata/influxdb/v2/rand"
|
|
"github.com/influxdata/influxdb/v2/snowflake"
|
|
)
|
|
|
|
const MaxIDGenerationN = 100
|
|
const ReservedIDs = 1000
|
|
|
|
type Store struct {
|
|
kvStore kv.Store
|
|
IDGen influxdb.IDGenerator
|
|
OrgBucketIDGen influxdb.IDGenerator
|
|
urmByUserIndex *kv.Index
|
|
}
|
|
|
|
func NewStore(kvStore kv.Store) *Store {
|
|
return &Store{
|
|
kvStore: kvStore,
|
|
IDGen: snowflake.NewDefaultIDGenerator(),
|
|
OrgBucketIDGen: rand.NewOrgBucketID(time.Now().UnixNano()),
|
|
urmByUserIndex: kv.NewIndex(kv.URMByUserIndexMapping, kv.WithIndexReadPathEnabled),
|
|
}
|
|
}
|
|
|
|
// View opens up a transaction that will not write to any data. Implementing interfaces
|
|
// should take care to ensure that all view transactions do not mutate any data.
|
|
func (s *Store) View(ctx context.Context, fn func(kv.Tx) error) error {
|
|
return s.kvStore.View(ctx, fn)
|
|
}
|
|
|
|
// Update opens up a transaction that will mutate data.
|
|
func (s *Store) Update(ctx context.Context, fn func(kv.Tx) error) error {
|
|
return s.kvStore.Update(ctx, fn)
|
|
}
|
|
|
|
// generateSafeID attempts to create ids for buckets
|
|
// and orgs that are without backslash, commas, and spaces, BUT ALSO do not already exist.
|
|
func (s *Store) generateSafeOrgBucketID(ctx context.Context, tx kv.Tx, bucket []byte) (influxdb.ID, error) {
|
|
for i := 0; i < MaxIDGenerationN; i++ {
|
|
id := s.OrgBucketIDGen.ID()
|
|
|
|
// TODO: this is probably unnecessary but for testing we need to keep it in.
|
|
// After KV is cleaned out we can update the tests and remove this.
|
|
if id < ReservedIDs {
|
|
continue
|
|
}
|
|
|
|
err := s.uniqueID(ctx, tx, bucket, id)
|
|
if err == nil {
|
|
return id, nil
|
|
}
|
|
|
|
if err == NotUniqueIDError {
|
|
continue
|
|
}
|
|
|
|
return influxdb.InvalidID(), err
|
|
}
|
|
return influxdb.InvalidID(), ErrFailureGeneratingID
|
|
}
|
|
|
|
func (s *Store) uniqueID(ctx context.Context, tx kv.Tx, bucket []byte, id influxdb.ID) error {
|
|
span, _ := tracing.StartSpanFromContext(ctx)
|
|
defer span.Finish()
|
|
|
|
encodedID, err := id.Encode()
|
|
if err != nil {
|
|
return &influxdb.Error{
|
|
Code: influxdb.EInvalid,
|
|
Err: err,
|
|
}
|
|
}
|
|
|
|
b, err := tx.Bucket(bucket)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = b.Get(encodedID)
|
|
if kv.IsNotFound(err) {
|
|
return nil
|
|
}
|
|
|
|
return NotUniqueIDError
|
|
}
|