go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/tools/cmd/spannerdataupdater/main.go (about)

     1  // Copyright 2022 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Provides a method for executing long-running Partitioned DML statements,
    16  // as needed to perform data backfills after schema changes to Spanner tables.
    17  //
    18  // Provided because the gcloud command:
    19  //
    20  // gcloud spanner databases execute-sql --enable-partitioned-dml --sql="<SQL>"
    21  //
    22  // does not support long running statements (beyond ~40 mins or so), regardless
    23  // of the timeout paramater specified.
    24  //
    25  // Example usage:
    26  //
    27  //	go run main.go \
    28  //	 -sql='UPDATE TestResults SET <Something> WHERE <Something Else>' \
    29  //	 -project=chops-weetbix-dev -instance=dev -database=chops-weetbix-dev
    30  package main
    31  
    32  import (
    33  	"context"
    34  	"flag"
    35  	"fmt"
    36  	"log"
    37  	"os"
    38  	"strings"
    39  	"time"
    40  
    41  	"cloud.google.com/go/spanner"
    42  	"google.golang.org/api/option"
    43  
    44  	"go.chromium.org/luci/auth"
    45  	"go.chromium.org/luci/common/errors"
    46  	"go.chromium.org/luci/common/logging/gologger"
    47  	"go.chromium.org/luci/hardcoded/chromeinfra"
    48  )
    49  
    50  var errCanceledByUser = errors.Reason("operation canceled by user").Err()
    51  
    52  type flags struct {
    53  	sql      string
    54  	project  string
    55  	instance string
    56  	database string
    57  }
    58  
    59  func parseFlags() (*flags, error) {
    60  	var f flags
    61  	flag.StringVar(&f.sql, "sql", "", "SQL Statement to run. This must satisfy the requirements of a parititoned DML statement: https://cloud.google.com/spanner/docs/dml-partitioned")
    62  	flag.StringVar(&f.project, "project", "", "GCP Project, e.g. luci-resultdb or chops-weetbix.")
    63  	flag.StringVar(&f.instance, "instance", "", "The Spanner instance name, e.g. dev or prod.")
    64  	flag.StringVar(&f.database, "database", "", "The Spanner database name.")
    65  	flag.Parse()
    66  
    67  	switch {
    68  	case len(flag.Args()) > 0:
    69  		return nil, fmt.Errorf("unexpected arguments: %q", flag.Args())
    70  	case f.sql == "":
    71  		return nil, fmt.Errorf("-sql is required")
    72  	case f.project == "":
    73  		return nil, fmt.Errorf("-project is required")
    74  	case f.instance == "":
    75  		return nil, fmt.Errorf("-instance is required")
    76  	case f.database == "":
    77  		return nil, fmt.Errorf("-database is required")
    78  	}
    79  	return &f, nil
    80  }
    81  
    82  func run(ctx context.Context) error {
    83  	flags, err := parseFlags()
    84  	if err != nil {
    85  		return errors.Annotate(err, "failed to parse flags").Err()
    86  	}
    87  
    88  	// Create an Authenticator and use it for Spanner operations.
    89  	authOpts := chromeinfra.DefaultAuthOptions()
    90  	authOpts.Scopes = []string{
    91  		"https://www.googleapis.com/auth/cloud-platform",
    92  		"https://www.googleapis.com/auth/userinfo.email",
    93  	}
    94  	authenticator := auth.NewAuthenticator(ctx, auth.InteractiveLogin, authOpts)
    95  
    96  	authTS, err := authenticator.TokenSource()
    97  	if err != nil {
    98  		return errors.Annotate(err, "could not get authentication credentials").Err()
    99  	}
   100  
   101  	database := fmt.Sprintf("projects/%s/instances/%s/databases/%s", flags.project, flags.instance, flags.database)
   102  	c, err := spanner.NewClient(ctx, database, option.WithTokenSource(authTS))
   103  	if err != nil {
   104  		return errors.Annotate(err, "could not create Spanner client").Err()
   105  	}
   106  	fmt.Println("The following partitioned DML will be applied.")
   107  	fmt.Println(strings.Repeat("=", 80))
   108  	fmt.Println("GCP Project:   ", flags.project)
   109  	fmt.Println("Instance:      ", flags.instance)
   110  	fmt.Println("Database:      ", flags.database)
   111  	fmt.Println("SQL Statement: ", flags.sql)
   112  	fmt.Println(strings.Repeat("=", 80))
   113  
   114  	confirm := confirm("Continue")
   115  	if !confirm {
   116  		return errCanceledByUser
   117  	}
   118  
   119  	// Apply one-week timeout.
   120  	ctx, cancel := context.WithTimeout(ctx, time.Hour*7*24)
   121  	defer cancel()
   122  	count, err := c.PartitionedUpdate(ctx, spanner.NewStatement(flags.sql))
   123  	if err != nil {
   124  		return errors.Annotate(err, "applying partitioned update").Err()
   125  	}
   126  	log.Printf("%v rows updated", count)
   127  	return nil
   128  }
   129  
   130  func main() {
   131  	ctx := gologger.StdConfig.Use(context.Background())
   132  	switch err := run(ctx); {
   133  	case errCanceledByUser == err:
   134  		os.Exit(1)
   135  	case err != nil:
   136  		log.Fatal(err)
   137  	}
   138  }
   139  
   140  // confirm asks for a user confirmation for an action, with No as default.
   141  // Only "y" or "Y" responses is treated as yes.
   142  func confirm(action string) (response bool) {
   143  	fmt.Printf("%s? [y/N] ", action)
   144  	var res string
   145  	fmt.Scanln(&res)
   146  	return res == "y" || res == "Y"
   147  }