src/pkg/gcs/gcs_client.go - cos/tools - Git at Google

 // Copyright 2021 Google Inc. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package gcs

 import (
 	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"net/url"
 	"os"
 	"strings"

 	"cloud.google.com/go/storage"
 	"github.com/golang/glog"
 	"google.golang.org/api/iterator"
 )

 const schemeGCS = "gs"
 const schemeHttps = "https"
 const gcsHostName = "storage.googleapis.com"

 func readGCSObject(ctx context.Context, gcsClient *storage.Client, inputURL string) (*storage.Reader, error) {
 	gcsBucket, name, err := getGCSVariables(inputURL)
 	if err != nil {
 		glog.Errorf("Unable to get gcs variables - gcsBucket, gcsPath with the input url: %s", inputURL)
 		return nil, err
 	}
 	return gcsClient.Bucket(gcsBucket).Object(name).NewReader(ctx)
 }

 // DownloadGCSObjectFromURL downloads the object at inputURL and saves it at destinationPath
 func DownloadGCSObjectFromURL(ctx context.Context, gcsClient *storage.Client, inputURL, destinationPath string) error {
 	// GCS Client Reader.
 	gcsBucket, objectName, err := getGCSVariables(inputURL)
 	if err != nil {
 		glog.Errorf("Unable to get gcs variables - gcsBucket, gcsPath with the input url: %s", inputURL)
 		return err
 	}
 	if err := DownloadGCSObject(ctx, gcsClient, gcsBucket, objectName, destinationPath); err != nil {
 		return err
 	}
 	glog.Infof("Sucessfully downloaded gcs object from %s to %s.", inputURL, destinationPath)
 	return nil
 }

 // DownloadGCSObject downloads the object at bucket, objectName and saves it at destinationPath
 func DownloadGCSObject(ctx context.Context, gcsClient *storage.Client, gcsBucket, objectName, destinationPath string) error {
 	glog.Infof("Downloading gcs object from bucket: %s, object: %s to %s.", gcsBucket, objectName, destinationPath)
 	rc, err := gcsClient.Bucket(gcsBucket).Object(objectName).NewReader(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to create the reader from GCS client: %w", err)
 	}
 	defer rc.Close()
 	// Create output file.
 	destinationFile, err := os.Create(destinationPath)
 	if err != nil {
 		return fmt.Errorf("failed to create output file: %s with error: %w", destinationPath, err)
 	}
 	defer destinationFile.Close()
 	// Write the content from storage reader to output file.
 	if _, err := io.Copy(destinationFile, rc); err != nil {
 		destinationFile.Close()
 		os.Remove(destinationPath)
 		return fmt.Errorf("failed to copy the content from GCS Reader to outputFile: %s when downloading from bucket: %s, object: %s with error: %w", destinationPath, gcsBucket, objectName, err)
 	}
 	glog.Infof("Sucessfully downloaded gcs object from bucket: %s, object: %s to %s.", gcsBucket, objectName, destinationPath)
 	return nil
 }

 // DownloadGCSObjectString downloads the object at inputURL and saves the contents of the object file to a string
 func DownloadGCSObjectString(ctx context.Context,
 	gcsClient *storage.Client, inputURL string) (string, error) {
 	r, err := readGCSObject(ctx, gcsClient, inputURL)
 	if err != nil {
 		return "", err
 	}
 	defer r.Close()

 	ret, err := ioutil.ReadAll(r)
 	if err != nil {
 		return "", fmt.Errorf("unable to read file from gcs bucket: %v", err)
 	}
 	return string(ret), nil
 }

 // GCSObjectExistsFromURL checks if an object exists at inputURL
 func GCSObjectExistsFromURL(ctx context.Context, gcsClient *storage.Client, inputURL string) (bool, error) {
 	gcsBucket, objectName, err := getGCSVariables(inputURL)
 	if err != nil {
 		glog.Errorf("Unable to get gcs variables - gcsBucket, gcsPath with the input url: %s", inputURL)
 		return false, err
 	}
 	return GCSObjectExists(ctx, gcsClient, gcsBucket, objectName)
 }

 // GCSObjectExists checks if an object with objectName exists at gcsBucket objectName.
 func GCSObjectExists(ctx context.Context, gcsClient *storage.Client, gcsBucket, objectName string) (bool, error) {
 	_, err := gcsClient.Bucket(gcsBucket).Object(objectName).Attrs(ctx)
 	if err == storage.ErrObjectNotExist {
 		glog.Infof("Can not find object: %s from bucket: %s", objectName, gcsBucket)
 		return false, nil
 	}
 	if err != nil {
 		glog.Errorf("Fail to find object: %s from bucket: %s with error: %v", objectName, gcsBucket, err)
 		return false, err
 	}
 	glog.Infof("Successfully find object: %s from bucket: %s", objectName, gcsBucket)
 	return true, nil
 }

 // ListGCSBucket lists all the objectNames in gcsBucket with prefix.
 func ListGCSBucket(ctx context.Context, gcsClient *storage.Client, gcsBucket, prefix string) ([]string, error) {
 	var objects []string
 	query := &storage.Query{Prefix: prefix}
 	query.SetAttrSelection([]string{"Name"})
 	it := gcsClient.Bucket(gcsBucket).Objects(ctx, query)
 	for {
 		attrs, err := it.Next()
 		if err == iterator.Done {
 			break
 		}
 		if err != nil {
 			glog.Errorf("Error when listing bucket in bucket: %s, prefix: %s with error: %v", gcsBucket, prefix, err)
 			return nil, fmt.Errorf("error when listing bucket in bucket: %s, prefix: %s with error: %w", gcsBucket, prefix, err)
 		}
 		objects = append(objects, attrs.Name)
 	}
 	return objects, nil
 }

 // UploadGCSObject uploads an object at inputPath to destination URL
 func UploadGCSObject(ctx context.Context, gcsClient *storage.Client, inputPath, destinationURL string) error {
 	fileReader, err := os.Open(inputPath)
 	if err != nil {
 		return err
 	}
 	return uploadGCSObject(ctx, gcsClient, fileReader, destinationURL)
 }

 // UploadGCSObjectString uploads an input string as a file to destination URL
 func UploadGCSObjectString(ctx context.Context, gcsClient *storage.Client, inputStr, destinationURL string) error {
 	reader := strings.NewReader(inputStr)
 	return uploadGCSObject(ctx, gcsClient, reader, destinationURL)
 }

 func uploadGCSObject(ctx context.Context,
 	gcsClient *storage.Client, reader io.Reader, destinationURL string) error {
 	gcsBucket, name, err := getGCSVariables(destinationURL)
 	if err != nil {
 		return fmt.Errorf("error parsing destination URL: %v", err)
 	}
 	w := gcsClient.Bucket(gcsBucket).Object(name).NewWriter(ctx)
 	if _, err := io.Copy(w, reader); err != nil {
 		return err
 	}
 	if err := w.Close(); err != nil {
 		return err
 	}
 	return nil
 }

 // DeleteGCSObject deletes an object at the input URL
 func DeleteGCSObject(ctx context.Context,
 	gcsClient *storage.Client, inputURL string) error {
 	gcsBucket, name, err := getGCSVariables(inputURL)
 	if err != nil {
 		return fmt.Errorf("error parsing input URL: %v", err)
 	}
 	return gcsClient.Bucket(gcsBucket).Object(name).Delete(ctx)
 }

 // getGCSVariables returns the GCSBucket, GCSPath, errorMsg) based on the gcsPath url input.
 func getGCSVariables(gcsPath string) (string, string, error) {
 	parsedURL, err := url.Parse(gcsPath)
 	if err != nil {
 		glog.Errorf("Error - %v: Failed to parse gcsPath - %s", gcsPath, err)
 		return "", "", err
 	}
 	if parsedURL.Scheme == schemeGCS {
 		// In this case, the gcsPath looks like: gs://cos-tools/16108.403.42/kernel-headers.tgz
 		// url.Hostname() returns "cos-tools" which is the GCSBucket
 		// url.EscapedPath() returns "/16108.403.42/kernel-headers.tgz" with a leading "/".
 		return parsedURL.Hostname(), strings.TrimLeft(parsedURL.EscapedPath(), "/"), nil
 	} else if parsedURL.Scheme == schemeHttps && strings.ToLower(parsedURL.Host) == gcsHostName {
 		// In this case, the gcsPath looks like: https://storage.googleapis.com/cos-tools/1700.00.0/lakitu/gpu_default_version
 		// Split the path into parts using "/"
 		pathParts := strings.Split(parsedURL.Path, "/")
 		// Check if there are enough parts in the path
 		if len(pathParts) < 3 {
 			glog.Errorf("Error: Invalid gcsPath input: %s - the number of path parts separated by '/' is smaller than 3.", gcsPath)
 			return "", "", fmt.Errorf("error: invalid gcs_path input: %s - the number of path parts separated by '/' is smaller than 3", gcsPath)
 		}
 		gcsBucket := pathParts[1]
 		gcsPath := strings.Join(pathParts[2:], "/")
 		return gcsBucket, gcsPath, nil
 	} else {
 		glog.Errorf("Invalid input gcsPath: %s - unable to parse it to get gcsBucket and gcsPath", gcsPath)
 		return "", "", fmt.Errorf("error: invalid input gcs_path: %s - unable to parse it to get gcsBucket and gcsPath", gcsPath)
 	}
 }
	// Copyright 2021 Google Inc. All Rights Reserved.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package gcs

	import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"strings"

	"cloud.google.com/go/storage"
	"github.com/golang/glog"
	"google.golang.org/api/iterator"
	)

	const schemeGCS = "gs"
	const schemeHttps = "https"
	const gcsHostName = "storage.googleapis.com"

	func readGCSObject(ctx context.Context, gcsClient storage.Client, inputURL string) (storage.Reader, error) {
	gcsBucket, name, err := getGCSVariables(inputURL)
	if err != nil {
	glog.Errorf("Unable to get gcs variables - gcsBucket, gcsPath with the input url: %s", inputURL)
	return nil, err
	}
	return gcsClient.Bucket(gcsBucket).Object(name).NewReader(ctx)
	}

	// DownloadGCSObjectFromURL downloads the object at inputURL and saves it at destinationPath
	func DownloadGCSObjectFromURL(ctx context.Context, gcsClient *storage.Client, inputURL, destinationPath string) error {
	// GCS Client Reader.
	gcsBucket, objectName, err := getGCSVariables(inputURL)
	if err != nil {
	glog.Errorf("Unable to get gcs variables - gcsBucket, gcsPath with the input url: %s", inputURL)
	return err
	}
	if err := DownloadGCSObject(ctx, gcsClient, gcsBucket, objectName, destinationPath); err != nil {
	return err
	}
	glog.Infof("Sucessfully downloaded gcs object from %s to %s.", inputURL, destinationPath)
	return nil
	}

	// DownloadGCSObject downloads the object at bucket, objectName and saves it at destinationPath
	func DownloadGCSObject(ctx context.Context, gcsClient *storage.Client, gcsBucket, objectName, destinationPath string) error {
	glog.Infof("Downloading gcs object from bucket: %s, object: %s to %s.", gcsBucket, objectName, destinationPath)
	rc, err := gcsClient.Bucket(gcsBucket).Object(objectName).NewReader(ctx)
	if err != nil {
	return fmt.Errorf("failed to create the reader from GCS client: %w", err)
	}
	defer rc.Close()
	// Create output file.
	destinationFile, err := os.Create(destinationPath)
	if err != nil {
	return fmt.Errorf("failed to create output file: %s with error: %w", destinationPath, err)
	}
	defer destinationFile.Close()
	// Write the content from storage reader to output file.
	if _, err := io.Copy(destinationFile, rc); err != nil {
	destinationFile.Close()
	os.Remove(destinationPath)
	return fmt.Errorf("failed to copy the content from GCS Reader to outputFile: %s when downloading from bucket: %s, object: %s with error: %w", destinationPath, gcsBucket, objectName, err)
	}
	glog.Infof("Sucessfully downloaded gcs object from bucket: %s, object: %s to %s.", gcsBucket, objectName, destinationPath)
	return nil
	}

	// DownloadGCSObjectString downloads the object at inputURL and saves the contents of the object file to a string
	func DownloadGCSObjectString(ctx context.Context,
	gcsClient *storage.Client, inputURL string) (string, error) {
	r, err := readGCSObject(ctx, gcsClient, inputURL)
	if err != nil {
	return "", err
	}
	defer r.Close()

	ret, err := ioutil.ReadAll(r)
	if err != nil {
	return "", fmt.Errorf("unable to read file from gcs bucket: %v", err)
	}
	return string(ret), nil
	}

	// GCSObjectExistsFromURL checks if an object exists at inputURL
	func GCSObjectExistsFromURL(ctx context.Context, gcsClient *storage.Client, inputURL string) (bool, error) {
	gcsBucket, objectName, err := getGCSVariables(inputURL)
	if err != nil {
	glog.Errorf("Unable to get gcs variables - gcsBucket, gcsPath with the input url: %s", inputURL)
	return false, err
	}
	return GCSObjectExists(ctx, gcsClient, gcsBucket, objectName)
	}

	// GCSObjectExists checks if an object with objectName exists at gcsBucket objectName.
	func GCSObjectExists(ctx context.Context, gcsClient *storage.Client, gcsBucket, objectName string) (bool, error) {
	_, err := gcsClient.Bucket(gcsBucket).Object(objectName).Attrs(ctx)
	if err == storage.ErrObjectNotExist {
	glog.Infof("Can not find object: %s from bucket: %s", objectName, gcsBucket)
	return false, nil
	}
	if err != nil {
	glog.Errorf("Fail to find object: %s from bucket: %s with error: %v", objectName, gcsBucket, err)
	return false, err
	}
	glog.Infof("Successfully find object: %s from bucket: %s", objectName, gcsBucket)
	return true, nil
	}

	// ListGCSBucket lists all the objectNames in gcsBucket with prefix.
	func ListGCSBucket(ctx context.Context, gcsClient *storage.Client, gcsBucket, prefix string) ([]string, error) {
	var objects []string
	query := &storage.Query{Prefix: prefix}
	query.SetAttrSelection([]string{"Name"})
	it := gcsClient.Bucket(gcsBucket).Objects(ctx, query)
	for {
	attrs, err := it.Next()
	if err == iterator.Done {
	break
	}
	if err != nil {
	glog.Errorf("Error when listing bucket in bucket: %s, prefix: %s with error: %v", gcsBucket, prefix, err)
	return nil, fmt.Errorf("error when listing bucket in bucket: %s, prefix: %s with error: %w", gcsBucket, prefix, err)
	}
	objects = append(objects, attrs.Name)
	}
	return objects, nil
	}

	// UploadGCSObject uploads an object at inputPath to destination URL
	func UploadGCSObject(ctx context.Context, gcsClient *storage.Client, inputPath, destinationURL string) error {
	fileReader, err := os.Open(inputPath)
	if err != nil {
	return err
	}
	return uploadGCSObject(ctx, gcsClient, fileReader, destinationURL)
	}

	// UploadGCSObjectString uploads an input string as a file to destination URL
	func UploadGCSObjectString(ctx context.Context, gcsClient *storage.Client, inputStr, destinationURL string) error {
	reader := strings.NewReader(inputStr)
	return uploadGCSObject(ctx, gcsClient, reader, destinationURL)
	}

	func uploadGCSObject(ctx context.Context,
	gcsClient *storage.Client, reader io.Reader, destinationURL string) error {
	gcsBucket, name, err := getGCSVariables(destinationURL)
	if err != nil {
	return fmt.Errorf("error parsing destination URL: %v", err)
	}
	w := gcsClient.Bucket(gcsBucket).Object(name).NewWriter(ctx)
	if _, err := io.Copy(w, reader); err != nil {
	return err
	}
	if err := w.Close(); err != nil {
	return err
	}
	return nil
	}

	// DeleteGCSObject deletes an object at the input URL
	func DeleteGCSObject(ctx context.Context,
	gcsClient *storage.Client, inputURL string) error {
	gcsBucket, name, err := getGCSVariables(inputURL)
	if err != nil {
	return fmt.Errorf("error parsing input URL: %v", err)
	}
	return gcsClient.Bucket(gcsBucket).Object(name).Delete(ctx)
	}

	// getGCSVariables returns the GCSBucket, GCSPath, errorMsg) based on the gcsPath url input.
	func getGCSVariables(gcsPath string) (string, string, error) {
	parsedURL, err := url.Parse(gcsPath)
	if err != nil {
	glog.Errorf("Error - %v: Failed to parse gcsPath - %s", gcsPath, err)
	return "", "", err
	}
	if parsedURL.Scheme == schemeGCS {
	// In this case, the gcsPath looks like: gs://cos-tools/16108.403.42/kernel-headers.tgz
	// url.Hostname() returns "cos-tools" which is the GCSBucket
	// url.EscapedPath() returns "/16108.403.42/kernel-headers.tgz" with a leading "/".
	return parsedURL.Hostname(), strings.TrimLeft(parsedURL.EscapedPath(), "/"), nil
	} else if parsedURL.Scheme == schemeHttps && strings.ToLower(parsedURL.Host) == gcsHostName {
	// In this case, the gcsPath looks like: https://storage.googleapis.com/cos-tools/1700.00.0/lakitu/gpu_default_version
	// Split the path into parts using "/"
	pathParts := strings.Split(parsedURL.Path, "/")
	// Check if there are enough parts in the path
	if len(pathParts) < 3 {
	glog.Errorf("Error: Invalid gcsPath input: %s - the number of path parts separated by '/' is smaller than 3.", gcsPath)
	return "", "", fmt.Errorf("error: invalid gcs_path input: %s - the number of path parts separated by '/' is smaller than 3", gcsPath)
	}
	gcsBucket := pathParts[1]
	gcsPath := strings.Join(pathParts[2:], "/")
	return gcsBucket, gcsPath, nil
	} else {
	glog.Errorf("Invalid input gcsPath: %s - unable to parse it to get gcsBucket and gcsPath", gcsPath)
	return "", "", fmt.Errorf("error: invalid input gcs_path: %s - unable to parse it to get gcsBucket and gcsPath", gcsPath)
	}
	}