Feature/sequence (arana-db#322)

* [ISSUE arana-db#94] Feat sequence (arana-db#211) * feat: support sequence * feat: none * feat: snowflake seqneuce finish * fix: add sequence unit test * feat: support sequence * chore: remove test.log * fix: fix cr issue * fix: fix integration_test fail * style: fix code style * docs: update code comment * fix: fix pr issue * fix: fix sql scriip style * test: add unit test * docs: fix sql script comment * refactor: update auto_increment to sequence * fix: fix cr issue * refactor: refactor sequence create mgr * style: fmt code style * fix: fix pr issue * docs: fix code comment * fix: fix cr issue * fix: fix sequence create table sql (arana-db#266) * fix: compatible with new optimizer structure * fix ut * fix: fix sequence bug (arana-db#302) * fix: fix sequence bug * refactor: use lazy creat sequence * style: fix code style * docs: update code comment * fix: fix pr issue * fix: remove unuse test file * Fix new sequence (arana-db#316) * fix: fix sequence bug * refactor: use lazy creat sequence * style: fix code style * docs: update code comment * fix: fix pr issue * fix: remove unuse test file * fix: fix snowflake all even number bug * Merge master (arana-db#320) * add unit test. (arana-db#306) * Add: add Range interface mock. * Add: add test case for iterator. * fix: format import. * fix: fix test case error. * Add: add test case for shard_expr. * Add: add license header. * Add: add new test case. * fix: fix import format. * refactor schemanager cache to lru cache (arana-db#246) Co-authored-by: Dong Jianhui <[email protected]> * optimize maxAllowedPacket config fix: arana-db#292 arana-db#281 (arana-db#297) * insert multi value at once request, throw error msg when any row error fix: arana-db#287 (arana-db#313) * feat: support show status (arana-db#309) * feat: max_allowed_packet is set by the user (arana-db#312) Co-authored-by: Zonglei Dong <[email protected]> Co-authored-by: Mulavar <[email protected]> Co-authored-by: Dong Jianhui <[email protected]> Co-authored-by: adair peng <[email protected]> * style: fix code style * ci: fix ci bug * ci: fix ci bug * ci: fix ci error * fix: fix ci bug * fix: fix runtime.Runtime use Co-authored-by: Jeffsky <[email protected]> Co-authored-by: Zonglei Dong <[email protected]> Co-authored-by: Mulavar <[email protected]> Co-authored-by: Dong Jianhui <[email protected]> Co-authored-by: adair peng <[email protected]>
IoCing · Jul 29, 2022 · e0d604b · e0d604b
1 parent b04366f
commit e0d604b
Show file tree

Hide file tree

Showing 28 changed files with 1,206 additions and 107 deletions.
diff --git a/conf/config.yaml b/conf/config.yaml
@@ -95,6 +95,9 @@ data:
     tables:
       - name: employees.student
         allow_full_scan: true
+        sequence:
+          type: snowflake
+          option:
         db_rules:
           - column: uid
             type: scriptExpr

diff --git a/integration_test/config/db/config.yaml b/integration_test/config/db/config.yaml
@@ -82,6 +82,9 @@ data:
     tables:
       - name: employees.student
         allow_full_scan: true
+        sequence:
+          type: snowflake
+          option:
         db_rules:
           - column: uid
             type: scriptExpr

diff --git a/integration_test/config/db_tbl/config.yaml b/integration_test/config/db_tbl/config.yaml
@@ -89,6 +89,9 @@ data:
     tables:
       - name: employees.student
         allow_full_scan: true
+        sequence:
+          type: snowflake
+          option:
         db_rules:
           - column: uid
             type: scriptExpr

diff --git a/integration_test/config/tbl/config.yaml b/integration_test/config/tbl/config.yaml
@@ -55,6 +55,9 @@ data:
     tables:
       - name: employees.student
         allow_full_scan: true
+        sequence:
+          type: snowflake
+          option:
         db_rules:
           - column: uid
             type: scriptExpr

diff --git a/pkg/boot/boot.go b/pkg/boot/boot.go
@@ -29,6 +29,7 @@ import (
 	"github.com/arana-db/arana/pkg/config"
 	"github.com/arana-db/arana/pkg/proto/rule"
 	"github.com/arana-db/arana/pkg/runtime"
+	rcontext "github.com/arana-db/arana/pkg/runtime/context"
 	"github.com/arana-db/arana/pkg/runtime/namespace"
 	_ "github.com/arana-db/arana/pkg/schema"
 	"github.com/arana-db/arana/pkg/security"
@@ -54,6 +55,9 @@ func Boot(ctx context.Context, provider Discovery) error {
 		if c, err = provider.GetCluster(ctx, cluster); err != nil {
 			continue
 		}
+
+		ctx = rcontext.WithTenant(ctx, c.Tenant)
+
 		if ns, err = buildNamespace(ctx, provider, cluster); err != nil {
 			log.Errorf("build namespace %s failed: %v", cluster, err)
 			continue
@@ -143,5 +147,5 @@ func buildNamespace(ctx context.Context, provider Discovery, clusterName string)
 	}
 	initCmds = append(initCmds, namespace.UpdateRule(&ru))
 
-	return namespace.New(clusterName, initCmds...), nil
+	return namespace.New(clusterName, initCmds...)
 }
diff --git a/pkg/boot/discovery.go b/pkg/boot/discovery.go
@@ -308,9 +308,9 @@ func (fp *discovery) GetTable(ctx context.Context, cluster, tableName string) (*
 	if !ok {
 		return nil, nil
 	}
-	var vt rule.VTable
 
 	var (
+		vt                 rule.VTable
 		topology           rule.Topology
 		dbFormat, tbFormat string
 		dbBegin, tbBegin   int
@@ -436,11 +436,18 @@ func (fp *discovery) GetTable(ctx context.Context, cluster, tableName string) (*
 	if table.AllowFullScan {
 		vt.SetAllowFullScan(true)
 	}
+	if table.Sequence != nil {
+		vt.SetAutoIncrement(&rule.AutoIncrement{
+			Type:   table.Sequence.Type,
+			Option: table.Sequence.Option,
+		})
+	}
 
 	// TODO: process attributes
 	_ = table.Attributes["sql_max_limit"]
 
 	vt.SetTopology(&topology)
+	vt.SetName(tableName)
 
 	return &vt, nil
 }

diff --git a/pkg/boot/plugin.go b/pkg/boot/plugin.go
@@ -21,4 +21,7 @@ import (
 	_ "github.com/arana-db/arana/pkg/config/etcd"
 	_ "github.com/arana-db/arana/pkg/config/file"
 	_ "github.com/arana-db/arana/pkg/config/nacos"
+	_ "github.com/arana-db/arana/pkg/sequence"
+	_ "github.com/arana-db/arana/pkg/sequence/group"
+	_ "github.com/arana-db/arana/pkg/sequence/snowflake"
 )
diff --git a/pkg/config/model.go b/pkg/config/model.go
@@ -123,6 +123,7 @@ type (
 
 	Table struct {
 		Name           string            `validate:"required" yaml:"name" json:"name"`
+		Sequence       *Sequence         `yaml:"sequence" json:"sequence"`
 		AllowFullScan  bool              `yaml:"allow_full_scan" json:"allow_full_scan,omitempty"`
 		DbRules        []*Rule           `yaml:"db_rules" json:"db_rules"`
 		TblRules       []*Rule           `yaml:"tbl_rules" json:"tbl_rules"`
@@ -131,6 +132,11 @@ type (
 		Attributes     map[string]string `yaml:"attributes" json:"attributes"`
 	}
 
+	Sequence struct {
+		Type   string            `yaml:"type" json:"type"`
+		Option map[string]string `yaml:"option" json:"option"`
+	}
+
 	Rule struct {
 		Column string `validate:"required" yaml:"column" json:"column"`
 		Type   string `validate:"required" yaml:"type" json:"type"`

diff --git a/pkg/proto/rule/rule.go b/pkg/proto/rule/rule.go
@@ -55,8 +55,22 @@ const (
 // VTable represents a virtual/logical table.
 type VTable struct {
 	attributes
-	topology *Topology
-	shards   map[string][2]*ShardMetadata // column -> [db shard metadata,table shard metadata]
+	name          string // TODO: set name
+	autoIncrement *AutoIncrement
+	topology      *Topology
+	shards        map[string][2]*ShardMetadata // column -> [db shard metadata,table shard metadata]
+}
+
+func (vt *VTable) Name() string {
+	return vt.name
+}
+
+func (vt *VTable) GetAutoIncrement() *AutoIncrement {
+	return vt.autoIncrement
+}
+
+func (vt *VTable) SetAutoIncrement(seq *AutoIncrement) {
+	vt.autoIncrement = seq
 }
 
 func (vt *VTable) SetAllowFullScan(allow bool) {
@@ -136,6 +150,10 @@ func (vt *VTable) SetTopology(topology *Topology) {
 	vt.topology = topology
 }
 
+func (vt *VTable) SetName(name string) {
+	vt.name = name
+}
+
 // Rule represents sharding rule, a Rule contains multiple logical tables.
 type Rule struct {
 	mu    sync.RWMutex
@@ -203,3 +221,15 @@ func (ru *Rule) MustVTable(name string) *VTable {
 	}
 	return v
 }
+
+// Range ranges each VTable
+func (ru *Rule) Range(f func(table string, vt *VTable) bool) {
+	ru.mu.RLock()
+	defer ru.mu.RUnlock()
+
+	for k, v := range ru.vtabs {
+		if !f(k, v) {
+			break
+		}
+	}
+}
diff --git a/pkg/proto/rule/sequence.go b/pkg/proto/rule/sequence.go
@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package rule
+
+type AutoIncrement struct {
+	Type   string
+	Option map[string]string
+}
diff --git a/pkg/proto/sequence.go b/pkg/proto/sequence.go
@@ -15,21 +15,103 @@
  * limitations under the License.
  */
 
-//go:generate mockgen -destination=../../testdata/mock_sequence.go -package=testdata . Sequence,Sequencer
+//go:generate mockgen -destination=../../testdata/mock_sequence.go -package=testdata . Sequence,SequenceManager
 package proto
 
 import (
 	"context"
+	"errors"
+	"fmt"
 )
 
-// Sequence represents a global unique id generator.
-type Sequence interface {
-	// Next generates a next value in int64.
-	Next(ctx context.Context) (int64, error)
+var (
+	ErrorNotSequenceType  = errors.New("sequence type not found")
+	ErrorNotFoundSequence = errors.New("sequence instance not found")
+
+	_defaultSequenceManager SequenceManager
+)
+
+func RegisterSequenceManager(l SequenceManager) {
+	_defaultSequenceManager = l
+}
+
+func LoadSequenceManager() SequenceManager {
+	cur := _defaultSequenceManager
+	if cur == nil {
+		return noopSequenceManager{}
+	}
+	return cur
+}
+
+func BuildAutoIncrementName(table string) string {
+	return fmt.Sprintf("__arana_incr_%s", table)
+}
+
+type (
+	RuntimeCtxKey = struct{}
+
+	// SequenceSupplier Create the creator of Sequence
+	SequenceSupplier func() EnhancedSequence
+
+	SequenceConfig struct {
+		Name   string
+		Type   string
+		Option map[string]string
+	}
+
+	// Sequence represents a global unique id generator.
+	Sequence interface {
+		// Acquire generates a next value in int64.
+		Acquire(ctx context.Context) (int64, error)
+		Reset() error
+		Update() error
+	}
+
+	// EnhancedSequence represents a global unique id generator.
+	EnhancedSequence interface {
+		Sequence
+		// Start start sequence instance.
+		Start(ctx context.Context, option SequenceConfig) error
+		// CurrentVal get sequence current id.
+		CurrentVal() int64
+		// Stop stops sequence.
+		Stop() error
+	}
+
+	// SequenceManager represents the factory to create a Sequence by table name.
+	SequenceManager interface {
+		// CreateSequence creates one sequence instance
+		CreateSequence(ctx context.Context, tenant, schema string, opt SequenceConfig) (Sequence, error)
+		// GetSequence gets sequence instance by name
+		GetSequence(ctx context.Context, tenant, schema, name string) (Sequence, error)
+	}
+)
+
+var (
+	// Record the list of Sequence plug -in through the registered self -registered
+	suppliersRegistry = make(map[string]SequenceSupplier)
+)
+
+// RegisterSequence Register a Sequence plugin
+func RegisterSequence(name string, supplier SequenceSupplier) {
+	suppliersRegistry[name] = supplier
+}
+
+// GetSequenceSupplier returns SequenceSupplier.
+func GetSequenceSupplier(name string) (SequenceSupplier, bool) {
+	val, ok := suppliersRegistry[name]
+	return val, ok
+}
+
+type noopSequenceManager struct {
+}
+
+// CreateSequence creates one sequence instance
+func (n noopSequenceManager) CreateSequence(ctx context.Context, tenant, schema string, opt SequenceConfig) (Sequence, error) {
+	return nil, nil
 }
 
-// Sequencer represents the factory to create a Sequence by table name.
-type Sequencer interface {
-	// Sequence returns the Sequence of table.
-	Sequence(ctx context.Context, table string) (Sequence, error)
+// GetSequence gets sequence instance by name
+func (n noopSequenceManager) GetSequence(ctx context.Context, tenant, schema, name string) (Sequence, error) {
+	return nil, nil
 }
diff --git a/pkg/runtime/context/context.go b/pkg/runtime/context/context.go
@@ -70,8 +70,8 @@ func WithSchema(ctx context.Context, data string) context.Context {
 	return context.WithValue(ctx, keySchema{}, data)
 }
 
-// WithSequencer binds a sequencer.
-func WithSequencer(ctx context.Context, sequencer proto.Sequencer) context.Context {
+// WithSequenceManager binds a sequencer.
+func WithSequenceManager(ctx context.Context, sequencer proto.SequenceManager) context.Context {
 	return context.WithValue(ctx, keySequence{}, sequencer)
 }
 
@@ -91,8 +91,8 @@ func WithHints(ctx context.Context, hints []*hint.Hint) context.Context {
 }
 
 // Sequencer extracts the sequencer.
-func Sequencer(ctx context.Context) proto.Sequencer {
-	s, ok := ctx.Value(keySequence{}).(proto.Sequencer)
+func Sequencer(ctx context.Context) proto.SequenceManager {
+	s, ok := ctx.Value(keySequence{}).(proto.SequenceManager)
 	if !ok {
 		return nil
 	}