blob: 743d13dc651dce71f8a2a70b09b188fb140af733 [file] [log] [blame]
// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package vulndbreqs
import (
"context"
"fmt"
"time"
"cloud.google.com/go/civil"
"golang.org/x/pkgsite-metrics/internal/bigquery"
"golang.org/x/pkgsite-metrics/internal/derrors"
)
const (
// Vuln DB requests live in their own dataset that doesn't vary.
DatasetName = "vulndb"
RequestCountTableName = "requests"
IPRequestCountTableName = "ip-requests"
)
func init() {
s, err := bigquery.InferSchema(RequestCount{})
if err != nil {
panic(err)
}
bigquery.AddTable(RequestCountTableName, s)
s, err = bigquery.InferSchema(IPRequestCount{})
if err != nil {
panic(err)
}
bigquery.AddTable(IPRequestCountTableName, s)
}
// RequestCount holds the number of requests made on a date.
type RequestCount struct {
CreatedAt time.Time `bigquery:"created_at"`
Date civil.Date `bigquery:"date"` // year-month-day without a timezone
Count int `bigquery:"count"`
}
// SetUploadTime is used by Client.Upload.
func (r *RequestCount) SetUploadTime(t time.Time) { r.CreatedAt = t }
// IPRequestCount holds the number of requests for a single IP on a date.
type IPRequestCount struct {
CreatedAt time.Time `bigquery:"created_at"`
Date civil.Date `bigquery:"date"` // year-month-day without a timezone
IP string `bigquery:"ip"` // obfuscated IP address
Count int `bigquery:"count"`
}
// SetUploadTime is used by Client.Upload.
func (r *IPRequestCount) SetUploadTime(t time.Time) { r.CreatedAt = t }
// writeToBigQuery writes request counts to BigQuery.
func writeToBigQuery(ctx context.Context, client *bigquery.Client, rcs []*RequestCount, ircs []*IPRequestCount) (err error) {
defer derrors.Wrap(&err, "vulndbreqs.writeToBigQuery")
if err := client.CreateTable(ctx, RequestCountTableName); err != nil {
return err
}
if err := bigquery.UploadMany(ctx, client, RequestCountTableName, rcs, 0); err != nil {
return err
}
if err := client.CreateTable(ctx, IPRequestCountTableName); err != nil {
return err
}
return bigquery.UploadMany(ctx, client, IPRequestCountTableName, ircs, 0)
}
// ReadRequestCountsFromBigQuery returns daily counts for requests to the vuln DB, most recent first.
func ReadRequestCountsFromBigQuery(ctx context.Context, client *bigquery.Client) (_ []*RequestCount, err error) {
defer derrors.Wrap(&err, "readFromBigQuery")
// Select the most recently inserted row for each date.
q := fmt.Sprintf("(%s) ORDER BY date DESC", bigquery.PartitionQuery{
From: "`" + client.FullTableName(RequestCountTableName) + "`",
PartitionOn: "date",
OrderBy: "created_at DESC",
})
iter, err := client.Query(ctx, q)
if err != nil {
return nil, err
}
return bigquery.All[RequestCount](iter)
}