github.com/humans-group/gqlgen@v0.7.2/client/websocket.go (about)

     1  package client
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"github.com/gorilla/websocket"
     9  	"github.com/vektah/gqlparser/gqlerror"
    10  )
    11  
    12  const (
    13  	connectionInitMsg = "connection_init" // Client -> Server
    14  	startMsg          = "start"           // Client -> Server
    15  	connectionAckMsg  = "connection_ack"  // Server -> Client
    16  	dataMsg           = "data"            // Server -> Client
    17  	errorMsg          = "error"           // Server -> Client
    18  )
    19  
    20  type operationMessage struct {
    21  	Payload json.RawMessage `json:"payload,omitempty"`
    22  	ID      string          `json:"id,omitempty"`
    23  	Type    string          `json:"type"`
    24  }
    25  
    26  type Subscription struct {
    27  	Close func() error
    28  	Next  func(response interface{}) error
    29  }
    30  
    31  func errorSubscription(err error) *Subscription {
    32  	return &Subscription{
    33  		Close: func() error { return nil },
    34  		Next: func(response interface{}) error {
    35  			return err
    36  		},
    37  	}
    38  }
    39  
    40  func (p *Client) Websocket(query string, options ...Option) *Subscription {
    41  	return p.WebsocketWithPayload(query, nil, options...)
    42  }
    43  
    44  func (p *Client) WebsocketWithPayload(query string, initPayload map[string]interface{}, options ...Option) *Subscription {
    45  	r := p.mkRequest(query, options...)
    46  	requestBody, err := json.Marshal(r)
    47  	if err != nil {
    48  		return errorSubscription(fmt.Errorf("encode: %s", err.Error()))
    49  	}
    50  
    51  	url := strings.Replace(p.url, "http://", "ws://", -1)
    52  	url = strings.Replace(url, "https://", "wss://", -1)
    53  
    54  	c, _, err := websocket.DefaultDialer.Dial(url, nil)
    55  	if err != nil {
    56  		return errorSubscription(fmt.Errorf("dial: %s", err.Error()))
    57  	}
    58  
    59  	initMessage := operationMessage{Type: connectionInitMsg}
    60  	if initPayload != nil {
    61  		initMessage.Payload, err = json.Marshal(initPayload)
    62  		if err != nil {
    63  			return errorSubscription(fmt.Errorf("parse payload: %s", err.Error()))
    64  		}
    65  	}
    66  
    67  	if err = c.WriteJSON(initMessage); err != nil {
    68  		return errorSubscription(fmt.Errorf("init: %s", err.Error()))
    69  	}
    70  
    71  	var ack operationMessage
    72  	if err = c.ReadJSON(&ack); err != nil {
    73  		return errorSubscription(fmt.Errorf("ack: %s", err.Error()))
    74  	}
    75  	if ack.Type != connectionAckMsg {
    76  		return errorSubscription(fmt.Errorf("expected ack message, got %#v", ack))
    77  	}
    78  
    79  	if err = c.WriteJSON(operationMessage{Type: startMsg, ID: "1", Payload: requestBody}); err != nil {
    80  		return errorSubscription(fmt.Errorf("start: %s", err.Error()))
    81  	}
    82  
    83  	return &Subscription{
    84  		Close: c.Close,
    85  		Next: func(response interface{}) error {
    86  			var op operationMessage
    87  			c.ReadJSON(&op)
    88  			if op.Type != dataMsg {
    89  				if op.Type == errorMsg {
    90  					return fmt.Errorf(string(op.Payload))
    91  				} else {
    92  					return fmt.Errorf("expected data message, got %#v", op)
    93  				}
    94  			}
    95  
    96  			respDataRaw := map[string]interface{}{}
    97  			err = json.Unmarshal(op.Payload, &respDataRaw)
    98  			if err != nil {
    99  				return fmt.Errorf("decode: %s", err.Error())
   100  			}
   101  
   102  			if respDataRaw["errors"] != nil {
   103  				var errs []*gqlerror.Error
   104  				if err = unpack(respDataRaw["errors"], &errs); err != nil {
   105  					return err
   106  				}
   107  				if len(errs) > 0 {
   108  					return fmt.Errorf("errors: %s", errs)
   109  				}
   110  			}
   111  
   112  			return unpack(respDataRaw["data"], response)
   113  		},
   114  	}
   115  }