github.com/ecodeclub/eorm@v0.0.2-0.20231001112437-dae71da914d0/internal/sharding/hash/shadow_hash.go (about)

     1  // Copyright 2021 ecodeclub
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  // http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package hash
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"strings"
    21  
    22  	"github.com/ecodeclub/eorm/internal/errs"
    23  	"github.com/ecodeclub/eorm/internal/sharding"
    24  )
    25  
    26  // ShadowHash TODO experiemntal
    27  type ShadowHash struct {
    28  	*Hash
    29  	Prefix string
    30  }
    31  
    32  func (h *ShadowHash) Broadcast(ctx context.Context) []sharding.Dst {
    33  	res := make([]sharding.Dst, 0, 8)
    34  	for i := 0; i < h.DBPattern.Base; i++ {
    35  		dbName := fmt.Sprintf(h.Prefix+h.DBPattern.Name, i)
    36  		for j := 0; j < h.TablePattern.Base; j++ {
    37  			res = append(res, sharding.Dst{
    38  				Name:  h.DsPattern.Name,
    39  				DB:    dbName,
    40  				Table: fmt.Sprintf(h.Prefix+h.TablePattern.Name, j),
    41  			})
    42  		}
    43  	}
    44  	return res
    45  }
    46  
    47  func (h *ShadowHash) Sharding(ctx context.Context, req sharding.Request) (sharding.Response, error) {
    48  	if h.ShardingKey == "" {
    49  		return sharding.EmptyResp, errs.ErrMissingShardingKey
    50  	}
    51  	skVal, ok := req.SkValues[h.ShardingKey]
    52  	if !ok {
    53  		return sharding.Response{Dsts: h.Broadcast(ctx)}, nil
    54  	}
    55  	dbName := h.DBPattern.Name
    56  	if !h.DBPattern.NotSharding && strings.Contains(dbName, "%d") {
    57  		dbName = fmt.Sprintf(dbName, skVal.(int)%h.DBPattern.Base)
    58  	}
    59  	tbName := h.TablePattern.Name
    60  	if !h.TablePattern.NotSharding && strings.Contains(tbName, "%d") {
    61  		tbName = fmt.Sprintf(tbName, skVal.(int)%h.TablePattern.Base)
    62  	}
    63  	dsName := h.DsPattern.Name
    64  	if !h.DsPattern.NotSharding && strings.Contains(dsName, "%d") {
    65  		dsName = fmt.Sprintf(dsName, skVal.(int)%h.DsPattern.Base)
    66  	}
    67  	if isSourceKey(ctx) {
    68  		dsName = h.Prefix + dsName
    69  	}
    70  	if isDBKey(ctx) {
    71  		dbName = h.Prefix + dbName
    72  	}
    73  
    74  	if isTableKey(ctx) {
    75  		tbName = h.Prefix + tbName
    76  	}
    77  
    78  	return sharding.Response{
    79  		Dsts: []sharding.Dst{{Name: dsName, DB: dbName, Table: tbName}},
    80  	}, nil
    81  }
    82  
    83  type sourceKey struct{}
    84  
    85  type dbKey struct{}
    86  
    87  type tableKey struct{}
    88  
    89  func CtxWithTableKey(ctx context.Context) context.Context {
    90  	return context.WithValue(ctx, tableKey{}, true)
    91  }
    92  
    93  func CtxWithDBKey(ctx context.Context) context.Context {
    94  	return context.WithValue(ctx, dbKey{}, true)
    95  }
    96  
    97  func CtxWithSourceKey(ctx context.Context) context.Context {
    98  	return context.WithValue(ctx, sourceKey{}, true)
    99  }
   100  
   101  func isSourceKey(ctx context.Context) bool {
   102  	return ctx.Value(sourceKey{}) != nil
   103  }
   104  
   105  func isDBKey(ctx context.Context) bool {
   106  	return ctx.Value(dbKey{}) != nil
   107  }
   108  
   109  func isTableKey(ctx context.Context) bool {
   110  	return ctx.Value(tableKey{}) != nil
   111  }