refactor: refact topology to support sharding db (#234)

* refactor: refact topology to support sharding db

* Update pkg/topo/topo.go

Co-authored-by: hehe.bu <hehebu33@gmail.com>

Co-authored-by: hehe.bu <hehebu33@gmail.com>

pkg/misc/strings.go

@@ -17,10 +17,19 @@
 package misc
 
 import (
+	"regexp"
 	"strings"
 	"unicode"
 )
 
+const (
+	Numeric = "^[0-9]+$"
+)
+
+var (
+	rxNumeric = regexp.MustCompile(Numeric)
+)
+
 func FirstNonEmptyString(first string, second string, others ...string) string {
 	if len(first) > 0 {
 		return first
@@ -148,3 +157,11 @@ func IsBlank(s string) bool {
 		return !unicode.IsSpace(r)
 	}) == -1
 }
+
+// IsNumeric checks if the string contains only numbers. Empty string is valid.
+func IsNumeric(str string) bool {
+	if len(str) == 0 {
+		return true
+	}
+	return rxNumeric.MatchString(str)
+}

pkg/misc/strings_test.go

@@ -441,3 +441,36 @@ func TestIsBlank(t *testing.T) {
 		})
 	}
 }
+
+func TestIsNumeric(t *testing.T) {
+	cases := map[string]struct {
+		input    string
+		expected bool
+	}{
+		"success-1": {
+			input:    "",
+			expected: true,
+		},
+		"success-2": {
+			input:    "1",
+			expected: true,
+		},
+		"false-3": {
+			input:    "a",
+			expected: false,
+		},
+		"false-4": {
+			input:    "1-9",
+			expected: false,
+		},
+	}
+
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			actual := IsNumeric(tc.input)
+			if actual != tc.expected {
+				t.Errorf("expected %t, got %t", tc.expected, actual)
+			}
+		})
+	}
+}

pkg/topo/topo.go

@@ -23,6 +23,8 @@ import (
 
 	"github.com/cznic/mathutil"
 	"github.com/pkg/errors"
+
+	"github.com/cectc/dbpack/pkg/misc"
 )
 
 const (
@@ -55,31 +57,48 @@ func ParseTopology(dbName, tableName string, topology map[int]string) (*Topology
 
 	for i := 0; i < dbLen; i++ {
 		tp := topology[i]
-		realDB := fmt.Sprintf("%s_%d", dbName, i)
-		params := topologyRegexp.FindStringSubmatch(tp)
-		if len(params) != 3 {
+		if tp == "" {
 			return nil, errors.Errorf("incorrect topology format")
 		}
-		begin, err := strconv.Atoi(params[1])
-		if err != nil {
-			return nil, err
-		}
-		end, err := strconv.Atoi(params[2])
-		if err != nil {
-			return nil, err
-		}
-		if begin >= end {
-			return nil, errors.Errorf("incorrect topology, begin index must less than end index")
-		}
+		realDB := fmt.Sprintf("%s_%d", dbName, i)
 		tableSlice := make([]string, 0)
-		for j := begin; j <= end; j++ {
-			index := j
+
+		if misc.IsNumeric(tp) {
+			index, err := strconv.Atoi(tp)
+			if err != nil {
+				return nil, err
+			}
 			realTable := fmt.Sprintf("%s_%d", tableName, index)
 			tables[realTable] = realDB
 			tableIndexMap[index] = realTable
 			tableIndexSlice = append(tableIndexSlice, index)
 			tableSlice = append(tableSlice, realTable)
 			max = mathutil.Max(max, index)
+		} else {
+			params := topologyRegexp.FindStringSubmatch(tp)
+			if len(params) != 3 {
+				return nil, errors.Errorf("incorrect topology format")
+			}
+			begin, err := strconv.Atoi(params[1])
+			if err != nil {
+				return nil, err
+			}
+			end, err := strconv.Atoi(params[2])
+			if err != nil {
+				return nil, err
+			}
+			if begin >= end {
+				return nil, errors.Errorf("incorrect topology, begin index must be less than end index")
+			}
+			for j := begin; j <= end; j++ {
+				index := j
+				realTable := fmt.Sprintf("%s_%d", tableName, index)
+				tables[realTable] = realDB
+				tableIndexMap[index] = realTable
+				tableIndexSlice = append(tableIndexSlice, index)
+				tableSlice = append(tableSlice, realTable)
+				max = mathutil.Max(max, index)
+			}
 		}
 		dbs[realDB] = tableSlice
 	}

pkg/topo/topo_test.go

@@ -38,4 +38,14 @@ func Test_ParseTopology(t *testing.T) {
 
 	assert.Nil(t, err)
 	assert.Equal(t, 100, tp.TableSliceLen)
+
+	tp, err = ParseTopology("school", "student", map[int]string{
+		0: "0",
+		1: "1",
+		2: "2",
+		3: "3",
+	})
+
+	assert.Nil(t, err)
+	assert.Equal(t, 4, tp.TableSliceLen)
 }