github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/testutil/zero.go (about)

     1  /*
     2   * Copyright 2019 Dgraph Labs, Inc. and Contributors
     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 testutil
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/json"
    22  	"fmt"
    23  	"io/ioutil"
    24  	"net/http"
    25  	"strconv"
    26  	"strings"
    27  
    28  	"github.com/dgraph-io/dgo"
    29  	"github.com/dgraph-io/dgo/protos/api"
    30  	"github.com/pkg/errors"
    31  	"google.golang.org/grpc"
    32  )
    33  
    34  // StateResponse represents the structure of the JSON object returned by calling
    35  // the /state endpoint in zero.
    36  type StateResponse struct {
    37  	Groups map[string]struct {
    38  		Members map[string]struct {
    39  			Addr       string `json:"addr"`
    40  			GroupID    int    `json:"groupId"`
    41  			ID         string `json:"id"`
    42  			LastUpdate string `json:"lastUpdate"`
    43  			Leader     bool   `json:"leader"`
    44  		} `json:"members"`
    45  		Tablets map[string]struct {
    46  			GroupID   int    `json:"groupId"`
    47  			Predicate string `json:"predicate"`
    48  		} `json:"tablets"`
    49  	} `json:"groups"`
    50  	Removed []struct {
    51  		Addr    string `json:"addr"`
    52  		GroupID int    `json:"groupId"`
    53  		ID      string `json:"id"`
    54  	} `json:"removed"`
    55  }
    56  
    57  // GetState queries the /state endpoint in zero and returns the response.
    58  func GetState() (*StateResponse, error) {
    59  	resp, err := http.Get("http://" + SockAddrZeroHttp + "/state")
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  	defer resp.Body.Close()
    64  
    65  	b, err := ioutil.ReadAll(resp.Body)
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  
    70  	if bytes.Contains(b, []byte("Error")) {
    71  		return nil, errors.Errorf("Failed to get state: %s", string(b))
    72  	}
    73  
    74  	var st StateResponse
    75  	if err := json.Unmarshal(b, &st); err != nil {
    76  		return nil, err
    77  	}
    78  	return &st, nil
    79  }
    80  
    81  // GetClientToGroup returns a dgraph client connected to an alpha in the given group.
    82  func GetClientToGroup(groupID string) (*dgo.Dgraph, error) {
    83  	state, err := GetState()
    84  	if err != nil {
    85  		return nil, err
    86  	}
    87  
    88  	group, ok := state.Groups[groupID]
    89  	if !ok {
    90  		return nil, errors.Errorf("group %s does not exist", groupID)
    91  	}
    92  
    93  	if len(group.Members) == 0 {
    94  		return nil, errors.Errorf("the group %s has no members", groupID)
    95  	}
    96  
    97  	member := group.Members["1"]
    98  	parts := strings.Split(member.Addr, ":")
    99  	if len(parts) != 2 {
   100  		return nil, errors.Errorf("the member has an invalid address: %v", member.Addr)
   101  	}
   102  	// internalPort is used for communication between alpha nodes
   103  	internalPort, err := strconv.Atoi(parts[1])
   104  	if err != nil {
   105  		return nil, errors.Errorf("unable to parse the port number from %s", parts[1])
   106  	}
   107  
   108  	// externalPort is for handling connections from clients
   109  	externalPort := internalPort + 2000
   110  
   111  	conn, err := grpc.Dial(fmt.Sprintf("localhost:%d", externalPort), grpc.WithInsecure())
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  	return dgo.NewDgraphClient(api.NewDgraphClient(conn)), nil
   116  }