增加自己定义生成字符串

icowan · Mar 28, 2020 · 60651e5 · 60651e5
1 parent 9769e79
commit 60651e5
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/cmd/service/service.go b/cmd/service/service.go
@@ -46,6 +46,7 @@ var (
 	devCors       = fs.String("dev-cors", "false", "is develop")
 	rateBucketNum = fs.Int("rate-bucket", 10, "rate bucket num")
 	maxLength     = fs.Int("max-length", -1, "code length")
+	alphabet      = fs.String("alphabet", "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ", "alphabet length")
 	err           error
 )
 
@@ -64,6 +65,7 @@ func Run() {
 	logPath = envString("LOG_PATH", logPath)
 	logLevel = envString("LOG_LEVEL", logLevel)
 	devCors = envString("DEV_CORS", devCors)
+	alphabet = envString("ALPHABET", alphabet)
 	rateBucketNum = envInt("RATE_BUCKET", rateBucketNum)
 	maxLength = envInt("MAX_LENGTH", maxLength)
 
@@ -86,7 +88,7 @@ func Run() {
 		}
 	}
 
-	svc := service.New(getServiceMiddleware(logger), logger, repo, *shortUri, *maxLength)
+	svc := service.New(getServiceMiddleware(logger), logger, repo, *shortUri, *maxLength, *alphabet)
 	eps := endpoint.New(svc, getEndpointMiddleware(logger))
 	g := createService(eps)
 	initCancelInterrupt(g)

diff --git a/install/docker-compose/docker-compose.yaml b/install/docker-compose/docker-compose.yaml
@@ -18,6 +18,7 @@ services:
       MONGO_ADDR: "mongodb://root:admin@mongo:27017/?authSource=admin"
       SHORT_URI: "http://localhost:8080"
       LOG_LEVEL: "debug"
+      MAX_LENGTH: 6
     restart: always
     depends_on:
       - mongo

diff --git a/pkg/service/service.go b/pkg/service/service.go
@@ -31,21 +31,22 @@ type service struct {
 	logger     log.Logger
 	shortUri   string
 	maxLen     int
+	alphabet   string
 }
 
-func New(middleware []Middleware, logger log.Logger, repository Repository, shortUri string, maxLength int) Service {
-	var svc Service = NewService(logger, repository, shortUri, maxLength)
+func New(middleware []Middleware, logger log.Logger, repository Repository, shortUri string, maxLength int, alphabet string) Service {
+	var svc Service = NewService(logger, repository, shortUri, maxLength, alphabet)
 	for _, m := range middleware {
 		svc = m(svc)
 	}
 	return svc
 }
 
-func NewService(logger log.Logger, repository Repository, shortUri string, maxLength int) Service {
+func NewService(logger log.Logger, repository Repository, shortUri string, maxLength int, alphabet string) Service {
 	if maxLength > 9 {
 		maxLength = 9
 	}
-	return &service{repository: repository, shortUri: shortUri, logger: logger, maxLen: maxLength}
+	return &service{repository: repository, shortUri: shortUri, logger: logger, maxLen: maxLength, alphabet: alphabet}
 }
 
 func (s *service) Get(ctx context.Context, code string) (redirect *Redirect, err error) {
@@ -55,7 +56,7 @@ func (s *service) Get(ctx context.Context, code string) (redirect *Redirect, err
 func (s *service) Post(ctx context.Context, domain string) (redirect *Redirect, err error) {
 	now := time.Now()
 	now = now.In(time.Local)
-	code := shortid.MustGenerate()
+	code, _ := shortid.MustNew(0, s.alphabet, 1).Generate()
 	// todo 考虑如何处理垃圾数据的问题 得复的url 不同的code
 
 	if s.maxLen > 0 {