vitess.io/vitess@v0.16.2/go/vt/mysqlctl/gcsbackupstorage/gcs.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package gcsbackupstorage implements the BackupStorage interface
    18  // for Google Cloud Storage.
    19  package gcsbackupstorage
    20  
    21  import (
    22  	"context"
    23  	"fmt"
    24  	"io"
    25  	"sort"
    26  	"strings"
    27  	"sync"
    28  
    29  	"cloud.google.com/go/storage"
    30  	"github.com/spf13/pflag"
    31  	"golang.org/x/oauth2/google"
    32  	"google.golang.org/api/iterator"
    33  	"google.golang.org/api/option"
    34  
    35  	"vitess.io/vitess/go/trace"
    36  	"vitess.io/vitess/go/vt/concurrency"
    37  	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
    38  	"vitess.io/vitess/go/vt/servenv"
    39  )
    40  
    41  var (
    42  	// bucket is where the backups will go.
    43  	bucket string
    44  
    45  	// root is a prefix added to all object names.
    46  	root string
    47  )
    48  
    49  func registerFlags(fs *pflag.FlagSet) {
    50  	fs.StringVar(&bucket, "gcs_backup_storage_bucket", "", "Google Cloud Storage bucket to use for backups.")
    51  	fs.StringVar(&root, "gcs_backup_storage_root", "", "Root prefix for all backup-related object names.")
    52  }
    53  
    54  func init() {
    55  	servenv.OnParseFor("vtbackup", registerFlags)
    56  	servenv.OnParseFor("vtctl", registerFlags)
    57  	servenv.OnParseFor("vtctld", registerFlags)
    58  	servenv.OnParseFor("vttablet", registerFlags)
    59  }
    60  
    61  // GCSBackupHandle implements BackupHandle for Google Cloud Storage.
    62  type GCSBackupHandle struct {
    63  	client   *storage.Client
    64  	bs       *GCSBackupStorage
    65  	dir      string
    66  	name     string
    67  	readOnly bool
    68  	errors   concurrency.AllErrorRecorder
    69  }
    70  
    71  // RecordError is part of the concurrency.ErrorRecorder interface.
    72  func (bh *GCSBackupHandle) RecordError(err error) {
    73  	bh.errors.RecordError(err)
    74  }
    75  
    76  // HasErrors is part of the concurrency.ErrorRecorder interface.
    77  func (bh *GCSBackupHandle) HasErrors() bool {
    78  	return bh.errors.HasErrors()
    79  }
    80  
    81  // Error is part of the concurrency.ErrorRecorder interface.
    82  func (bh *GCSBackupHandle) Error() error {
    83  	return bh.errors.Error()
    84  }
    85  
    86  // Directory implements BackupHandle.
    87  func (bh *GCSBackupHandle) Directory() string {
    88  	return bh.dir
    89  }
    90  
    91  // Name implements BackupHandle.
    92  func (bh *GCSBackupHandle) Name() string {
    93  	return bh.name
    94  }
    95  
    96  // AddFile implements BackupHandle.
    97  func (bh *GCSBackupHandle) AddFile(ctx context.Context, filename string, filesize int64) (io.WriteCloser, error) {
    98  	if bh.readOnly {
    99  		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
   100  	}
   101  	object := objName(bh.dir, bh.name, filename)
   102  	return bh.client.Bucket(bucket).Object(object).NewWriter(ctx), nil
   103  }
   104  
   105  // EndBackup implements BackupHandle.
   106  func (bh *GCSBackupHandle) EndBackup(ctx context.Context) error {
   107  	if bh.readOnly {
   108  		return fmt.Errorf("EndBackup cannot be called on read-only backup")
   109  	}
   110  	return nil
   111  }
   112  
   113  // AbortBackup implements BackupHandle.
   114  func (bh *GCSBackupHandle) AbortBackup(ctx context.Context) error {
   115  	if bh.readOnly {
   116  		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
   117  	}
   118  	return bh.bs.RemoveBackup(ctx, bh.dir, bh.name)
   119  }
   120  
   121  // ReadFile implements BackupHandle.
   122  func (bh *GCSBackupHandle) ReadFile(ctx context.Context, filename string) (io.ReadCloser, error) {
   123  	if !bh.readOnly {
   124  		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
   125  	}
   126  	object := objName(bh.dir, bh.name, filename)
   127  	return bh.client.Bucket(bucket).Object(object).NewReader(ctx)
   128  }
   129  
   130  // GCSBackupStorage implements BackupStorage for Google Cloud Storage.
   131  type GCSBackupStorage struct {
   132  	// client is the instance of the Google Cloud Storage Go client.
   133  	// Once this field is set, it must not be written again/unset to nil.
   134  	_client *storage.Client
   135  	// mu guards all fields.
   136  	mu sync.Mutex
   137  }
   138  
   139  // ListBackups implements BackupStorage.
   140  func (bs *GCSBackupStorage) ListBackups(ctx context.Context, dir string) ([]backupstorage.BackupHandle, error) {
   141  	c, err := bs.client(ctx)
   142  	if err != nil {
   143  		return nil, err
   144  	}
   145  
   146  	// List prefixes that begin with dir (i.e. list subdirs).
   147  	var subdirs []string
   148  
   149  	var searchPrefix string
   150  	if dir == "/" {
   151  		searchPrefix = ""
   152  	} else {
   153  		searchPrefix = objName(dir, "" /* include trailing slash */)
   154  	}
   155  
   156  	query := &storage.Query{
   157  		Delimiter: "/",
   158  		Prefix:    searchPrefix,
   159  	}
   160  
   161  	it := c.Bucket(bucket).Objects(ctx, query)
   162  	for {
   163  		obj, err := it.Next()
   164  		if err == iterator.Done {
   165  			break
   166  		}
   167  		if err != nil {
   168  			return nil, err
   169  		}
   170  		// Each returned prefix is a subdir.
   171  		// Strip parent dir from full path.
   172  		if obj.Prefix != "" {
   173  			subdir := strings.TrimPrefix(obj.Prefix, searchPrefix)
   174  			subdir = strings.TrimSuffix(subdir, "/")
   175  			subdirs = append(subdirs, subdir)
   176  		}
   177  	}
   178  
   179  	// Backups must be returned in order, oldest first.
   180  	sort.Strings(subdirs)
   181  
   182  	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
   183  	for _, subdir := range subdirs {
   184  		result = append(result, &GCSBackupHandle{
   185  			client:   c,
   186  			bs:       bs,
   187  			dir:      dir,
   188  			name:     subdir,
   189  			readOnly: true,
   190  		})
   191  	}
   192  	return result, nil
   193  }
   194  
   195  // StartBackup implements BackupStorage.
   196  func (bs *GCSBackupStorage) StartBackup(ctx context.Context, dir, name string) (backupstorage.BackupHandle, error) {
   197  	c, err := bs.client(ctx)
   198  	if err != nil {
   199  		return nil, err
   200  	}
   201  
   202  	return &GCSBackupHandle{
   203  		client:   c,
   204  		bs:       bs,
   205  		dir:      dir,
   206  		name:     name,
   207  		readOnly: false,
   208  	}, nil
   209  }
   210  
   211  // RemoveBackup implements BackupStorage.
   212  func (bs *GCSBackupStorage) RemoveBackup(ctx context.Context, dir, name string) error {
   213  	c, err := bs.client(ctx)
   214  	if err != nil {
   215  		return err
   216  	}
   217  
   218  	// Find all objects with the right prefix.
   219  	query := &storage.Query{
   220  		Prefix: objName(dir, name, "" /* include trailing slash */),
   221  	}
   222  	// Delete all the found objects.
   223  	it := c.Bucket(bucket).Objects(ctx, query)
   224  	for {
   225  		obj, err := it.Next()
   226  		if err == iterator.Done {
   227  			break
   228  		}
   229  		if err != nil {
   230  			return err
   231  		}
   232  		if err := c.Bucket(bucket).Object(obj.Name).Delete(ctx); err != nil {
   233  			return fmt.Errorf("unable to delete %q from bucket %q: %v", obj.Name, bucket, err)
   234  		}
   235  	}
   236  	return nil
   237  }
   238  
   239  // Close implements BackupStorage.
   240  func (bs *GCSBackupStorage) Close() error {
   241  	bs.mu.Lock()
   242  	defer bs.mu.Unlock()
   243  
   244  	if bs._client != nil {
   245  		// If client.Close() fails, we still clear bs._client,
   246  		// so we know to create a new client the next time one
   247  		// is needed.
   248  		client := bs._client
   249  		bs._client = nil
   250  		if err := client.Close(); err != nil {
   251  			return err
   252  		}
   253  	}
   254  	return nil
   255  }
   256  
   257  // client returns the GCS Storage client instance.
   258  // If there isn't one yet, it tries to create one.
   259  func (bs *GCSBackupStorage) client(ctx context.Context) (*storage.Client, error) {
   260  	bs.mu.Lock()
   261  	defer bs.mu.Unlock()
   262  
   263  	if bs._client == nil {
   264  		// The context needs to be valid for longer than just
   265  		// the creation context, so we create a new one, but
   266  		// keep the span information.
   267  		ctx = trace.CopySpan(context.Background(), ctx)
   268  		authClient, err := google.DefaultClient(ctx, storage.ScopeFullControl)
   269  		if err != nil {
   270  			return nil, err
   271  		}
   272  		client, err := storage.NewClient(ctx, option.WithHTTPClient(authClient))
   273  		if err != nil {
   274  			return nil, err
   275  		}
   276  		bs._client = client
   277  	}
   278  	return bs._client, nil
   279  }
   280  
   281  // objName joins path parts into an object name.
   282  // Unlike path.Join, it doesn't collapse ".." or strip trailing slashes.
   283  // It also adds the value of the --gcs_backup_storage_root flag if set.
   284  func objName(parts ...string) string {
   285  	if root != "" {
   286  		return root + "/" + strings.Join(parts, "/")
   287  	}
   288  	return strings.Join(parts, "/")
   289  }
   290  
   291  func init() {
   292  	backupstorage.BackupStorageMap["gcs"] = &GCSBackupStorage{}
   293  }