github.com/ecodeclub/eorm@v0.0.2-0.20231001112437-dae71da914d0/internal/sharding/hash/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  
    21  	"github.com/ecodeclub/eorm/internal/errs"
    22  	operator "github.com/ecodeclub/eorm/internal/operator"
    23  	"github.com/ecodeclub/eorm/internal/sharding"
    24  )
    25  
    26  // Hash TODO experiemntal
    27  type Hash struct {
    28  	// Base         int
    29  	ShardingKey  string
    30  	DBPattern    *Pattern
    31  	TablePattern *Pattern
    32  	// Datasource Pattern
    33  	DsPattern *Pattern
    34  }
    35  
    36  func (h *Hash) Broadcast(ctx context.Context) []sharding.Dst {
    37  	if !h.DBPattern.NotSharding && h.TablePattern.NotSharding && h.DsPattern.NotSharding { // 只分库
    38  		return h.onlyDBroadcast(ctx)
    39  	} else if h.DBPattern.NotSharding && !h.TablePattern.NotSharding && h.DsPattern.NotSharding { // 只分表
    40  		return h.onlyTableBroadcast(ctx)
    41  	} else if h.DBPattern.NotSharding && h.TablePattern.NotSharding && !h.DsPattern.NotSharding { // 只分集群
    42  		return h.onlyDataSourceBroadcast(ctx)
    43  	} else if !h.DBPattern.NotSharding && !h.TablePattern.NotSharding && !h.DsPattern.NotSharding { // 分集群分库分表
    44  		return h.allBroadcast(ctx)
    45  	}
    46  	// 分库分表
    47  	return h.defaultBroadcast(ctx)
    48  }
    49  
    50  func (h *Hash) defaultBroadcast(ctx context.Context) []sharding.Dst {
    51  	res := make([]sharding.Dst, 0, 8)
    52  	for i := 0; i < h.DBPattern.Base; i++ {
    53  		dbName := fmt.Sprintf(h.DBPattern.Name, i)
    54  		for j := 0; j < h.TablePattern.Base; j++ {
    55  			res = append(res, sharding.Dst{
    56  				Name:  h.DsPattern.Name,
    57  				DB:    dbName,
    58  				Table: fmt.Sprintf(h.TablePattern.Name, j),
    59  			})
    60  		}
    61  	}
    62  	return res
    63  }
    64  
    65  func (h *Hash) allBroadcast(ctx context.Context) []sharding.Dst {
    66  	res := make([]sharding.Dst, 0, 8)
    67  	for s := 0; s < h.DsPattern.Base; s++ {
    68  		dsName := fmt.Sprintf(h.DsPattern.Name, s)
    69  		for i := 0; i < h.DBPattern.Base; i++ {
    70  			dbName := fmt.Sprintf(h.DBPattern.Name, i)
    71  			for j := 0; j < h.TablePattern.Base; j++ {
    72  				res = append(res, sharding.Dst{
    73  					Name: dsName, DB: dbName,
    74  					Table: fmt.Sprintf(h.TablePattern.Name, j),
    75  				})
    76  			}
    77  		}
    78  	}
    79  	return res
    80  }
    81  
    82  func (h *Hash) onlyDBroadcast(ctx context.Context) []sharding.Dst {
    83  	res := make([]sharding.Dst, 0, 8)
    84  	for i := 0; i < h.DBPattern.Base; i++ {
    85  		res = append(res, sharding.Dst{
    86  			Name:  h.DsPattern.Name,
    87  			Table: h.TablePattern.Name,
    88  			DB:    fmt.Sprintf(h.DBPattern.Name, i),
    89  		})
    90  	}
    91  	return res
    92  }
    93  
    94  func (h *Hash) onlyTableBroadcast(ctx context.Context) []sharding.Dst {
    95  	res := make([]sharding.Dst, 0, 8)
    96  	for j := 0; j < h.TablePattern.Base; j++ {
    97  		res = append(res, sharding.Dst{
    98  			Name:  h.DsPattern.Name,
    99  			DB:    h.DBPattern.Name,
   100  			Table: fmt.Sprintf(h.TablePattern.Name, j),
   101  		})
   102  	}
   103  	return res
   104  }
   105  
   106  func (h *Hash) onlyDataSourceBroadcast(ctx context.Context) []sharding.Dst {
   107  	res := make([]sharding.Dst, 0, 8)
   108  	for j := 0; j < h.DsPattern.Base; j++ {
   109  		res = append(res, sharding.Dst{
   110  			Name:  fmt.Sprintf(h.DsPattern.Name, j),
   111  			DB:    h.DBPattern.Name,
   112  			Table: h.TablePattern.Name,
   113  		})
   114  	}
   115  	return res
   116  }
   117  
   118  func (h *Hash) Sharding(ctx context.Context, req sharding.Request) (sharding.Response, error) {
   119  	if h.ShardingKey == "" {
   120  		return sharding.EmptyResp, errs.ErrMissingShardingKey
   121  	}
   122  	skVal, ok := req.SkValues[h.ShardingKey]
   123  	if !ok {
   124  		return sharding.Response{Dsts: h.Broadcast(ctx)}, nil
   125  	}
   126  	switch req.Op {
   127  	case operator.OpEQ:
   128  		dbName := h.DBPattern.Name
   129  		if !h.DBPattern.NotSharding {
   130  			dbName = fmt.Sprintf(dbName, skVal.(int)%h.DBPattern.Base)
   131  		}
   132  		tbName := h.TablePattern.Name
   133  		if !h.TablePattern.NotSharding {
   134  			tbName = fmt.Sprintf(tbName, skVal.(int)%h.TablePattern.Base)
   135  		}
   136  		dsName := h.DsPattern.Name
   137  		if !h.DsPattern.NotSharding {
   138  			dsName = fmt.Sprintf(dsName, skVal.(int)%h.DsPattern.Base)
   139  		}
   140  		return sharding.Response{
   141  			Dsts: []sharding.Dst{{Name: dsName, DB: dbName, Table: tbName}},
   142  		}, nil
   143  	case operator.OpGT, operator.OpLT, operator.OpGTEQ,
   144  		operator.OpLTEQ, operator.OpNEQ, operator.OpNotIN:
   145  		return sharding.Response{Dsts: h.Broadcast(ctx)}, nil
   146  	default:
   147  		return sharding.EmptyResp, errs.NewUnsupportedOperatorError(req.Op.Text)
   148  	}
   149  }
   150  
   151  func (h *Hash) ShardingKeys() []string {
   152  	return []string{h.ShardingKey}
   153  }
   154  
   155  type Pattern struct {
   156  	Base        int
   157  	Name        string
   158  	NotSharding bool
   159  }