diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f6abe19..6459e1d1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,7 +11,7 @@
 - Convert usernaems to lowercase when constructing scope names etc. (#322)
 - Initial support for Adaptive Rate Limits.
 - Many accidentally exported names have been unexported. (#327)
-- Add new GetEntitySchema function to the Connector interface (#335)
+- Add new GetEntitySchema function to the Connector interface (#335 and #338)
 - Remove unused PluginFunc argument from the routing connector (#337)
 
 ## v2.6.0 (2018-04-16)
diff --git a/connector.go b/connector.go
index bd2e0437..463b5fe3 100644
--- a/connector.go
+++ b/connector.go
@@ -190,7 +190,7 @@ type Connector interface {
 	// CheckSchemaStatus checks the status of the schema whether it is accepted or in progress of application.
 	CheckSchemaStatus(ctx context.Context, scope string, namePrefix string, version int32) (*SchemaStatus, error)
 	// GetEntitySchema returns the entity info for a given entity in a given scope and prefix.
-	GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*EntityInfo, error)
+	GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*EntityDefinition, error)
 
 	// Datastore management
 	// CreateScope creates a scope for storage of data, usually implemented by a keyspace for this data
diff --git a/connectors/base/base.go b/connectors/base/base.go
index cb8ed0fb..ee5f4688 100644
--- a/connectors/base/base.go
+++ b/connectors/base/base.go
@@ -168,7 +168,7 @@ func (c *Connector) CheckSchemaStatus(ctx context.Context, scope string, namePre
 }
 
 // GetEntitySchema calls next
-func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityInfo, error) {
+func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityDefinition, error) {
 	if c.Next == nil {
 		return nil, NewErrNoMoreConnector()
 	}
diff --git a/connectors/devnull/devnull.go b/connectors/devnull/devnull.go
index 90a4c0a6..af517254 100644
--- a/connectors/devnull/devnull.go
+++ b/connectors/devnull/devnull.go
@@ -116,8 +116,8 @@ func (c *Connector) CheckSchemaStatus(ctx context.Context, scope, namePrefix str
 }
 
 // GetEntitySchema always returns a blank EntityInfo and no error.
-func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityInfo, error) {
-	return &dosa.EntityInfo{}, nil
+func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityDefinition, error) {
+	return &dosa.EntityDefinition{}, nil
 }
 
 // CreateScope returns success
diff --git a/connectors/random/random.go b/connectors/random/random.go
index 532c83f1..ce26dac0 100644
--- a/connectors/random/random.go
+++ b/connectors/random/random.go
@@ -188,8 +188,8 @@ func (c *Connector) CheckSchemaStatus(ctx context.Context, scope, namePrefix str
 }
 
 // GetEntitySchema always returns a blank EntityInfo and no error.
-func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityInfo, error) {
-	return &dosa.EntityInfo{}, nil
+func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityDefinition, error) {
+	return &dosa.EntityDefinition{}, nil
 }
 
 // CreateScope returns success
diff --git a/connectors/routing/connector.go b/connectors/routing/connector.go
index c2f12085..55a9aeec 100644
--- a/connectors/routing/connector.go
+++ b/connectors/routing/connector.go
@@ -188,7 +188,7 @@ func (rc *Connector) CheckSchemaStatus(ctx context.Context, scope string, namePr
 }
 
 // GetEntitySchema calls the selected connector
-func (rc *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityInfo, error) {
+func (rc *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityDefinition, error) {
 	connector, err := rc.getConnector(scope, namePrefix)
 	if err != nil {
 		return nil, err
diff --git a/connectors/yarpc/yarpc.go b/connectors/yarpc/yarpc.go
index 0289de5e..37fde4d7 100644
--- a/connectors/yarpc/yarpc.go
+++ b/connectors/yarpc/yarpc.go
@@ -592,7 +592,7 @@ func (c *Connector) CheckSchemaStatus(ctx context.Context, scope, namePrefix str
 }
 
 // GetEntitySchema gets the schema for the specified entity.
-func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityInfo, error) {
+func (c *Connector) GetEntitySchema(ctx context.Context, scope, namePrefix, entityName string, version int32) (*dosa.EntityDefinition, error) {
 	// We're not going to implement this at the moment since it's not needed. However, it could be easily
 	// implemented by adding a new method to the thrift idl
 	panic("Not implemented")
diff --git a/mocks/client.go b/mocks/client.go
index 20c5c287..bd229ab8 100644
--- a/mocks/client.go
+++ b/mocks/client.go
@@ -26,9 +26,10 @@ package mocks
 
 import (
 	context "context"
+	reflect "reflect"
+
 	gomock "github.com/golang/mock/gomock"
 	dosa "github.com/uber-go/dosa"
-	reflect "reflect"
 )
 
 // MockClient is a mock of Client interface
diff --git a/mocks/connector.go b/mocks/connector.go
index 6d0e5b9f..eb665da7 100644
--- a/mocks/connector.go
+++ b/mocks/connector.go
@@ -131,9 +131,9 @@ func (mr *MockConnectorMockRecorder) DropScope(arg0, arg1 interface{}) *gomock.C
 }
 
 // GetEntitySchema mocks base method
-func (m *MockConnector) GetEntitySchema(arg0 context.Context, arg1, arg2, arg3 string, arg4 int32) (*dosa.EntityInfo, error) {
+func (m *MockConnector) GetEntitySchema(arg0 context.Context, arg1, arg2, arg3 string, arg4 int32) (*dosa.EntityDefinition, error) {
 	ret := m.ctrl.Call(m, "GetEntitySchema", arg0, arg1, arg2, arg3, arg4)
-	ret0, _ := ret[0].(*dosa.EntityInfo)
+	ret0, _ := ret[0].(*dosa.EntityDefinition)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }