subscriptions: Add matching of subscriptions against entities

Author: Jannis Pohlmann

ast.go

@@ -0,0 +1,57 @@
+package graphqlws
+
+import (
+	"github.com/graphql-go/graphql/language/ast"
+)
+
+func operationDefinitionsWithOperation(
+	doc *ast.Document,
+	op string,
+) []*ast.OperationDefinition {
+	defs := []*ast.OperationDefinition{}
+	for _, node := range doc.Definitions {
+		if node.GetKind() == "OperationDefinition" {
+			if def, ok := node.(*ast.OperationDefinition); ok {
+				defs = append(defs, def)
+			}
+		}
+	}
+	return defs
+}
+
+func selectionSetsForOperationDefinitions(
+	defs []*ast.OperationDefinition,
+) []*ast.SelectionSet {
+	sets := []*ast.SelectionSet{}
+	for _, def := range defs {
+		if set := def.GetSelectionSet(); set != nil {
+			sets = append(sets, set)
+		}
+	}
+	return sets
+}
+
+func nameForSelectionSet(set *ast.SelectionSet) (string, bool) {
+	if len(set.Selections) >= 1 {
+		if field, ok := set.Selections[0].(*ast.Field); ok {
+			return field.Name.Value, true
+		}
+	}
+	return "", false
+}
+
+func namesForSelectionSets(sets []*ast.SelectionSet) []string {
+	names := []string{}
+	for _, set := range sets {
+		if name, ok := nameForSelectionSet(set); ok {
+			names = append(names, name)
+		}
+	}
+	return names
+}
+
+func subscriptionFieldNamesFromDocument(doc *ast.Document) []string {
+	defs := operationDefinitionsWithOperation(doc, "subscription")
+	sets := selectionSetsForOperationDefinitions(defs)
+	return namesForSelectionSets(sets)
+}

subscriptions.go

@@ -31,8 +31,27 @@ type Subscription struct {
 	Query         string
 	Variables     map[string]interface{}
 	OperationName string
-	SendData      SubscriptionSendDataFunc
 	Document      *ast.Document
+	Fields        []string
+	Connection    Connection
+	SendData      SubscriptionSendDataFunc
+}
+
+// MatchesField returns true if the subscription is for data that
+// belongs to the given field.
+func (s *Subscription) MatchesField(field string) bool {
+	if s.Document == nil || len(s.Fields) == 0 {
+		return false
+	}
+
+	// The subscription matches the field if any of the queries have
+	// the same name as the field
+	for _, name := range s.Fields {
+		if name == field {
+			return true
+		}
+	}
+	return false
 }
 
 // ConnectionSubscriptions defines a map of all subscriptions of
@@ -113,6 +132,9 @@ func (m *subscriptionManager) AddSubscription(
 	// Remember the query document for later
 	subscription.Document = document
 
+	// Extract query names from the document (typically, there should only be one)
+	subscription.Fields = subscriptionFieldNamesFromDocument(document)
+
 	// Allocate the connection's map of subscription IDs to
 	// subscriptions on demand
 	if m.subscriptions[conn] == nil {