github.com/fortexxx/gqlgen@v0.10.3-0.20191216030626-ca5ea8b21ead/example/chat/resolvers.go (about)

     1  //go:generate go run ../../testdata/gqlgen.go
     2  
     3  package chat
     4  
     5  import (
     6  	"context"
     7  	"math/rand"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/99designs/gqlgen/graphql"
    12  )
    13  
    14  type resolver struct {
    15  	Rooms map[string]*Chatroom
    16  	mu    sync.Mutex // nolint: structcheck
    17  }
    18  
    19  func (r *resolver) Mutation() MutationResolver {
    20  	return &mutationResolver{r}
    21  }
    22  
    23  func (r *resolver) Query() QueryResolver {
    24  	return &queryResolver{r}
    25  }
    26  
    27  func (r *resolver) Subscription() SubscriptionResolver {
    28  	return &subscriptionResolver{r}
    29  }
    30  
    31  func New() Config {
    32  	return Config{
    33  		Resolvers: &resolver{
    34  			Rooms: map[string]*Chatroom{},
    35  		},
    36  		Directives: DirectiveRoot{
    37  			User: func(ctx context.Context, obj interface{}, next graphql.Resolver, username string) (res interface{}, err error) {
    38  				return next(context.WithValue(ctx, "username", username))
    39  			},
    40  		},
    41  	}
    42  }
    43  
    44  func getUsername(ctx context.Context) string {
    45  	if username, ok := ctx.Value("username").(string); ok {
    46  		return username
    47  	}
    48  	return ""
    49  }
    50  
    51  type Chatroom struct {
    52  	Name      string
    53  	Messages  []Message
    54  	Observers map[string]struct {
    55  		Username string
    56  		Message  chan *Message
    57  	}
    58  }
    59  
    60  type mutationResolver struct{ *resolver }
    61  
    62  func (r *mutationResolver) Post(ctx context.Context, text string, username string, roomName string) (*Message, error) {
    63  	r.mu.Lock()
    64  	room := r.Rooms[roomName]
    65  	if room == nil {
    66  		room = &Chatroom{
    67  			Name: roomName,
    68  			Observers: map[string]struct {
    69  				Username string
    70  				Message  chan *Message
    71  			}{},
    72  		}
    73  		r.Rooms[roomName] = room
    74  	}
    75  	r.mu.Unlock()
    76  
    77  	message := Message{
    78  		ID:        randString(8),
    79  		CreatedAt: time.Now(),
    80  		Text:      text,
    81  		CreatedBy: username,
    82  	}
    83  
    84  	room.Messages = append(room.Messages, message)
    85  	r.mu.Lock()
    86  	for _, observer := range room.Observers {
    87  		if observer.Username == "" || observer.Username == message.CreatedBy {
    88  			observer.Message <- &message
    89  		}
    90  	}
    91  	r.mu.Unlock()
    92  	return &message, nil
    93  }
    94  
    95  type queryResolver struct{ *resolver }
    96  
    97  func (r *queryResolver) Room(ctx context.Context, name string) (*Chatroom, error) {
    98  	r.mu.Lock()
    99  	room := r.Rooms[name]
   100  	if room == nil {
   101  		room = &Chatroom{
   102  			Name: name,
   103  			Observers: map[string]struct {
   104  				Username string
   105  				Message  chan *Message
   106  			}{},
   107  		}
   108  		r.Rooms[name] = room
   109  	}
   110  	r.mu.Unlock()
   111  
   112  	return room, nil
   113  }
   114  
   115  type subscriptionResolver struct{ *resolver }
   116  
   117  func (r *subscriptionResolver) MessageAdded(ctx context.Context, roomName string) (<-chan *Message, error) {
   118  	r.mu.Lock()
   119  	room := r.Rooms[roomName]
   120  	if room == nil {
   121  		room = &Chatroom{
   122  			Name: roomName,
   123  			Observers: map[string]struct {
   124  				Username string
   125  				Message  chan *Message
   126  			}{},
   127  		}
   128  		r.Rooms[roomName] = room
   129  	}
   130  	r.mu.Unlock()
   131  
   132  	id := randString(8)
   133  	events := make(chan *Message, 1)
   134  
   135  	go func() {
   136  		<-ctx.Done()
   137  		r.mu.Lock()
   138  		delete(room.Observers, id)
   139  		r.mu.Unlock()
   140  	}()
   141  
   142  	r.mu.Lock()
   143  	room.Observers[id] = struct {
   144  		Username string
   145  		Message  chan *Message
   146  	}{Username: getUsername(ctx), Message: events}
   147  	r.mu.Unlock()
   148  
   149  	return events, nil
   150  }
   151  
   152  var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
   153  
   154  func randString(n int) string {
   155  	b := make([]rune, n)
   156  	for i := range b {
   157  		b[i] = letterRunes[rand.Intn(len(letterRunes))]
   158  	}
   159  	return string(b)
   160  }