diff --git a/controller/go.mod b/controller/go.mod
index 75de018..955f4b8 100644
--- a/controller/go.mod
+++ b/controller/go.mod
@@ -3,6 +3,7 @@
 go 1.14
 
 require (
+	github.com/bradleyjkemp/memviz v0.2.3
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/itaysk/regogo v0.0.0-20200418072509-74b59e1875c2
 	github.com/vektah/gqlparser v1.3.1
diff --git a/controller/go.sum b/controller/go.sum
index 55b74b5..eb9f888 100644
--- a/controller/go.sum
+++ b/controller/go.sum
@@ -49,6 +49,9 @@
 github.com/blevesearch/go-porterstemmer v1.0.2/go.mod h1:haWQqFT3RdOGz7PJuM3or/pWNJS1pKkoZJWCkWu0DVA=
 github.com/blevesearch/segment v0.0.0-20160915185041-762005e7a34f/go.mod h1:IInt5XRvpiGE09KOk9mmCMLjHhydIhNPKPPFLFBB7L8=
 github.com/blevesearch/snowballstem v0.0.0-20180110192139-26b06a2c243d/go.mod h1:cdytUvf6FKWA9NpXJihYdZq8TN2AiQ5HOS0UZUz0C9g=
+github.com/bradleyjkemp/cupaloy/v2 v2.5.0/go.mod h1:TD5UU0rdYTbu/TtuwFuWrtiRARuN7mtRipvs/bsShSE=
+github.com/bradleyjkemp/memviz v0.2.3 h1:8fqKnV1xQz4NQkDy5Gklhm9fGtUK+R3oW0z1unBDFGY=
+github.com/bradleyjkemp/memviz v0.2.3/go.mod h1:meU694rvawW7NqtNLtlg+TEU+UqAjrbJayEPZQUSOBs=
 github.com/cenkalti/backoff v2.1.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
diff --git a/controller/main.go b/controller/main.go
index 4802b3e..70e4af0 100644
--- a/controller/main.go
+++ b/controller/main.go
@@ -163,28 +163,38 @@
 	if err != nil {
 		panic(err)
 	}
-	// err = gqlClient.SetSchema(`
-	// type Image {
-	//      id: ID!
-	//      objectPath: String! @search(by: [exact])
-	// }
+	err = gqlClient.SetSchema(`
+enum EventState {
+  NEW
+  PROCESSING
+  DONE
+}
 
-	// type ImageSegment {
-	//      id: ID!
-	//      upperLeftX: Int!
-	//      upperLeftY: Int!
-	//      lowerRightX: Int!
-	//      lowerRightY: Int!
-	//      sourceImage: Image!
-	//      objectPath: String
-	// }
+type Foo { bar: Int }`)
+	if err != nil {
+		panic(err)
+	}
+	err = gqlClient.AddSchema(`
+	type Image {
+	     id: ID!
+	     objectPath: String! @search(by: [exact])
+	}
 
-	// extend type Image {
-	//      segments: [ImageSegment] @hasInverse(field: sourceImage)
-	// }`)
-	// if err != nil {
-	// 	panic(err)
-	// }
+	type ImageSegment {
+	     id: ID!
+	     upperLeftX: Float!
+	     upperLeftY: Float!
+	     lowerRightX: Float!
+	     lowerRightY: Float!
+	     sourceImage: Image! @hasInverse(field: segments)
+	}
+
+	extend type Image {
+	     segments: [ImageSegment] @hasInverse(field: sourceImage)
+	}`)
+	if err != nil {
+		panic(err)
+	}
 	mw := MinioWebhook{gqlClient, pods}
 	http.HandleFunc("/minio_webhook", mw.minioHandler)
 	http.HandleFunc("/graphql", mw.graphqlHandler)
diff --git a/controller/schema/dgraph_schema_store.go b/controller/schema/dgraph_schema_store.go
index 6725cbe..941a16b 100644
--- a/controller/schema/dgraph_schema_store.go
+++ b/controller/schema/dgraph_schema_store.go
@@ -4,6 +4,7 @@
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"strings"
@@ -12,6 +13,8 @@
 	"github.com/itaysk/regogo"
 	"github.com/vektah/gqlparser"
 	"github.com/vektah/gqlparser/ast"
+	"github.com/vektah/gqlparser/formatter"
+	//"github.com/bradleyjkemp/memviz"
 )
 
 const jsonContentType = "application/json"
@@ -21,6 +24,16 @@
 const getSchemaQuery = `{ getGQLSchema() { schema generatedSchema } }`
 const setSchemaQuery = `mutation { updateGQLSchema(input: { set: { schema: "%s" } }) { gqlSchema { schema generatedSchema } } }`
 const runQuery = `{ "query": "%s" }`
+const eventTmpl = `
+type %sEvent {
+  id: ID!
+  state: EventState!
+  node: %s! @hasInverse(field: events)
+}
+
+extend type %s {
+  events: [%sEvent] @hasInverse(field: node)
+}`
 
 type DgraphClient struct {
 	gqlAddress      string
@@ -47,19 +60,28 @@
 }
 
 func (s *DgraphClient) AddSchema(gqlSchema string) error {
-	return s.SetSchema(s.userSchema + gqlSchema)
+	extendedSchema, err := s.extendSchema(gqlSchema)
+	if err != nil {
+		return err
+	} else {
+		return s.SetSchema(extendedSchema)
+	}
 }
 
 func (s *DgraphClient) SetSchema(gqlSchema string) error {
 	glog.Info("Setting GraphQL schema")
 	glog.Info(gqlSchema)
 	resp, err := s.runQuery(
-		fmt.Sprintf(setSchemaQuery, sanitizeSchema(gqlSchema)),
+		fmt.Sprintf(setSchemaQuery, gqlSchema),
 		s.schemaAddress)
 	if err != nil {
 		return err
 	}
-	return s.updateSchema(resp)
+	data, err := regogo.Get(resp, "input.updateGQLSchema.gqlSchema")
+	if err != nil {
+		return err
+	}
+	return s.updateSchema(data.JSON())
 }
 
 func (s *DgraphClient) fetchSchema() error {
@@ -68,15 +90,19 @@
 	if err != nil {
 		return err
 	}
-	return s.updateSchema(resp)
-}
-
-func (s *DgraphClient) updateSchema(resp string) error {
-	userSchema, err := regogo.Get(resp, "input.getGQLSchema.schema")
+	data, err := regogo.Get(resp, "input.getGQLSchema")
 	if err != nil {
 		return err
 	}
-	generatedSchema, err := regogo.Get(resp, "input.getGQLSchema.generatedSchema")
+	return s.updateSchema(data.JSON())
+}
+
+func (s *DgraphClient) updateSchema(resp string) error {
+	userSchema, err := regogo.Get(resp, "input.schema")
+	if err != nil {
+		return err
+	}
+	generatedSchema, err := regogo.Get(resp, "input.generatedSchema")
 	if err != nil {
 		return err
 	}
@@ -91,16 +117,21 @@
 }
 
 func (s *DgraphClient) RunQuery(query string) (string, error) {
-	_, gqlErr := gqlparser.LoadQuery(s.Schema(), query)
+	q, gqlErr := gqlparser.LoadQuery(s.Schema(), query)
 	if gqlErr != nil {
 		return "", errors.New(gqlErr.Error())
 	}
+	rewritten := rewriteQuery(q, s.Schema())
+	var b strings.Builder
+	// TODO(giolekva): gqlparser should be reporting error back
+	formatter.NewFormatter(&b).FormatQueryDocument(rewritten)
+	query = b.String()
 	return s.runQuery(query, s.gqlAddress)
 }
 
 func (s *DgraphClient) runQuery(query string, onAddr string) (string, error) {
 	glog.Infof("Running GraphQL query: %s", query)
-	queryJson := fmt.Sprintf(runQuery, sanitizeQuery(query))
+	queryJson := fmt.Sprintf(runQuery, fixWhitespaces(escapeQuery(query)))
 	resp, err := http.Post(
 		onAddr,
 		jsonContentType,
@@ -125,11 +156,160 @@
 	return data.JSON(), nil
 }
 
-func sanitizeSchema(schema string) string {
+func (s *DgraphClient) extendSchema(schema string) (string, error) {
+	try := s.generatedSchema + "   " + schema
+	parsed, gqlErr := gqlparser.LoadSchema(&ast.Source{Input: try})
+	if gqlErr != nil {
+		return "", errors.New(gqlErr.Error())
+	}
+	var extended strings.Builder
+	_, err := io.WriteString(&extended, s.userSchema)
+	if err != nil {
+		return "", err
+	}
+	_, err = io.WriteString(&extended, schema)
+	if err != nil {
+		return "", err
+	}
+	for _, t := range parsed.Types {
+		if shouldIgnoreDefinition(t) {
+			continue
+		}
+		_, err := fmt.Fprintf(&extended, eventTmpl, t.Name, t.Name, t.Name, t.Name)
+		if err != nil {
+			return "", err
+		}
+	}
+	return extended.String(), nil
+}
+
+func findDefinitionWithName(name string, s *ast.Schema) *ast.Definition {
+	for n, d := range s.Types {
+		if n == name {
+			return d
+		}
+	}
+	panic(fmt.Sprintf("Expected event input definiton for  %s", name))
+}
+
+func findEventnputDefinitionFor(d *ast.Definition, s *ast.Schema) *ast.Definition {
+	if !strings.HasSuffix(d.Name, "Input") {
+		panic(fmt.Sprintf("Expected input definiton, got %s", d.Name))
+	}
+	eventInput := fmt.Sprintf("%sEventInput", strings.TrimSuffix(d.Name, "Input"))
+	return findDefinitionWithName(eventInput, s)
+}
+
+func newEventStateValue(s *ast.Schema) *ast.ChildValue {
+	return &ast.ChildValue{
+		Name:     "state",
+		Position: nil,
+		Value: &ast.Value{
+			Raw:                "NEW",
+			Children:           ast.ChildValueList{},
+			Kind:               ast.EnumValue,
+			Position:           nil,
+			Definition:         findDefinitionWithName("EventState", s),
+			VariableDefinition: nil,
+			ExpectedType:       nil,
+		},
+	}
+}
+
+func newEventListValue(d *ast.Definition, s *ast.Schema) *ast.ChildValue {
+	return &ast.ChildValue{
+		Name:     "events",
+		Position: nil,
+		Value: &ast.Value{
+			Raw:                "",
+			Children:           ast.ChildValueList{newEventValue(d, s)},
+			Kind:               ast.ListValue,
+			Position:           nil,
+			Definition:         findEventnputDefinitionFor(d, s),
+			VariableDefinition: nil,
+			ExpectedType:       nil,
+		},
+	}
+}
+
+func newEventValue(d *ast.Definition, s *ast.Schema) *ast.ChildValue {
+	return &ast.ChildValue{
+		Name:     "events",
+		Position: nil,
+		Value: &ast.Value{
+			Raw:                "",
+			Children:           ast.ChildValueList{newEventStateValue(s)},
+			Kind:               ast.ObjectValue,
+			Position:           nil,
+			Definition:         findEventnputDefinitionFor(d, s),
+			VariableDefinition: nil,
+			ExpectedType:       nil,
+		},
+	}
+}
+
+func rewriteValue(v *ast.Value, s *ast.Schema) {
+	if v == nil {
+		panic("Received nil value")
+	}
+	switch v.Kind {
+	case ast.Variable:
+	case ast.IntValue:
+	case ast.FloatValue:
+	case ast.StringValue:
+	case ast.BlockValue:
+	case ast.BooleanValue:
+	case ast.NullValue:
+	case ast.EnumValue:
+	case ast.ListValue:
+		for _, c := range v.Children {
+			rewriteValue(c.Value, s)
+		}
+	case ast.ObjectValue:
+		for _, c := range v.Children {
+			rewriteValue(c.Value, s)
+		}
+		if v.Definition.Kind == ast.InputObject &&
+			!strings.HasSuffix(v.Definition.Name, "Event") {
+			v.Children = append(v.Children, newEventListValue(v.Definition, s))
+		}
+	}
+}
+
+func rewriteQuery(q *ast.QueryDocument, s *ast.Schema) *ast.QueryDocument {
+	for _, op := range q.Operations {
+		if op.Operation != ast.Mutation {
+			continue
+		}
+		for _, sel := range op.SelectionSet {
+			field, ok := sel.(*ast.Field)
+			if !ok {
+				panic(sel)
+			}
+			for _, arg := range field.Arguments {
+				rewriteValue(arg.Value, s)
+			}
+		}
+	}
+	return q
+
+}
+
+// TODO(giolekva): will be safer to use directive instead
+func shouldIgnoreDefinition(d *ast.Definition) bool {
+	return d.Kind != ast.Object ||
+		d.Name == "Query" ||
+		d.Name == "Mutation" ||
+		strings.HasPrefix(d.Name, "__") ||
+		strings.HasSuffix(d.Name, "Payload") ||
+		strings.HasSuffix(d.Name, "Event")
+}
+
+func fixWhitespaces(schema string) string {
 	return strings.ReplaceAll(
 		strings.ReplaceAll(schema, "\n", " "), "\t", " ")
 }
 
-func sanitizeQuery(query string) string {
+func escapeQuery(query string) string {
 	return strings.ReplaceAll(query, "\"", "\\\"")
 }
diff --git a/controller/tests/query_test.go b/controller/tests/query_test.go
index b4f9898..73e5bdf 100644
--- a/controller/tests/query_test.go
+++ b/controller/tests/query_test.go
@@ -1,33 +1,111 @@
 package tests
 
 import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
 	"testing"
 
 	"github.com/vektah/gqlparser"
 	"github.com/vektah/gqlparser/ast"
+
+	"github.com/bradleyjkemp/memviz"
 )
 
 var gqlSchema = `#######################\n# Input Schema\n#######################\n\ntype Image {\n\tid: ID!\n\tobjectPath: String! @search(by: [exact])\n\tsegments(filter: ImageSegmentFilter, order: ImageSegmentOrder, first: Int, offset: Int): [ImageSegment] @hasInverse(field: sourceImage)\n}\n\ntype ImageSegment {\n\tid: ID!\n\tupperLeftX: Float!\n\tupperLeftY: Float!\n\tlowerRightX: Float!\n\tlowerRightY: Float!\n\tsourceImage(filter: ImageFilter): Image! @hasInverse(field: segments)\n\tobjectPath: String\n}\n\n#######################\n# Extended Definitions\n#######################\n\nscalar DateTime\n\nenum DgraphIndex {\n\tint\n\tfloat\n\tbool\n\thash\n\texact\n\tterm\n\tfulltext\n\ttrigram\n\tregexp\n\tyear\n\tmonth\n\tday\n\thour\n}\n\ndirective @hasInverse(field: String!) on FIELD_DEFINITION\ndirective @search(by: [DgraphIndex!]) on FIELD_DEFINITION\ndirective @dgraph(type: String, pred: String) on OBJECT | INTERFACE | FIELD_DEFINITION\ndirective @id on FIELD_DEFINITION\ndirective @secret(field: String!, pred: String) on OBJECT | INTERFACE\n\ninput IntFilter {\n\teq: Int\n\tle: Int\n\tlt: Int\n\tge: Int\n\tgt: Int\n}\n\ninput FloatFilter {\n\teq: Float\n\tle: Float\n\tlt: Float\n\tge: Float\n\tgt: Float\n}\n\ninput DateTimeFilter {\n\teq: DateTime\n\tle: DateTime\n\tlt: DateTime\n\tge: DateTime\n\tgt: DateTime\n}\n\ninput StringTermFilter {\n\tallofterms: String\n\tanyofterms: String\n}\n\ninput StringRegExpFilter {\n\tregexp: String\n}\n\ninput StringFullTextFilter {\n\talloftext: String\n\tanyoftext: String\n}\n\ninput StringExactFilter {\n\teq: String\n\tle: String\n\tlt: String\n\tge: String\n\tgt: String\n}\n\ninput StringHashFilter {\n\teq: String\n}\n\n#######################\n# Generated Types\n#######################\n\ntype AddImagePayload {\n\timage(filter: ImageFilter, order: ImageOrder, first: Int, offset: Int): [Image]\n\tnumUids: Int\n}\n\ntype AddImageSegmentPayload {\n\timagesegment(filter: ImageSegmentFilter, order: ImageSegmentOrder, first: Int, offset: Int): [ImageSegment]\n\tnumUids: Int\n}\n\ntype DeleteImagePayload {\n\tmsg: String\n\tnumUids: Int\n}\n\ntype DeleteImageSegmentPayload {\n\tmsg: String\n\tnumUids: Int\n}\n\ntype UpdateImagePayload {\n\timage(filter: ImageFilter, order: ImageOrder, first: Int, offset: Int): [Image]\n\tnumUids: Int\n}\n\ntype UpdateImageSegmentPayload {\n\timagesegment(filter: ImageSegmentFilter, order: ImageSegmentOrder, first: Int, offset: Int): [ImageSegment]\n\tnumUids: Int\n}\n\n#######################\n# Generated Enums\n#######################\n\nenum ImageOrderable {\n\tobjectPath\n}\n\nenum ImageSegmentOrderable {\n\tupperLeftX\n\tupperLeftY\n\tlowerRightX\n\tlowerRightY\n\tobjectPath\n}\n\n#######################\n# Generated Inputs\n#######################\n\ninput AddImageInput {\n\tobjectPath: String!\n\tsegments: [ImageSegmentRef]\n}\n\ninput AddImageSegmentInput {\n\tupperLeftX: Float!\n\tupperLeftY: Float!\n\tlowerRightX: Float!\n\tlowerRightY: Float!\n\tsourceImage: ImageRef!\n\tobjectPath: String\n}\n\ninput ImageFilter {\n\tid: [ID!]\n\tobjectPath: StringExactFilter\n\tand: ImageFilter\n\tor: ImageFilter\n\tnot: ImageFilter\n}\n\ninput ImageOrder {\n\tasc: ImageOrderable\n\tdesc: ImageOrderable\n\tthen: ImageOrder\n}\n\ninput ImagePatch {\n\tobjectPath: String\n\tsegments: [ImageSegmentRef]\n}\n\ninput ImageRef {\n\tid: ID\n\tobjectPath: String\n\tsegments: [ImageSegmentRef]\n}\n\ninput ImageSegmentFilter {\n\tid: [ID!]\n\tnot: ImageSegmentFilter\n}\n\ninput ImageSegmentOrder {\n\tasc: ImageSegmentOrderable\n\tdesc: ImageSegmentOrderable\n\tthen: ImageSegmentOrder\n}\n\ninput ImageSegmentPatch {\n\tupperLeftX: Float\n\tupperLeftY: Float\n\tlowerRightX: Float\n\tlowerRightY: Float\n\tsourceImage: ImageRef\n\tobjectPath: String\n}\n\ninput ImageSegmentRef {\n\tid: ID\n\tupperLeftX: Float\n\tupperLeftY: Float\n\tlowerRightX: Float\n\tlowerRightY: Float\n\tsourceImage: ImageRef\n\tobjectPath: String\n}\n\ninput UpdateImageInput {\n\tfilter: ImageFilter!\n\tset: ImagePatch\n\tremove: ImagePatch\n}\n\ninput UpdateImageSegmentInput {\n\tfilter: ImageSegmentFilter!\n\tset: ImageSegmentPatch\n\tremove: ImageSegmentPatch\n}\n\n#######################\n# Generated Query\n#######################\n\ntype Query {\n\tgetImage(id: ID!): Image\n\tqueryImage(filter: ImageFilter, order: ImageOrder, first: Int, offset: Int): [Image]\n\tgetImageSegment(id: ID!): ImageSegment\n\tqueryImageSegment(filter: ImageSegmentFilter, order: ImageSegmentOrder, first: Int, offset: Int): [ImageSegment]\n}\n\n#######################\n# Generated Mutations\n#######################\n\ntype Mutation {\n\taddImage(input: [AddImageInput!]!): AddImagePayload\n\tupdateImage(input: UpdateImageInput!): UpdateImagePayload\n\tdeleteImage(filter: ImageFilter!): DeleteImagePayload\n\taddImageSegment(input: [AddImageSegmentInput!]!): AddImageSegmentPayload\n\tupdateImageSegment(input: UpdateImageSegmentInput!): UpdateImageSegmentPayload\n\tdeleteImageSegment(filter: ImageSegmentFilter!): DeleteImageSegmentPayload\n}\n`
 
+var simpleSchema = `
+type ABC {
+  id: ID!
+  x: Int
+}
+
+input ABCInput {
+  x: Int
+}
+
+extend type ABC {
+  event: ID
+}
+`
+
+var enumSchema = `
+enum EventState {
+  NEW
+  PROCESSING
+  DONE
+}
+
+input Event {
+  id: ID
+  state: EventState!
+}
+
+type E {
+  foo: String
+}
+
+type Mutation {
+  addEvent(input: Event!): E
+}
+`
+
+// TODO(giolekva): will be safer to use directive instead
+func shouldIgnoreDefinition(d *ast.Definition) bool {
+	return d.Kind != ast.Object ||
+		d.Name == "Query" ||
+		d.Name == "Mutation" ||
+		strings.HasPrefix(d.Name, "__") ||
+		strings.HasSuffix(d.Name, "Payload")
+}
+
+func checkErr(err error) {
+	if err != nil {
+		panic(err)
+	}
+}
+
+func TestParseSchema(t *testing.T) {
+	schema := getSchema(gqlSchema)
+	var extended strings.Builder
+	_, err := io.WriteString(&extended, gqlSchema)
+	checkErr(err)
+	for _, n := range schema.Types {
+		if shouldIgnoreDefinition(n) {
+			continue
+		}
+		_, err := fmt.Fprintf(&extended, "\nextend type %s {\n event: ID\n}", n.Name)
+		checkErr(err)
+	}
+	fmt.Print(&extended)
+	f, _ := os.Create("schema.dot")
+	memviz.Map(f, schema)
+	f.Close()
+}
+
 func TestParseQuery(t *testing.T) {
-	schema := getSchema()
-	query, err := gqlparser.LoadQuery(schema, `{
-getImage(id: "0x2") {
-  id
-  objectPath
+	schema := getSchema(enumSchema)
+	query, err := gqlparser.LoadQuery(schema, `mutation {
+addEvent(input: {
+  state: NEW
+}) {
+  foo
 }
 }`)
 	if err != nil {
 		panic(err)
 	}
-	print(ast.Dump(schema.Mutation))
-	print(ast.Dump(query))
+	f, _ := os.Create("query.dot")
+	memviz.Map(f, query)
+	f.Close()
 }
 
-func getSchema() *ast.Schema {
-	schema, err := gqlparser.LoadSchema(&ast.Source{Input: gqlSchema})
+func getSchema(schema string) *ast.Schema {
+	s, err := gqlparser.LoadSchema(&ast.Source{
+		Input: strings.ReplaceAll(strings.ReplaceAll(schema, "\\n", "\n"), "\\t", "\t")})
 	if err != nil {
 		panic(err)
 	}
-	return schema
+	return s
 }