cmd/ejobs: support min/max import semantics for module filtering

Change-Id: I008b092499705b20be9dfdc42a6d0a1a920ab57a
Reviewed-on: https://go-review.googlesource.com/c/pkgsite-metrics/+/679355
Reviewed-by: Jonathan Amsterdam <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: thatnealpatel

cmd/ejobs/main.go

@@ -47,6 +47,7 @@ var (
 
 var (
 	minImporters int           // for start
+	maxImporters int           // for start
 	noDeps       bool          // for start
 	moduleFile   string        // for start
 	waitInterval time.Duration // for wait
@@ -69,8 +70,10 @@ var commands = []command{
 		"start a job",
 		doStart,
 		func(fs *flag.FlagSet) {
-			fs.IntVar(&minImporters, "min", -1,
+			fs.IntVar(&maxImporters, "min", -1,
 				"run on modules with at least this many importers (<0: use server default of 10)")
+			fs.IntVar(&minImporters, "max", -1,
+				"run on modules with at most this many importers (<0: use server default of math.MaxInt)")
 			fs.StringVar(&moduleFile, "file", "",
 				"file with modules to use: each line is MODULE_PATH VERSION NUM_IMPORTERS")
 			fs.BoolVar(&noDeps, "nodeps", false, "do not download dependencies for modules")
@@ -342,6 +345,9 @@ func doStart(ctx context.Context, args []string) error {
 	if minImporters >= 0 {
 		u += fmt.Sprintf("&min=%d", minImporters)
 	}
+	if maxImporters >= 0 {
+		u += fmt.Sprintf("&max=%d", maxImporters)
+	}
 	if gcsPath != "" {
 		gurl := "gs://" + gcsPath
 		u += fmt.Sprintf("&file=%s", url.QueryEscape(gurl))

internal/analysis/analysis.go

@@ -45,6 +45,7 @@ type EnqueueParams struct {
 	Args     string // command-line arguments to binary; split on whitespace
 	Insecure bool   // if true, run outside sandbox
 	Min      int    // minimum import-by count for a module to be included
+	Max      int    // maximum import-by count for a module to be included
 	File     string // path to file containing modules; if missing, use DB
 	Suffix   string // appended to task queue IDs to generate unique tasks
 	User     string // user initiating enqueue

internal/pkgsitedb/db.go

@@ -52,15 +52,15 @@ func redactPassword(dbinfo string) string {
 // ModuleSpecs retrieves all modules that contain packages that are
 // imported by minImportedByCount or more packages.
 // It looks for the information in the search_documents table of the given pkgsite DB.
-func ModuleSpecs(ctx context.Context, db *sql.DB, minImportedByCount int) (specs []scan.ModuleSpec, err error) {
+func ModuleSpecs(ctx context.Context, db *sql.DB, minImports, maxImports int) (specs []scan.ModuleSpec, err error) {
 	defer derrors.Wrap(&err, "moduleSpecsFromDB")
 	query := `
 		SELECT module_path, version, max(imported_by_count)
 		FROM search_documents
 		GROUP BY module_path, version
-		HAVING max(imported_by_count) >= $1
-		ORDER by max(imported_by_count) desc`
-	rows, err := db.QueryContext(ctx, query, minImportedByCount)
+		HAVING max(imported_by_count) >= $1 AND max(imported_by_count) <= $2
+		ORDER BY max(imported_by_count) desc`
+	rows, err := db.QueryContext(ctx, query, minImports, maxImports)
 	if err != nil {
 		return nil, err
 	}

internal/pkgsitedb/db_test.go

@@ -12,6 +12,7 @@ import (
 	"database/sql"
 	"flag"
 	"fmt"
+	"math"
 	"net/url"
 	"strings"
 	"testing"
@@ -50,7 +51,7 @@ func TestModuleSpecs(t *testing.T) {
 	if err := db.PingContext(ctx); err != nil {
 		t.Fatal(err)
 	}
-	got, err := ModuleSpecs(ctx, db, 1000)
+	got, err := ModuleSpecs(ctx, db, 1000, math.MaxInt)
 	if err != nil {
 		t.Fatal(err)
 	}

internal/scan/parse.go

@@ -35,7 +35,7 @@ type ModuleSpec struct {
 	ImportedBy    int
 }
 
-func ParseCorpusFile(filename string, minImportedByCount int) (ms []ModuleSpec, err error) {
+func ParseCorpusFile(filename string, minImports, maxImports int) (ms []ModuleSpec, err error) {
 	defer derrors.Wrap(&err, "ParseCorpusFile(%q)", filename)
 	lines, err := ReadFileLines(filename)
 	if err != nil {
@@ -60,7 +60,7 @@ func ParseCorpusFile(filename string, minImportedByCount int) (ms []ModuleSpec,
 		if err != nil {
 			return nil, fmt.Errorf("%v on line %q", err, line)
 		}
-		if n >= minImportedByCount {
+		if minImports <= n && n <= maxImports {
 			ms = append(ms, ModuleSpec{Path: path, Version: vers, ImportedBy: n})
 		}
 	}

internal/scan/parse_test.go

@@ -6,6 +6,7 @@ package scan
 
 import (
 	"flag"
+	"math"
 	"net/http"
 	"reflect"
 	"slices"
@@ -100,7 +101,7 @@ func TestModuleURLPathError(t *testing.T) {
 
 func TestParseCorpusFile(t *testing.T) {
 	const file = "testdata/modules.txt"
-	got, err := ParseCorpusFile(file, 1)
+	got, err := ParseCorpusFile(file, 1, math.MaxInt)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -114,7 +115,7 @@ func TestParseCorpusFile(t *testing.T) {
 		t.Errorf("\n got %v\nwant %v", got, want)
 	}
 
-	got, err = ParseCorpusFile(file, 10)
+	got, err = ParseCorpusFile(file, 10, math.MaxInt)
 	if err != nil {
 		t.Fatal(err)
 	}

internal/worker/analysis.go

@@ -426,10 +426,13 @@ func readSource(file string, line int, nContext int) (_ string, err error) {
 func (s *analysisServer) handleEnqueue(w http.ResponseWriter, r *http.Request) (err error) {
 	defer derrors.Wrap(&err, "analysisServer.handleEnqueue")
 	ctx := r.Context()
-	params := &analysis.EnqueueParams{Min: defaultMinImportedByCount}
+	params := &analysis.EnqueueParams{Min: defaultMinImportedByCount, Max: defaultMaxImportedByCount}
 	if err := scan.ParseParams(r, params); err != nil {
 		return fmt.Errorf("%w: %v", derrors.InvalidArgument, err)
 	}
+	if params.Min > params.Max {
+		return fmt.Errorf("%w: analysis: bad min/max range", derrors.InvalidArgument)
+	}
 	if params.Binary == "" {
 		return fmt.Errorf("%w: analysis: missing binary", derrors.InvalidArgument)
 	}
@@ -446,7 +449,7 @@ func (s *analysisServer) handleEnqueue(w http.ResponseWriter, r *http.Request) (
 	if err != nil {
 		return err
 	}
-	mods, err := readModules(ctx, s.cfg, params.File, params.Min)
+	mods, err := readModules(ctx, s.cfg, params.File, params.Min, params.Max)
 	if err != nil {
 		return err
 	}

internal/worker/enqueue.go

@@ -6,6 +6,7 @@ package worker
 
 import (
 	"context"
+	"math"
 	"sync"
 
 	"golang.org/x/pkgsite-metrics/internal/config"
@@ -16,24 +17,27 @@ import (
 	"golang.org/x/pkgsite-metrics/internal/scan"
 )
 
-const defaultMinImportedByCount = 10
+const (
+	defaultMinImportedByCount = 10
+	defaultMaxImportedByCount = math.MaxInt
+)
 
-func readModules(ctx context.Context, cfg *config.Config, file string, minImpCount int) ([]scan.ModuleSpec, error) {
+func readModules(ctx context.Context, cfg *config.Config, file string, minImports, maxImports int) ([]scan.ModuleSpec, error) {
 	if file != "" {
 		log.Infof(ctx, "reading modules from file %s", file)
-		return scan.ParseCorpusFile(file, minImpCount)
+		return scan.ParseCorpusFile(file, minImports, maxImports)
 	}
 	log.Infof(ctx, "reading modules from DB %s", cfg.PkgsiteDBName)
-	return readFromDB(ctx, cfg, minImpCount)
+	return readFromDB(ctx, cfg, minImports, maxImports)
 }
 
-func readFromDB(ctx context.Context, cfg *config.Config, minImportedByCount int) ([]scan.ModuleSpec, error) {
+func readFromDB(ctx context.Context, cfg *config.Config, minImports, maxImports int) ([]scan.ModuleSpec, error) {
 	db, err := pkgsitedb.Open(ctx, cfg)
 	if err != nil {
 		return nil, err
 	}
 	defer db.Close()
-	return pkgsitedb.ModuleSpecs(ctx, db, minImportedByCount)
+	return pkgsitedb.ModuleSpecs(ctx, db, minImports, maxImports)
 }
 
 func enqueueTasks(ctx context.Context, tasks []queue.Task, q queue.Queue, opts *queue.Options) (err error) {

internal/worker/govulncheck_enqueue.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"net/http"
 	"sort"
 	"strings"
@@ -80,7 +81,7 @@ func createGovulncheckQueueTasks(ctx context.Context, cfg *config.Config, params
 	)
 	for _, mode := range modes {
 		if modspecs == nil {
-			modspecs, err = readModules(ctx, cfg, params.File, params.Min)
+			modspecs, err = readModules(ctx, cfg, params.File, params.Min, math.MaxInt)
 			if err != nil {
 				return nil, err
 			}