@@ -40,6 +40,8 @@ type OwnerMetadata struct {
4040}
4141
4242// GetOwnerMetadata gets metadata about an owner account
43+ //
44+ // todo: assumes the core mint accounts are always opened first
4345func GetOwnerMetadata (ctx context.Context , data code_data.Provider , owner * Account ) (* OwnerMetadata , error ) {
4446 mtdt := & OwnerMetadata {
4547 Account : owner ,
@@ -82,23 +84,25 @@ func GetOwnerMetadata(ctx context.Context, data code_data.Provider, owner *Accou
8284//
8385// todo: Needs tests here, but most already exist in account service
8486func GetOwnerManagementState (ctx context.Context , data code_data.Provider , owner * Account ) (OwnerManagementState , error ) {
85- recordsByType , err := GetLatestTokenAccountRecordsForOwner (ctx , data , owner )
87+ recordsByMintAndType , err := GetLatestTokenAccountRecordsForOwner (ctx , data , owner )
8688 if err != nil {
8789 return OwnerManagementStateUnknown , err
8890 }
8991
9092 // Has an account ever been opened with the owner? If not, the owner is not a Code account.
9193 // SubmitIntent guarantees all accounts are opened, so there's no need to do anything more
9294 // than an empty check.
93- if len (recordsByType ) == 0 {
95+ if len (recordsByMintAndType ) == 0 {
9496 return OwnerManagementStateNotFound , nil
9597 }
9698
9799 // Are all opened accounts managed by Code? If not, the owner is not a Code account.
98- for _ , batchAccountRecords := range recordsByType {
99- for _ , accountRecords := range batchAccountRecords {
100- if accountRecords .IsTimelock () && ! accountRecords .IsManagedByCode (ctx ) {
101- return OwnerManagementStateUnlocked , nil
100+ for _ , recordsByType := range recordsByMintAndType {
101+ for _ , recordsList := range recordsByType {
102+ for _ , records := range recordsList {
103+ if records .IsTimelock () && ! records .IsManagedByCode (ctx ) {
104+ return OwnerManagementStateUnlocked , nil
105+ }
102106 }
103107 }
104108 }
@@ -108,23 +112,25 @@ func GetOwnerManagementState(ctx context.Context, data code_data.Provider, owner
108112
109113// GetLatestTokenAccountRecordsForOwner gets DB records for the latest set of
110114// token accounts for an owner account.
111- func GetLatestTokenAccountRecordsForOwner (ctx context.Context , data code_data.Provider , owner * Account ) (map [commonpb.AccountType ][]* AccountRecords , error ) {
112- res := make (map [commonpb.AccountType ][]* AccountRecords )
115+ func GetLatestTokenAccountRecordsForOwner (ctx context.Context , data code_data.Provider , owner * Account ) (map [string ] map [ commonpb.AccountType ][]* AccountRecords , error ) {
116+ res := make (map [string ] map [ commonpb.AccountType ][]* AccountRecords )
113117
114- infoRecordsByType , err := data .GetLatestAccountInfosByOwnerAddress (ctx , owner .publicKey .ToBase58 ())
118+ infoRecordsByMintAndType , err := data .GetLatestAccountInfosByOwnerAddress (ctx , owner .publicKey .ToBase58 ())
115119 if err != account .ErrAccountInfoNotFound && err != nil {
116120 return nil , err
117121 }
118122
119- if len (infoRecordsByType ) == 0 {
123+ if len (infoRecordsByMintAndType ) == 0 {
120124 return res , nil
121125 }
122126
123127 var timelockAccounts []string
124- for _ , infoRecords := range infoRecordsByType {
125- for _ , infoRecord := range infoRecords {
126- if infoRecord .IsTimelock () {
127- timelockAccounts = append (timelockAccounts , infoRecord .TokenAccount )
128+ for _ , infoRecordsByType := range infoRecordsByMintAndType {
129+ for _ , infoRecords := range infoRecordsByType {
130+ for _ , infoRecord := range infoRecords {
131+ if infoRecord .IsTimelock () {
132+ timelockAccounts = append (timelockAccounts , infoRecord .TokenAccount )
133+ }
128134 }
129135 }
130136 }
@@ -134,32 +140,38 @@ func GetLatestTokenAccountRecordsForOwner(ctx context.Context, data code_data.Pr
134140 return nil , err
135141 }
136142
137- for _ , generalRecords := range infoRecordsByType {
138- for _ , generalRecord := range generalRecords {
139- var timelockRecord * timelock.Record
140- var ok bool
141- if generalRecord .IsTimelock () {
142- timelockRecord , ok = timelockRecordsByVault [generalRecord .TokenAccount ]
143- if ! ok {
144- return nil , errors .New ("timelock record unexpectedly doesn't exist" )
143+ for _ , infoRecordsByType := range infoRecordsByMintAndType {
144+ for _ , infoRecords := range infoRecordsByType {
145+ for _ , infoRecord := range infoRecords {
146+ var timelockRecord * timelock.Record
147+ var ok bool
148+ if infoRecord .IsTimelock () {
149+ timelockRecord , ok = timelockRecordsByVault [infoRecord .TokenAccount ]
150+ if ! ok {
151+ return nil , errors .New ("timelock record unexpectedly doesn't exist" )
152+ }
145153 }
146- }
147154
148- // Filter out account records for accounts that have completed their
149- // full lifecycle
150- //
151- // todo: This needs tests
152- switch generalRecord .AccountType {
153- case commonpb .AccountType_POOL :
154- if timelockRecord .IsClosed () {
155- continue
155+ // Filter out account records for accounts that have completed their
156+ // full lifecycle
157+ //
158+ // todo: This needs tests
159+ switch infoRecord .AccountType {
160+ case commonpb .AccountType_POOL :
161+ if timelockRecord .IsClosed () {
162+ continue
163+ }
156164 }
157- }
158165
159- res [generalRecord .AccountType ] = append (res [generalRecord .AccountType ], & AccountRecords {
160- General : generalRecord ,
161- Timelock : timelockRecord ,
162- })
166+ if _ , ok := res [infoRecord .MintAccount ]; ! ok {
167+ res [infoRecord .MintAccount ] = make (map [commonpb.AccountType ][]* AccountRecords )
168+ }
169+
170+ res [infoRecord.MintAccount ][infoRecord.AccountType ] = append (res [infoRecord.MintAccount ][infoRecord.AccountType ], & AccountRecords {
171+ General : infoRecord ,
172+ Timelock : timelockRecord ,
173+ })
174+ }
163175 }
164176 }
165177
@@ -168,18 +180,24 @@ func GetLatestTokenAccountRecordsForOwner(ctx context.Context, data code_data.Pr
168180
169181// GetLatestCodeTimelockAccountRecordsForOwner is a utility wrapper over GetLatestTokenAccountRecordsForOwner
170182// that filters for Code Timelock accounts.
171- func GetLatestCodeTimelockAccountRecordsForOwner (ctx context.Context , data code_data.Provider , owner * Account ) (map [commonpb.AccountType ][]* AccountRecords , error ) {
172- res := make (map [commonpb.AccountType ][]* AccountRecords )
183+ func GetLatestCodeTimelockAccountRecordsForOwner (ctx context.Context , data code_data.Provider , owner * Account ) (map [string ] map [ commonpb.AccountType ][]* AccountRecords , error ) {
184+ res := make (map [string ] map [ commonpb.AccountType ][]* AccountRecords )
173185
174- recordsByType , err := GetLatestTokenAccountRecordsForOwner (ctx , data , owner )
186+ recordsByMintAndType , err := GetLatestTokenAccountRecordsForOwner (ctx , data , owner )
175187 if err != nil {
176188 return nil , err
177189 }
178190
179- for _ , recordsList := range recordsByType {
180- for _ , records := range recordsList {
181- if records .IsTimelock () {
182- res [records .General .AccountType ] = append (res [records .General .AccountType ], records )
191+ for _ , recordsByType := range recordsByMintAndType {
192+ for _ , recordsList := range recordsByType {
193+ for _ , records := range recordsList {
194+ if records .IsTimelock () {
195+ if _ , ok := res [records .General .MintAccount ]; ! ok {
196+ res [records .General .MintAccount ] = make (map [commonpb.AccountType ][]* AccountRecords )
197+ }
198+
199+ res [records .General .MintAccount ][records .General .AccountType ] = append (res [records .General .MintAccount ][records .General .AccountType ], records )
200+ }
183201 }
184202 }
185203 }
0 commit comments