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

     1  //go:generate compogen readme --connector ./config ./README.mdx
     2  package bigquery
     3  
     4  import (
     5  	"context"
     6  	_ "embed"
     7  	"errors"
     8  	"fmt"
     9  	"sync"
    10  
    11  	"cloud.google.com/go/bigquery"
    12  	"go.uber.org/zap"
    13  	"google.golang.org/api/option"
    14  	"google.golang.org/protobuf/types/known/structpb"
    15  
    16  	"github.com/instill-ai/component/pkg/base"
    17  )
    18  
    19  const (
    20  	taskInsert = "TASK_INSERT"
    21  )
    22  
    23  //go:embed config/definition.json
    24  var definitionJSON []byte
    25  
    26  //go:embed config/tasks.json
    27  var tasksJSON []byte
    28  
    29  var once sync.Once
    30  var con *connector
    31  
    32  type connector struct {
    33  	base.BaseConnector
    34  }
    35  
    36  type execution struct {
    37  	base.BaseConnectorExecution
    38  }
    39  
    40  func Init(l *zap.Logger, u base.UsageHandler) *connector {
    41  	once.Do(func() {
    42  		con = &connector{
    43  			BaseConnector: base.BaseConnector{
    44  				Logger:       l,
    45  				UsageHandler: u,
    46  			},
    47  		}
    48  		err := con.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
    49  		if err != nil {
    50  			panic(err)
    51  		}
    52  	})
    53  	return con
    54  }
    55  
    56  func (c *connector) CreateExecution(sysVars map[string]any, connection *structpb.Struct, task string) (*base.ExecutionWrapper, error) {
    57  	return &base.ExecutionWrapper{Execution: &execution{
    58  		BaseConnectorExecution: base.BaseConnectorExecution{Connector: c, SystemVariables: sysVars, Connection: connection, Task: task},
    59  	}}, nil
    60  }
    61  
    62  func NewClient(jsonKey, projectID string) (*bigquery.Client, error) {
    63  	return bigquery.NewClient(context.Background(), projectID, option.WithCredentialsJSON([]byte(jsonKey)))
    64  }
    65  
    66  func getJSONKey(config *structpb.Struct) string {
    67  	return config.GetFields()["json_key"].GetStringValue()
    68  }
    69  func getProjectID(config *structpb.Struct) string {
    70  	return config.GetFields()["project_id"].GetStringValue()
    71  }
    72  func getDatasetID(config *structpb.Struct) string {
    73  	return config.GetFields()["dataset_id"].GetStringValue()
    74  }
    75  func getTableName(config *structpb.Struct) string {
    76  	return config.GetFields()["table_name"].GetStringValue()
    77  }
    78  
    79  func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
    80  	outputs := []*structpb.Struct{}
    81  
    82  	client, err := NewClient(getJSONKey(e.Connection), getProjectID(e.Connection))
    83  	if err != nil || client == nil {
    84  		return nil, fmt.Errorf("error creating BigQuery client: %v", err)
    85  	}
    86  	defer client.Close()
    87  
    88  	for _, input := range inputs {
    89  		var output *structpb.Struct
    90  		switch e.Task {
    91  		case taskInsert, "":
    92  			datasetID := getDatasetID(e.Connection)
    93  			tableName := getTableName(e.Connection)
    94  			tableRef := client.Dataset(datasetID).Table(tableName)
    95  			metaData, err := tableRef.Metadata(context.Background())
    96  			if err != nil {
    97  				return nil, err
    98  			}
    99  			valueSaver, err := getDataSaver(input, metaData.Schema)
   100  			if err != nil {
   101  				return nil, err
   102  			}
   103  			err = insertDataToBigQuery(getProjectID(e.Connection), datasetID, tableName, valueSaver, client)
   104  			if err != nil {
   105  				return nil, err
   106  			}
   107  			output = &structpb.Struct{Fields: map[string]*structpb.Value{"status": {Kind: &structpb.Value_StringValue{StringValue: "success"}}}}
   108  		default:
   109  			return nil, fmt.Errorf("unsupported task: %s", e.Task)
   110  		}
   111  		outputs = append(outputs, output)
   112  	}
   113  	return outputs, nil
   114  }
   115  
   116  func (c *connector) Test(sysVars map[string]any, connection *structpb.Struct) error {
   117  
   118  	client, err := NewClient(getJSONKey(connection), getProjectID(connection))
   119  	if err != nil || client == nil {
   120  		return fmt.Errorf("error creating BigQuery client: %v", err)
   121  	}
   122  	defer client.Close()
   123  	if client.Project() == getProjectID(connection) {
   124  		return nil
   125  	}
   126  	return errors.New("project ID does not match")
   127  }