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