vitess.io/vitess@v0.16.2/go/vt/vtgate/vindexes/region_json.go (about)

     1  /*
     2  Copyright 2020 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package vindexes
    18  
    19  import (
    20  	"bytes"
    21  	"context"
    22  	"encoding/binary"
    23  	"encoding/json"
    24  	"fmt"
    25  	"os"
    26  	"strconv"
    27  
    28  	"vitess.io/vitess/go/vt/vtgate/evalengine"
    29  
    30  	"vitess.io/vitess/go/sqltypes"
    31  	"vitess.io/vitess/go/vt/key"
    32  	"vitess.io/vitess/go/vt/log"
    33  )
    34  
    35  var (
    36  	_ MultiColumn = (*RegionJSON)(nil)
    37  )
    38  
    39  func init() {
    40  	Register("region_json", NewRegionJSON)
    41  }
    42  
    43  // RegionMap is used to store mapping of country to region
    44  type RegionMap map[string]uint64
    45  
    46  // RegionJSON is a multi-column unique vindex
    47  // The first column is used to lookup the prefix part of the keyspace id, the second column is hashed,
    48  // and the two values are combined to produce the keyspace id.
    49  // RegionJson can be used for geo-partitioning because the first column can denote a region,
    50  // and it will dictate the shard range for that region.
    51  type RegionJSON struct {
    52  	name        string
    53  	regionMap   RegionMap
    54  	regionBytes int
    55  }
    56  
    57  // NewRegionJSON creates a RegionJson vindex.
    58  // The supplied map requires all the fields of "RegionExperimental".
    59  // Additionally, it requires a region_map argument representing the path to a json file
    60  // containing a map of country to region.
    61  func NewRegionJSON(name string, m map[string]string) (Vindex, error) {
    62  	rmPath := m["region_map"]
    63  	rmap := make(map[string]uint64)
    64  	data, err := os.ReadFile(rmPath)
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  	log.Infof("Loaded Region map from: %s", rmPath)
    69  	err = json.Unmarshal(data, &rmap)
    70  	if err != nil {
    71  		return nil, err
    72  	}
    73  	rb, err := strconv.Atoi(m["region_bytes"])
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	switch rb {
    78  	case 1, 2:
    79  	default:
    80  		return nil, fmt.Errorf("region_bytes must be 1 or 2: %v", rb)
    81  	}
    82  
    83  	return &RegionJSON{
    84  		name:        name,
    85  		regionMap:   rmap,
    86  		regionBytes: rb,
    87  	}, nil
    88  }
    89  
    90  // String returns the name of the vindex.
    91  func (rv *RegionJSON) String() string {
    92  	return rv.name
    93  }
    94  
    95  // Cost returns the cost of this index as 1.
    96  func (rv *RegionJSON) Cost() int {
    97  	return 1
    98  }
    99  
   100  // IsUnique returns true since the Vindex is unique.
   101  func (rv *RegionJSON) IsUnique() bool {
   102  	return true
   103  }
   104  
   105  // NeedsVCursor satisfies the Vindex interface.
   106  func (rv *RegionJSON) NeedsVCursor() bool {
   107  	return false
   108  }
   109  
   110  // Map satisfies MultiColumn.
   111  func (rv *RegionJSON) Map(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value) ([]key.Destination, error) {
   112  	destinations := make([]key.Destination, 0, len(rowsColValues))
   113  	for _, row := range rowsColValues {
   114  		if len(row) != 2 {
   115  			destinations = append(destinations, key.DestinationNone{})
   116  			continue
   117  		}
   118  		// Compute hash.
   119  		hn, err := evalengine.ToUint64(row[0])
   120  		if err != nil {
   121  			destinations = append(destinations, key.DestinationNone{})
   122  			continue
   123  		}
   124  		h := vhash(hn)
   125  
   126  		rn, ok := rv.regionMap[row[1].ToString()]
   127  		if !ok {
   128  			destinations = append(destinations, key.DestinationNone{})
   129  			continue
   130  		}
   131  		r := make([]byte, 2)
   132  		binary.BigEndian.PutUint16(r, uint16(rn))
   133  
   134  		// Concatenate and add to destinations.
   135  		if rv.regionBytes == 1 {
   136  			r = r[1:]
   137  		}
   138  		dest := append(r, h...)
   139  		destinations = append(destinations, key.DestinationKeyspaceID(dest))
   140  	}
   141  	return destinations, nil
   142  }
   143  
   144  // Verify satisfies MultiColumn
   145  func (rv *RegionJSON) Verify(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value, ksids [][]byte) ([]bool, error) {
   146  	result := make([]bool, len(rowsColValues))
   147  	destinations, _ := rv.Map(ctx, vcursor, rowsColValues)
   148  	for i, dest := range destinations {
   149  		destksid, ok := dest.(key.DestinationKeyspaceID)
   150  		if !ok {
   151  			continue
   152  		}
   153  		result[i] = bytes.Equal([]byte(destksid), ksids[i])
   154  	}
   155  	return result, nil
   156  }
   157  
   158  func (rv *RegionJSON) PartialVindex() bool {
   159  	return false
   160  }