github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/store/sqlstore/tokens_store.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package sqlstore
     5  
     6  import (
     7  	"database/sql"
     8  	"fmt"
     9  
    10  	sq "github.com/Masterminds/squirrel"
    11  	"github.com/pkg/errors"
    12  
    13  	"github.com/masterhung0112/hk_server/v5/model"
    14  	"github.com/masterhung0112/hk_server/v5/shared/mlog"
    15  	"github.com/masterhung0112/hk_server/v5/store"
    16  )
    17  
    18  type SqlTokenStore struct {
    19  	*SqlStore
    20  }
    21  
    22  func newSqlTokenStore(sqlStore *SqlStore) store.TokenStore {
    23  	s := &SqlTokenStore{sqlStore}
    24  
    25  	for _, db := range sqlStore.GetAllConns() {
    26  		table := db.AddTableWithName(model.Token{}, "Tokens").SetKeys(false, "Token")
    27  		table.ColMap("Token").SetMaxSize(64)
    28  		table.ColMap("Type").SetMaxSize(64)
    29  		table.ColMap("Extra").SetMaxSize(2048)
    30  	}
    31  
    32  	return s
    33  }
    34  
    35  func (s SqlTokenStore) createIndexesIfNotExists() {
    36  }
    37  
    38  func (s SqlTokenStore) Save(token *model.Token) error {
    39  	if err := token.IsValid(); err != nil {
    40  		return err
    41  	}
    42  
    43  	if err := s.GetMaster().Insert(token); err != nil {
    44  		return errors.Wrap(err, "failed to save Token")
    45  	}
    46  	return nil
    47  }
    48  
    49  func (s SqlTokenStore) Delete(token string) error {
    50  	if _, err := s.GetMaster().Exec("DELETE FROM Tokens WHERE Token = :Token", map[string]interface{}{"Token": token}); err != nil {
    51  		return errors.Wrapf(err, "failed to delete Token with value %s", token)
    52  	}
    53  	return nil
    54  }
    55  
    56  func (s SqlTokenStore) GetByToken(tokenString string) (*model.Token, error) {
    57  	token := &model.Token{}
    58  
    59  	if err := s.GetReplica().SelectOne(token, "SELECT * FROM Tokens WHERE Token = :Token", map[string]interface{}{"Token": tokenString}); err != nil {
    60  		if err == sql.ErrNoRows {
    61  			return nil, store.NewErrNotFound("Token", fmt.Sprintf("Token=%s", tokenString))
    62  		}
    63  
    64  		return nil, errors.Wrapf(err, "failed to get Token with value %s", tokenString)
    65  	}
    66  
    67  	return token, nil
    68  }
    69  
    70  func (s SqlTokenStore) Cleanup() {
    71  	mlog.Debug("Cleaning up token store.")
    72  	deltime := model.GetMillis() - model.MAX_TOKEN_EXPIRY_TIME
    73  	if _, err := s.GetMaster().Exec("DELETE FROM Tokens WHERE CreateAt < :DelTime", map[string]interface{}{"DelTime": deltime}); err != nil {
    74  		mlog.Error("Unable to cleanup token store.")
    75  	}
    76  }
    77  
    78  func (s SqlTokenStore) GetAllTokensByType(tokenType string) ([]*model.Token, error) {
    79  	tokens := []*model.Token{}
    80  	query, args, err := s.getQueryBuilder().Select("*").From("Tokens").Where(sq.Eq{"Type": tokenType}).ToSql()
    81  	if err != nil {
    82  		return nil, errors.Wrap(err, "could not build sql query to get all tokens by type")
    83  	}
    84  
    85  	if _, err := s.GetReplica().Select(&tokens, query, args...); err != nil {
    86  		return nil, errors.Wrapf(err, "failed to get all tokens of Type=%s", tokenType)
    87  	}
    88  	return tokens, nil
    89  }
    90  
    91  func (s SqlTokenStore) RemoveAllTokensByType(tokenType string) error {
    92  	if _, err := s.GetMaster().Exec("DELETE FROM Tokens WHERE Type = :TokenType", map[string]interface{}{"TokenType": tokenType}); err != nil {
    93  		return errors.Wrapf(err, "failed to remove all Tokens with Type=%s", tokenType)
    94  	}
    95  	return nil
    96  }