golang.org/x/build@v0.0.0-20240506185731-218518f32b70/perfdata/cmd/reindex/reindex.go (about)

     1  // Copyright 2017 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Reindex repopulates the perfdata SQL database from the original data files in Google Cloud Storage.
     6  //
     7  // Usage:
     8  //
     9  //	reindex [-v] [-db foo.bar/baz] [-bucket name] prefix...
    10  //
    11  // Reindex reindexes all the uploads with IDs starting with the given prefixes.
    12  
    13  //go:build cloud
    14  
    15  package main
    16  
    17  import (
    18  	"context"
    19  	"flag"
    20  	"fmt"
    21  	"log"
    22  	"os"
    23  	"strings"
    24  
    25  	"cloud.google.com/go/storage"
    26  	_ "github.com/GoogleCloudPlatform/cloudsql-proxy/proxy/dialers/mysql"
    27  	"golang.org/x/build/perfdata/db"
    28  	"golang.org/x/perf/storage/benchfmt"
    29  	"google.golang.org/api/iterator"
    30  )
    31  
    32  var (
    33  	dbName  = flag.String("db", "root:@cloudsql(golang-org:us-central1:golang-org)/perfdata?interpolateParams=true", "connect to MySQL `database`")
    34  	bucket  = flag.String("bucket", "golang-perfdata", "read from Google Cloud Storage `bucket`")
    35  	verbose = flag.Bool("v", false, "verbose")
    36  )
    37  
    38  func usage() {
    39  	fmt.Fprintf(os.Stderr, `Usage of reindex:
    40  	reindex [flags] prefix...
    41  `)
    42  	flag.PrintDefaults()
    43  	os.Exit(2)
    44  }
    45  
    46  func main() {
    47  	log.SetPrefix("reindex: ")
    48  	log.SetFlags(0)
    49  	flag.Usage = usage
    50  	flag.Parse()
    51  	if *verbose {
    52  		log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
    53  	}
    54  
    55  	ctx := context.Background()
    56  
    57  	prefixes := flag.Args()
    58  	if len(prefixes) == 0 {
    59  		log.Fatal("no prefixes to reindex")
    60  	}
    61  
    62  	d, err := db.OpenSQL("mysql", *dbName)
    63  	if err != nil {
    64  		log.Fatal(err)
    65  	}
    66  	defer d.Close()
    67  
    68  	client, err := storage.NewClient(ctx)
    69  	if err != nil {
    70  		log.Fatal(err)
    71  	}
    72  	bucket := client.Bucket(*bucket)
    73  
    74  	for _, prefix := range prefixes {
    75  		if strings.Index(prefix, "/") >= 0 {
    76  			log.Fatalf("prefix %q cannot contain /", prefix)
    77  		}
    78  		it := bucket.Objects(ctx, &storage.Query{Prefix: "uploads/" + prefix})
    79  		var lastUploadId string
    80  		var files []string
    81  		for {
    82  			objAttrs, err := it.Next()
    83  			if err == iterator.Done {
    84  				break
    85  			}
    86  			if err != nil {
    87  				log.Fatal(err)
    88  			}
    89  			name := strings.TrimPrefix(objAttrs.Name, "uploads/")
    90  			slash := strings.Index(name, "/")
    91  			if slash < 0 {
    92  				log.Printf("ignoring file %q", objAttrs.Name)
    93  			}
    94  			uploadID := name[:slash]
    95  			if lastUploadId != "" && uploadID != lastUploadId {
    96  				if err := reindex(ctx, d, bucket, lastUploadId, files); err != nil {
    97  					log.Fatal(err)
    98  				}
    99  				files = nil
   100  			}
   101  			files = append(files, objAttrs.Name)
   102  			lastUploadId = uploadID
   103  		}
   104  		if len(files) > 0 {
   105  			if err := reindex(ctx, d, bucket, lastUploadId, files); err != nil {
   106  				log.Fatal(err)
   107  			}
   108  		}
   109  	}
   110  }
   111  
   112  func reindex(ctx context.Context, db *db.DB, bucket *storage.BucketHandle, uploadID string, files []string) error {
   113  	if *verbose {
   114  		log.Printf("reindexing %q", uploadID)
   115  	}
   116  	u, err := db.ReplaceUpload(uploadID)
   117  	if err != nil {
   118  		return err
   119  	}
   120  	for _, name := range files {
   121  		if err := reindexOne(ctx, u, bucket, name); err != nil {
   122  			return err
   123  		}
   124  	}
   125  	return u.Commit()
   126  }
   127  
   128  func reindexOne(ctx context.Context, u *db.Upload, bucket *storage.BucketHandle, name string) error {
   129  	r, err := bucket.Object(name).NewReader(ctx)
   130  	if err != nil {
   131  		return err
   132  	}
   133  	defer r.Close()
   134  	br := benchfmt.NewReader(r)
   135  	for br.Next() {
   136  		if err := u.InsertRecord(br.Result()); err != nil {
   137  			return err
   138  		}
   139  	}
   140  	if err := br.Err(); err != nil {
   141  		return err
   142  	}
   143  	return nil
   144  }