111 lines
2.9 KiB
Go
111 lines
2.9 KiB
Go
package localstore
|
|
|
|
import (
|
|
"context"
|
|
"sort"
|
|
"time"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/templates"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/usage"
|
|
)
|
|
|
|
func (s *Store) ListTemplates(_ context.Context, owner string) ([]templates.Template, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
out := []templates.Template{}
|
|
for _, v := range s.templates {
|
|
if v.OwnerID == owner {
|
|
out = append(out, v)
|
|
}
|
|
}
|
|
sort.Slice(out, func(i, j int) bool {
|
|
if out[i].SortOrder != out[j].SortOrder {
|
|
return out[i].SortOrder < out[j].SortOrder
|
|
}
|
|
return out[i].UpdatedAt.After(out[j].UpdatedAt)
|
|
})
|
|
return out, nil
|
|
}
|
|
func (s *Store) CreateTemplate(_ context.Context, v templates.Template) (templates.Template, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if _, ok := s.templates[v.ID]; ok {
|
|
return templates.Template{}, templates.ErrInvalidTemplate
|
|
}
|
|
s.templates[v.ID] = v
|
|
return v, nil
|
|
}
|
|
func (s *Store) UpdateTemplate(_ context.Context, owner, id string, p templates.Patch, now time.Time) (templates.Template, bool, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
v, ok := s.templates[id]
|
|
if !ok || v.OwnerID != owner {
|
|
return templates.Template{}, false, nil
|
|
}
|
|
if p.Name != nil {
|
|
v.Name = *p.Name
|
|
}
|
|
if p.Description != nil {
|
|
v.Description = *p.Description
|
|
}
|
|
if p.Prompt != nil {
|
|
v.Prompt = *p.Prompt
|
|
}
|
|
if p.PreviewImageURL != nil {
|
|
v.PreviewImageURL = *p.PreviewImageURL
|
|
}
|
|
if p.Settings != nil {
|
|
v.Settings = *p.Settings
|
|
}
|
|
if p.SortOrder != nil {
|
|
v.SortOrder = *p.SortOrder
|
|
}
|
|
v.UpdatedAt = now
|
|
s.templates[id] = v
|
|
return v, true, nil
|
|
}
|
|
func (s *Store) DeleteTemplate(_ context.Context, owner, id string) (templates.Template, bool, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
v, ok := s.templates[id]
|
|
if !ok || v.OwnerID != owner {
|
|
return templates.Template{}, false, nil
|
|
}
|
|
delete(s.templates, id)
|
|
return v, true, nil
|
|
}
|
|
|
|
func (s *Store) Insert(e usage.Event) (usage.Event, bool, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if id, ok := s.usageJobIDs[e.JobID]; ok {
|
|
return s.usageEvents[id], false, nil
|
|
}
|
|
if _, ok := s.usageEvents[e.ID]; ok {
|
|
return usage.Event{}, false, nil
|
|
}
|
|
s.usageEvents[e.ID] = e
|
|
if e.JobID != "" {
|
|
s.usageJobIDs[e.JobID] = e.ID
|
|
}
|
|
return e, true, nil
|
|
}
|
|
func (s *Store) List(f usage.Filters) ([]usage.Event, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
out := []usage.Event{}
|
|
for _, e := range s.usageEvents {
|
|
if f.OwnerID != "" && e.OwnerID != f.OwnerID || f.OrganizationID != "" && e.OrganizationID != f.OrganizationID || f.Capability != "" && e.Capability != f.Capability || f.Provider != "" && e.Provider != f.Provider || f.From != "" && e.CreatedAt < f.From || f.To != "" && e.CreatedAt >= f.To {
|
|
continue
|
|
}
|
|
out = append(out, e)
|
|
}
|
|
sort.Slice(out, func(i, j int) bool {
|
|
if out[i].CreatedAt != out[j].CreatedAt {
|
|
return out[i].CreatedAt > out[j].CreatedAt
|
|
}
|
|
return out[i].ID > out[j].ID
|
|
})
|
|
return out, nil
|
|
}
|