refactor(go-website): connect list page to database #5839
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,352 @@ | ||
| package datastore | ||
|
|
||
| import ( | ||
| "context" | ||
| "errors" | ||
| "fmt" | ||
| "slices" | ||
| "strings" | ||
| "sync" | ||
| "time" | ||
|
|
||
| "cloud.google.com/go/datastore" | ||
| "github.com/google/osv.dev/go/internal/models" | ||
| "github.com/google/osv.dev/go/logger" | ||
| "github.com/ossf/osv-schema/bindings/go/osvschema" | ||
| "golang.org/x/sync/errgroup" | ||
| "google.golang.org/api/iterator" | ||
| ) | ||
|
|
||
| type VulnerabilitySearchStore struct { | ||
| client *datastore.Client | ||
|
|
||
| mu sync.RWMutex | ||
| cachedCounts []models.EcosystemCount | ||
| lastFetched time.Time | ||
| isRefreshing bool | ||
| } | ||
|
|
||
| var _ models.VulnerabilitySearchStore = (*VulnerabilitySearchStore)(nil) | ||
|
|
||
| func NewVulnerabilitySearchStore(client *datastore.Client) *VulnerabilitySearchStore { | ||
| store := &VulnerabilitySearchStore{client: client} | ||
| // Warm up ecosystem counts in the background so initial HTTP requests don't block | ||
| go func() { | ||
| if _, err := store.refreshEcosystemCounts(context.Background()); err != nil { | ||
| logger.Error("failed to warm up ecosystem counts in background", "error", err) | ||
| } | ||
| }() | ||
|
|
||
| return store | ||
| } | ||
|
|
||
| func (lv *ListedVulnerability) toModel(key *datastore.Key) *models.ListedVulnerability { | ||
| id := "" | ||
| if key != nil { | ||
| id = key.Name | ||
| } else if lv.Key != nil { | ||
| id = lv.Key.Name | ||
| } | ||
|
|
||
| pkgs := make([]models.Package, 0, len(lv.Packages)) | ||
| for _, p := range lv.Packages { | ||
| if strings.Contains(p, "/") { | ||
| // This is quite hacky, but there's not much point fixing it here if we want to migrate away from datastore. | ||
| // We joined the ecosystem name with the package name in datastore (because that's how it's rendered | ||
| // on the website), but this makes changing the format here difficult. Un-join them here to get | ||
| // joined back up on the website frontend side (these should remain separate in a postgres migration). | ||
| eco, name, _ := strings.Cut(p, "/") | ||
| // Check if this actually was a url (and thus a repo) | ||
| cut, _, _ := strings.Cut(eco, ":") // an ecosystem name can legitimately have a dot after the colon for the repository URL | ||
| if strings.Contains(cut, ".") && eco != "crates.io" { | ||
| pkgs = append(pkgs, models.Package{Repo: p}) | ||
| } else { | ||
| // Regular ecosystem/name | ||
| pkgs = append(pkgs, models.Package{ | ||
| Package: &osvschema.Package{ | ||
| Ecosystem: eco, | ||
| Name: name, | ||
| }, | ||
| }) | ||
| } | ||
| } else { | ||
| pkgs = append(pkgs, models.Package{ | ||
| Package: &osvschema.Package{ | ||
| Name: p, | ||
| }, | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| severities := make([]*osvschema.Severity, 0, len(lv.Severities)) | ||
| for _, s := range lv.Severities { | ||
| var sevType osvschema.Severity_Type | ||
| if val, ok := osvschema.Severity_Type_value[s.Type]; ok { | ||
| sevType = osvschema.Severity_Type(val) | ||
| } | ||
| severities = append(severities, &osvschema.Severity{ | ||
| Type: sevType, | ||
| Score: s.Score, | ||
| }) | ||
| } | ||
|
|
||
| return &models.ListedVulnerability{ | ||
| ID: id, | ||
| Published: lv.Published, | ||
| Packages: pkgs, | ||
| Summary: lv.Summary, | ||
| IsFixed: lv.IsFixed, | ||
| Severities: severities, | ||
| } | ||
| } | ||
|
|
||
| func (s *VulnerabilitySearchStore) Search(ctx context.Context, query models.VulnerabilitySearchQuery) (*models.VulnerabilitySearchResult, error) { | ||
| pageSize := query.PageSize | ||
| if pageSize <= 0 { | ||
| pageSize = 16 | ||
| } | ||
|
|
||
| searchString := strings.ToLower(strings.TrimSpace(query.Query)) | ||
| if len(searchString) > 300 { | ||
| return &models.VulnerabilitySearchResult{}, nil | ||
| } | ||
|
|
||
| q := datastore.NewQuery("ListedVulnerability") | ||
| if searchString != "" { | ||
| q = q.FilterField("search_indices", "=", searchString) | ||
| } | ||
| if query.Ecosystem != "" { | ||
| q = q.FilterField("ecosystems", "=", query.Ecosystem) | ||
| } | ||
| q = q.Order("-published").Order("__key__") | ||
|
|
||
| var items []*models.ListedVulnerability | ||
| var nextAfterTime time.Time | ||
| var nextAfterID string | ||
|
|
||
| afterTime := query.AfterTime.Truncate(time.Microsecond) | ||
|
|
||
| if !afterTime.IsZero() { | ||
| // Keyset pagination using AfterTime and AfterID | ||
| q = q.FilterField("published", "<=", afterTime) | ||
| it := s.client.Run(ctx, q) | ||
|
|
||
| for { | ||
| var lv ListedVulnerability | ||
| key, err := it.Next(&lv) | ||
| if errors.Is(err, iterator.Done) { | ||
| break | ||
| } | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to iterate listed vulnerabilities: %w", err) | ||
| } | ||
|
|
||
| // Skip entities at the boundary timestamp that were already seen | ||
| if lv.Published.Truncate(time.Microsecond).Equal(afterTime) && key.Name <= query.AfterID { | ||
| continue | ||
| } | ||
|
|
||
| items = append(items, lv.toModel(key)) | ||
| if len(items) == pageSize { | ||
| // Check if there is another entity after this page | ||
| var peekLv ListedVulnerability | ||
| _, peekErr := it.Next(&peekLv) | ||
| if peekErr == nil { | ||
| nextAfterTime = lv.Published.Truncate(time.Microsecond) | ||
| nextAfterID = key.Name | ||
| } | ||
|
|
||
| break | ||
| } | ||
| } | ||
| } else { | ||
| // First page or legacy offset pagination fallback | ||
| if query.Page > 1 { | ||
| offset := (query.Page - 1) * pageSize | ||
| q = q.Offset(offset) | ||
| } | ||
| q = q.Limit(pageSize + 1) | ||
| var entities []ListedVulnerability | ||
| keys, err := s.client.GetAll(ctx, q, &entities) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to fetch listed vulnerabilities: %w", err) | ||
| } | ||
| hasMore := len(entities) > pageSize | ||
| if hasMore { | ||
| entities = entities[:pageSize] | ||
| } | ||
| for i, e := range entities { | ||
| items = append(items, e.toModel(keys[i])) | ||
| } | ||
| if hasMore && len(items) > 0 { | ||
| lastIdx := len(entities) - 1 | ||
| nextAfterTime = entities[lastIdx].Published.Truncate(time.Microsecond) | ||
| nextAfterID = keys[lastIdx].Name | ||
| } | ||
| } | ||
|
|
||
| // Exact ID match prioritization on page 1 | ||
| if query.AfterTime.IsZero() && query.Page <= 1 && searchString != "" && len(items) > 1 { | ||
| for i, v := range items { | ||
| if strings.EqualFold(v.ID, query.Query) { | ||
| if i > 0 { | ||
| match := items[i] | ||
| copy(items[1:i+1], items[0:i]) | ||
| items[0] = match | ||
| } | ||
|
|
||
| break | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return &models.VulnerabilitySearchResult{ | ||
| Vulnerabilities: items, | ||
| NextAfterTime: nextAfterTime, | ||
| NextAfterID: nextAfterID, | ||
| }, nil | ||
| } | ||
|
|
||
| func (s *VulnerabilitySearchStore) Autocomplete(ctx context.Context, prefix string, limit int) ([]string, error) { | ||
| prefix = strings.ToLower(strings.TrimSpace(prefix)) | ||
| if prefix == "" || len(prefix) > 300 { | ||
| return []string{}, nil | ||
| } | ||
| if limit <= 0 { | ||
| limit = 10 | ||
| } | ||
|
|
||
| q := datastore.NewQuery("ListedVulnerability"). | ||
| FilterField("autocomplete_tags", ">=", prefix). | ||
| FilterField("autocomplete_tags", "<", prefix+"\ufffd"). | ||
| Order("autocomplete_tags"). | ||
| Project("autocomplete_tags"). | ||
| Limit(limit * 5) | ||
|
|
||
| var entities []ListedVulnerability | ||
| if _, err := s.client.GetAll(ctx, q, &entities); err != nil { | ||
| return nil, fmt.Errorf("failed to query autocomplete tags: %w", err) | ||
| } | ||
|
|
||
| tagSet := make(map[string]struct{}) | ||
| for _, e := range entities { | ||
| for _, tag := range e.AutocompleteTags { | ||
| if strings.HasPrefix(strings.ToLower(tag), prefix) { | ||
| tagSet[strings.ToUpper(tag)] = struct{}{} | ||
| } | ||
| } | ||
| } | ||
|
|
||
| suggestions := make([]string, 0, len(tagSet)) | ||
| for tag := range tagSet { | ||
| suggestions = append(suggestions, tag) | ||
| } | ||
| slices.Sort(suggestions) | ||
|
|
||
| if len(suggestions) > limit { | ||
| suggestions = suggestions[:limit] | ||
| } | ||
|
|
||
| return suggestions, nil | ||
| } | ||
|
|
||
| func (s *VulnerabilitySearchStore) EcosystemCounts(ctx context.Context) ([]models.EcosystemCount, error) { | ||
| s.mu.RLock() | ||
| cached := s.cachedCounts | ||
| age := time.Since(s.lastFetched) | ||
| s.mu.RUnlock() | ||
|
|
||
| // 1. Fresh cache: return immediately (< 0.001 ms) | ||
| if len(cached) > 0 && age < 30*time.Minute { | ||
| return cached, nil | ||
| } | ||
|
|
||
| // 2. Stale cache: return stale data immediately and refresh asynchronously in the background | ||
| if len(cached) > 0 { | ||
| s.mu.Lock() | ||
| if !s.isRefreshing { | ||
| s.isRefreshing = true | ||
| bgCtx := context.WithoutCancel(ctx) | ||
| go func() { | ||
| defer func() { | ||
| s.mu.Lock() | ||
| s.isRefreshing = false | ||
| s.mu.Unlock() | ||
| }() | ||
| if _, err := s.refreshEcosystemCounts(bgCtx); err != nil { | ||
| logger.Error("failed to refresh ecosystem counts in background", "error", err) | ||
| } | ||
| }() | ||
| } | ||
| s.mu.Unlock() | ||
|
|
||
| return cached, nil | ||
| } | ||
|
|
||
| // 3. Cold start fallback: fetch synchronously | ||
| return s.refreshEcosystemCounts(ctx) | ||
|
Contributor
There was a problem hiding this comment. This feels like a good place for singleflight, so we don't have a bunch of refreshes all running at the start of the server.
Member
Author
There was a problem hiding this comment. Made the singleflight change in #5840 (to avoid merge conflicts) |
||
| } | ||
|
|
||
| func (s *VulnerabilitySearchStore) refreshEcosystemCounts(ctx context.Context) ([]models.EcosystemCount, error) { | ||
| q := datastore.NewQuery("ListedVulnerability"). | ||
| Project("ecosystems"). | ||
| Distinct() | ||
|
|
||
| var entities []ListedVulnerability | ||
| if _, err := s.client.GetAll(ctx, q, &entities); err != nil { | ||
| return nil, fmt.Errorf("failed to get distinct ecosystems: %w", err) | ||
| } | ||
|
|
||
| ecoSet := make(map[string]struct{}) | ||
| for _, e := range entities { | ||
| for _, eco := range e.Ecosystems { | ||
| if strings.Contains(eco, ":") || eco == "[EMPTY]" || eco == "" { | ||
| continue | ||
| } | ||
| ecoSet[eco] = struct{}{} | ||
| } | ||
| } | ||
|
|
||
| ecosystems := make([]string, 0, len(ecoSet)) | ||
| for eco := range ecoSet { | ||
| ecosystems = append(ecosystems, eco) | ||
| } | ||
| slices.Sort(ecosystems) | ||
|
|
||
| counts := make([]models.EcosystemCount, len(ecosystems)) | ||
| g, gCtx := errgroup.WithContext(ctx) | ||
|
|
||
| for i, eco := range ecosystems { | ||
| g.Go(func() error { | ||
| countQ := datastore.NewQuery("ListedVulnerability").FilterField("ecosystems", "=", eco) | ||
| c, err := s.client.Count(gCtx, countQ) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to count ecosystem %q: %w", eco, err) | ||
| } | ||
| counts[i] = models.EcosystemCount{ | ||
| Name: eco, | ||
| Count: c, | ||
| } | ||
|
|
||
| return nil | ||
| }) | ||
| } | ||
|
|
||
| if err := g.Wait(); err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| result := make([]models.EcosystemCount, 0, len(counts)) | ||
| for _, c := range counts { | ||
| if c.Count > 0 { | ||
| result = append(result, c) | ||
| } | ||
| } | ||
|
|
||
| s.mu.Lock() | ||
| s.cachedCounts = result | ||
| s.lastFetched = time.Now() | ||
| s.mu.Unlock() | ||
|
|
||
| return result, nil | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is it fine to not encode the query anymore? Why was this changed?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please add a comment if this is intentional.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
html/templateactually automatically encodes these, and we don't want to double-encode them.