diff --git a/cmd/service/service.go b/cmd/service/service.go index e79024a..3ae2266 100644 --- 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 index db8b2d4..5f8405a 100644 --- 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 index a4cc0bb..e98826c 100644 --- 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 {