github.com/quay/claircore@v1.5.28/datastore/postgres/get.go (about)

     1  package postgres
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strconv"
     7  	"time"
     8  
     9  	"github.com/jackc/pgx/v4"
    10  	"github.com/prometheus/client_golang/prometheus"
    11  	"github.com/prometheus/client_golang/prometheus/promauto"
    12  	"github.com/quay/zlog"
    13  
    14  	"github.com/quay/claircore"
    15  	"github.com/quay/claircore/datastore"
    16  )
    17  
    18  var (
    19  	getVulnerabilitiesCounter = promauto.NewCounterVec(
    20  		prometheus.CounterOpts{
    21  			Namespace: "claircore",
    22  			Subsystem: "vulnstore",
    23  			Name:      "getvulnerabilities_total",
    24  			Help:      "Total number of database queries issued in the get method.",
    25  		},
    26  		[]string{"query"},
    27  	)
    28  	getVulnerabilitiesDuration = promauto.NewHistogramVec(
    29  		prometheus.HistogramOpts{
    30  			Namespace: "claircore",
    31  			Subsystem: "vulnstore",
    32  			Name:      "getvulnerabilities_duration_seconds",
    33  			Help:      "The duration of all queries issued in the get method",
    34  		},
    35  		[]string{"query"},
    36  	)
    37  )
    38  
    39  // Get implements vulnstore.Vulnerability.
    40  func (s *MatcherStore) Get(ctx context.Context, records []*claircore.IndexRecord, opts datastore.GetOpts) (map[string][]*claircore.Vulnerability, error) {
    41  	ctx = zlog.ContextWithValues(ctx, "component", "internal/vulnstore/postgres/Get")
    42  	tx, err := s.pool.Begin(ctx)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  	defer tx.Rollback(ctx)
    47  	// start a batch
    48  	batch := &pgx.Batch{}
    49  	for _, record := range records {
    50  		query, err := buildGetQuery(record, &opts)
    51  		if err != nil {
    52  			// if we cannot build a query for an individual record continue to the next
    53  			zlog.Debug(ctx).
    54  				Err(err).
    55  				Str("record", fmt.Sprintf("%+v", record)).
    56  				Msg("could not build query for record")
    57  			continue
    58  		}
    59  		// queue the select query
    60  		batch.Queue(query)
    61  	}
    62  	// send the batch
    63  
    64  	start := time.Now()
    65  	res := tx.SendBatch(ctx, batch)
    66  	// Can't just defer the close, because the batch must be fully handled
    67  	// before resolving the transaction. Maybe we can move this result handling
    68  	// into its own function to be able to just defer it.
    69  
    70  	// gather all the returned vulns for each queued select statement
    71  	results := make(map[string][]*claircore.Vulnerability)
    72  	vulnSet := make(map[string]map[string]struct{})
    73  	for _, record := range records {
    74  		rows, err := res.Query()
    75  		if err != nil {
    76  			res.Close()
    77  			return nil, err
    78  		}
    79  
    80  		// unpack all returned rows into claircore.Vulnerability structs
    81  		for rows.Next() {
    82  			// fully allocate vuln struct
    83  			v := &claircore.Vulnerability{
    84  				Package: &claircore.Package{},
    85  				Dist:    &claircore.Distribution{},
    86  				Repo:    &claircore.Repository{},
    87  			}
    88  
    89  			var id int64
    90  			err := rows.Scan(
    91  				&id,
    92  				&v.Name,
    93  				&v.Description,
    94  				&v.Issued,
    95  				&v.Links,
    96  				&v.Severity,
    97  				&v.NormalizedSeverity,
    98  				&v.Package.Name,
    99  				&v.Package.Version,
   100  				&v.Package.Module,
   101  				&v.Package.Arch,
   102  				&v.Package.Kind,
   103  				&v.Dist.DID,
   104  				&v.Dist.Name,
   105  				&v.Dist.Version,
   106  				&v.Dist.VersionCodeName,
   107  				&v.Dist.VersionID,
   108  				&v.Dist.Arch,
   109  				&v.Dist.CPE,
   110  				&v.Dist.PrettyName,
   111  				&v.ArchOperation,
   112  				&v.Repo.Name,
   113  				&v.Repo.Key,
   114  				&v.Repo.URI,
   115  				&v.FixedInVersion,
   116  				&v.Updater,
   117  			)
   118  			v.ID = strconv.FormatInt(id, 10)
   119  			if err != nil {
   120  				res.Close()
   121  				return nil, fmt.Errorf("failed to scan vulnerability: %v", err)
   122  			}
   123  
   124  			rid := record.Package.ID
   125  			if _, ok := vulnSet[rid]; !ok {
   126  				vulnSet[rid] = make(map[string]struct{})
   127  			}
   128  			if _, ok := vulnSet[rid][v.ID]; !ok {
   129  				vulnSet[rid][v.ID] = struct{}{}
   130  				results[rid] = append(results[rid], v)
   131  			}
   132  		}
   133  	}
   134  	if err := res.Close(); err != nil {
   135  		return nil, fmt.Errorf("some weird batch error: %v", err)
   136  	}
   137  
   138  	getVulnerabilitiesCounter.WithLabelValues("query_batch").Add(1)
   139  	getVulnerabilitiesDuration.WithLabelValues("query_batch").Observe(time.Since(start).Seconds())
   140  
   141  	if err := tx.Commit(ctx); err != nil {
   142  		return nil, fmt.Errorf("failed to commit tx: %v", err)
   143  	}
   144  	return results, nil
   145  }