github.com/instill-ai/component@v0.16.0-beta/pkg/connector/googlecloudstorage/v0/main.go (about)

     1  //go:generate compogen readme --connector ./config ./README.mdx
     2  package googlecloudstorage
     3  
     4  import (
     5  	"context"
     6  	_ "embed"
     7  	"fmt"
     8  	"sync"
     9  
    10  	"cloud.google.com/go/storage"
    11  	"go.uber.org/zap"
    12  	"google.golang.org/api/option"
    13  	"google.golang.org/protobuf/types/known/structpb"
    14  
    15  	"github.com/instill-ai/component/pkg/base"
    16  )
    17  
    18  const (
    19  	taskUpload = "TASK_UPLOAD"
    20  )
    21  
    22  //go:embed config/definition.json
    23  var definitionJSON []byte
    24  
    25  //go:embed config/tasks.json
    26  var tasksJSON []byte
    27  var once sync.Once
    28  var con *connector
    29  
    30  type connector struct {
    31  	base.BaseConnector
    32  }
    33  
    34  type execution struct {
    35  	base.BaseConnectorExecution
    36  }
    37  
    38  func Init(l *zap.Logger, u base.UsageHandler) *connector {
    39  	once.Do(func() {
    40  		con = &connector{
    41  			BaseConnector: base.BaseConnector{
    42  				Logger:       l,
    43  				UsageHandler: u,
    44  			},
    45  		}
    46  		err := con.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
    47  		if err != nil {
    48  			panic(err)
    49  		}
    50  	})
    51  	return con
    52  }
    53  
    54  func (c *connector) CreateExecution(sysVars map[string]any, connection *structpb.Struct, task string) (*base.ExecutionWrapper, error) {
    55  	return &base.ExecutionWrapper{Execution: &execution{
    56  		BaseConnectorExecution: base.BaseConnectorExecution{Connector: c, SystemVariables: sysVars, Connection: connection, Task: task},
    57  	}}, nil
    58  }
    59  
    60  func NewClient(jsonKey string) (*storage.Client, error) {
    61  	return storage.NewClient(context.Background(), option.WithCredentialsJSON([]byte(jsonKey)))
    62  }
    63  
    64  func getBucketName(config *structpb.Struct) string {
    65  	return config.GetFields()["bucket_name"].GetStringValue()
    66  }
    67  
    68  func getJSONKey(config *structpb.Struct) string {
    69  	return config.GetFields()["json_key"].GetStringValue()
    70  }
    71  
    72  func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
    73  	outputs := []*structpb.Struct{}
    74  
    75  	client, err := NewClient(getJSONKey(e.Connection))
    76  	if err != nil || client == nil {
    77  		return nil, fmt.Errorf("error creating GCS client: %v", err)
    78  	}
    79  	defer client.Close()
    80  	for _, input := range inputs {
    81  		var output *structpb.Struct
    82  		switch e.Task {
    83  		case taskUpload, "":
    84  			objectName := input.GetFields()["object_name"].GetStringValue()
    85  			data := input.GetFields()["data"].GetStringValue()
    86  			bucketName := getBucketName(e.Connection)
    87  			err = uploadToGCS(client, bucketName, objectName, data)
    88  			if err != nil {
    89  				return nil, err
    90  			}
    91  
    92  			gsutilURI := fmt.Sprintf("gs://%s/%s", bucketName, objectName)
    93  			authenticatedURL := fmt.Sprintf("https://storage.cloud.google.com/%s/%s?authuser=1", bucketName, objectName)
    94  			publicURL := ""
    95  
    96  			// Check whether the object is public or not
    97  			publicAccess, err := isObjectPublic(client, bucketName, objectName)
    98  			if err != nil {
    99  				return nil, err
   100  			}
   101  			if publicAccess {
   102  				publicURL = fmt.Sprintf("https://storage.googleapis.com/%s/%s", bucketName, objectName)
   103  			}
   104  
   105  			output = &structpb.Struct{Fields: map[string]*structpb.Value{
   106  				"authenticated_url": {Kind: &structpb.Value_StringValue{StringValue: authenticatedURL}},
   107  				"gsutil_uri":        {Kind: &structpb.Value_StringValue{StringValue: gsutilURI}},
   108  				"public_url":        {Kind: &structpb.Value_StringValue{StringValue: publicURL}},
   109  				"public_access":     {Kind: &structpb.Value_BoolValue{BoolValue: publicAccess}},
   110  				"status":            {Kind: &structpb.Value_StringValue{StringValue: "success"}}}}
   111  		}
   112  		outputs = append(outputs, output)
   113  	}
   114  	return outputs, nil
   115  }
   116  
   117  func (c *connector) Test(sysVars map[string]any, connection *structpb.Struct) error {
   118  
   119  	client, err := NewClient(getJSONKey(connection))
   120  	if err != nil {
   121  		return fmt.Errorf("error creating GCS client: %v", err)
   122  	}
   123  	if client == nil {
   124  		return fmt.Errorf("GCS client is nil")
   125  	}
   126  	defer client.Close()
   127  	return nil
   128  }