@@ -7,6 +7,8 @@ package gen
77import (
88 "fmt"
99 "reflect"
10+ "sort"
11+ "strings"
1012
1113 "entgo.io/ent/dialect/gremlin/graph/dsl"
1214 "entgo.io/ent/dialect/sql"
@@ -34,14 +36,15 @@ func (m SchemaMode) Support(mode SchemaMode) bool { return m&mode != 0 }
3436
3537// Storage driver type for codegen.
3638type Storage struct {
37- Name string // storage name.
38- Builder reflect.Type // query builder type.
39- Dialects []string // supported dialects.
40- IdentName string // identifier name (fields and funcs).
41- Imports []string // import packages needed.
42- SchemaMode SchemaMode // schema mode support.
43- Ops func (* Field ) []Op // storage specific operations.
44- OpCode func (Op ) string // operation code for predicates.
39+ Name string // storage name.
40+ Builder reflect.Type // query builder type.
41+ Dialects []string // supported dialects.
42+ IdentName string // identifier name (fields and funcs).
43+ Imports []string // import packages needed.
44+ SchemaMode SchemaMode // schema mode support.
45+ Ops func (* Field ) []Op // storage specific operations.
46+ OpCode func (Op ) string // operation code for predicates.
47+ Init func (* Graph ) error // optional init function.
4548}
4649
4750// StorageDrivers holds the storage driver options for entc.
@@ -66,6 +69,30 @@ var drivers = []*Storage{
6669 return nil
6770 },
6871 OpCode : opCodes (sqlCode [:]),
72+ Init : func (g * Graph ) error {
73+ var with , without []string
74+ for _ , n := range g .Nodes {
75+ if s , err := n .TableSchema (); err == nil && s != "" {
76+ with = append (with , n .Name )
77+ } else {
78+ without = append (without , n .Name )
79+ }
80+ }
81+ switch {
82+ case len (with ) == 0 :
83+ return nil
84+ case len (without ) > 0 :
85+ return fmt .Errorf ("missing schema annotation for %s" , strings .Join (without , ", " ))
86+ default :
87+ if ! g .featureEnabled (FeatureSchemaConfig ) {
88+ g .Features = append (g .Features , FeatureSchemaConfig )
89+ }
90+ if ! g .featureEnabled (featureMultiSchema ) {
91+ g .Features = append (g .Features , featureMultiSchema )
92+ }
93+ return nil
94+ }
95+ },
6996 },
7097 {
7198 Name : "gremlin" ,
@@ -82,6 +109,7 @@ var drivers = []*Storage{
82109 },
83110 SchemaMode : Unique ,
84111 OpCode : opCodes (gremlinCode [:]),
112+ Init : func (* Graph ) error { return nil }, // Noop.
85113 },
86114}
87115
@@ -126,3 +154,51 @@ func opCodes(codes []string) func(Op) string {
126154 return o .Name ()
127155 }
128156}
157+
158+ // TableSchemas returns all table schemas in ent/schema (intentionally exported).
159+ func (g * Graph ) TableSchemas () ([]string , error ) {
160+ all := make (map [string ]struct {})
161+ for _ , n := range g .Nodes {
162+ s , err := n .TableSchema ()
163+ if err != nil {
164+ return nil , err
165+ }
166+ all [s ] = struct {}{}
167+ for _ , e := range n .Edges {
168+ // {{- if and $e.M2M (not $e.Inverse) (not $e.Through) }}
169+ if e .M2M () && ! e .IsInverse () && e .Through == nil {
170+ s , err := e .TableSchema ()
171+ if err != nil {
172+ return nil , err
173+ }
174+ all [s ] = struct {}{}
175+ }
176+ }
177+ }
178+ names := make ([]string , 0 , len (all ))
179+ for s := range all {
180+ names = append (names , s )
181+ }
182+ sort .Strings (names )
183+ return names , nil
184+ }
185+
186+ // TableSchema returns the schema name of where the type table resides (intentionally exported).
187+ func (t * Type ) TableSchema () (string , error ) {
188+ switch ant := t .EntSQL (); {
189+ case ant == nil || ant .Schema == "" :
190+ return "" , fmt .Errorf ("atlas: missing schema annotation for node %q" , t .Name )
191+ default :
192+ return ant .Schema , nil
193+ }
194+ }
195+
196+ // TableSchema returns the schema name of where the type table resides (intentionally exported).
197+ func (e * Edge ) TableSchema () (string , error ) {
198+ switch ant := e .EntSQL (); {
199+ case ant == nil || ant .Schema == "" :
200+ return e .Owner .TableSchema ()
201+ default :
202+ return ant .Schema , nil
203+ }
204+ }
0 commit comments