feat: google sheets external table support for bigquery datastore (#45)

Co-authored-by: Favian Kharisma Hazman <[email protected]>

Author: vianhazman

ext/datastore/bigquery/adapter.go

@@ -84,6 +84,73 @@ func bqFieldModeTo(field BQField) (fieldMode, error) {
 	}
 	return fm, nil
 }
+func bqGoogleSheetsOptionsTo(m map[string]interface{}) *bqapi.GoogleSheetsOptions {
+	var skipLeadingRows int64
+	var sheetRange string
+
+	if val, ok := m["skip_leading_rows"]; ok {
+		skipLeadingRows = int64(val.(float64))
+	}
+
+	if val, ok := m["range"]; ok {
+		sheetRange = val.(string)
+	}
+
+	return &bqapi.GoogleSheetsOptions{
+		SkipLeadingRows: skipLeadingRows,
+		Range:           sheetRange,
+	}
+}
+
+func bqGoogleSheetsOptionsFrom(opt *bqapi.GoogleSheetsOptions) map[string]interface{} {
+	resultMap := make(map[string]interface{})
+
+	if opt.SkipLeadingRows != 0 {
+		// Map value of int has to be converted to float because of using interface{}
+		resultMap["skip_leading_rows"] = float64(opt.SkipLeadingRows)
+	}
+	if opt.Range != "" {
+		resultMap["range"] = opt.Range
+	}
+	return resultMap
+}
+
+func bqExternalDataConfigTo(es BQExternalSource) (*bqapi.ExternalDataConfig, error) {
+	var option bqapi.ExternalDataConfigOptions
+	var sourceType bqapi.DataFormat
+	switch bqapi.DataFormat(strings.ToUpper(es.SourceType)) {
+	case bqapi.GoogleSheets:
+		option = bqGoogleSheetsOptionsTo(es.Config)
+		sourceType = bqapi.GoogleSheets
+	default:
+		return &bqapi.ExternalDataConfig{}, fmt.Errorf("Source format not yet implemented %s", es.SourceType)
+	}
+
+	externalConfig := &bqapi.ExternalDataConfig{
+		SourceFormat: sourceType,
+		SourceURIs:   es.SourceURIs,
+		Options:      option,
+	}
+	return externalConfig, nil
+}
+
+func bqExternalDataConfigFrom(c *bqapi.ExternalDataConfig) (*BQExternalSource, error) {
+	var option map[string]interface{}
+
+	switch c.SourceFormat {
+	case bqapi.GoogleSheets:
+		option = bqGoogleSheetsOptionsFrom(c.Options.(*bqapi.GoogleSheetsOptions))
+	default:
+		return &BQExternalSource{}, fmt.Errorf("Source format not yet implemented %s", c.SourceFormat)
+	}
+
+	externalDataConfig := &BQExternalSource{
+		SourceType: string(c.SourceFormat),
+		SourceURIs: c.SourceURIs,
+		Config:     option,
+	}
+	return externalDataConfig, nil
+}
 
 func bqFieldModeFrom(fm fieldMode) string {
 	if fm.repeated {
@@ -159,6 +226,14 @@ func bqCreateTableMetaAdapter(t BQTable) (meta *bqapi.TableMetadata, err error)
 			meta.RangePartitioning = bqPartitioningRangeTo(*t.Metadata.Partition)
 		}
 	}
+
+	if t.Metadata.Source != nil {
+		meta.ExternalDataConfig, err = bqExternalDataConfigTo(*t.Metadata.Source)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if t.Metadata.ExpirationTime != "" {
 		expiryTime, err := time.Parse(time.RFC3339, t.Metadata.ExpirationTime)
 		if err != nil {
@@ -195,6 +270,7 @@ func bqUpdateTableMetaAdapter(t BQTable) (meta bqapi.TableMetadataToUpdate, err
 	for key, value := range t.Metadata.Labels {
 		meta.SetLabel(key, value)
 	}
+
 	if t.Metadata.ExpirationTime != "" {
 		expiryTime, err := time.Parse(time.RFC3339, t.Metadata.ExpirationTime)
 		if err != nil {

ext/datastore/bigquery/adapter_test.go

@@ -67,6 +67,29 @@ func TestAdapter(t *testing.T) {
 		},
 	}
 
+	t.Run("should convert from and to BQ ExternalDataConfig successfully", func(t *testing.T) {
+		externalDataSource := BQExternalSource{
+			SourceType: string(ExternalTableTypeGoogleSheets),
+			SourceURIs: []string{"http://googlesheets.com/1234"},
+			Config:     map[string]interface{}{"skip_leading_rows": 1.0, "range": "A!:A1:B1"},
+		}
+		expectedBQExternalDataConfig := &bigquery.ExternalDataConfig{
+			SourceFormat: bigquery.GoogleSheets,
+			SourceURIs:   []string{"http://googlesheets.com/1234"},
+			Options: &bigquery.GoogleSheetsOptions{
+				SkipLeadingRows: 1,
+				Range:           "A!:A1:B1",
+			},
+		}
+		bQExternalDataConfigResult, err := bqExternalDataConfigTo(externalDataSource)
+		assert.Nil(t, err)
+		assert.Equal(t, expectedBQExternalDataConfig, bQExternalDataConfigResult)
+
+		externalSourceResult, err := bqExternalDataConfigFrom(bQExternalDataConfigResult)
+		assert.Nil(t, err)
+		assert.Equal(t, &externalDataSource, externalSourceResult)
+	})
+
 	t.Run("should convert from and to BQ TimePartitioning successfully", func(t *testing.T) {
 		partitionField := "partition-field"
 		partitionExpiryInHours := int64(720)
@@ -150,6 +173,7 @@ func TestAdapter(t *testing.T) {
 		assert.Equal(t, schema, schemaResult)
 		assert.Nil(t, err)
 	})
+
 	t.Run("bqCreateTableMetaAdapter", func(t *testing.T) {
 		t.Run("should convert to BQ TableMetadata", func(t *testing.T) {
 			testingProject := "project"

ext/datastore/bigquery/bigquery.go

@@ -41,9 +41,10 @@ func (b BigQuery) Description() string {
 
 func (b BigQuery) Types() map[models.ResourceType]models.DatastoreTypeController {
 	return map[models.ResourceType]models.DatastoreTypeController{
-		models.ResourceTypeTable:   &tableSpec{},
-		models.ResourceTypeView:    &standardViewSpec{},
-		models.ResourceTypeDataset: &datasetSpec{},
+		models.ResourceTypeTable:         &tableSpec{},
+		models.ResourceTypeView:          &standardViewSpec{},
+		models.ResourceTypeDataset:       &datasetSpec{},
+		models.ResourceTypeExternalTable: &externalTableSpec{},
 	}
 }
 
@@ -65,6 +66,8 @@ func (b *BigQuery) CreateResource(ctx context.Context, request models.CreateReso
 		return createStandardView(ctx, request.Resource, client, false)
 	case models.ResourceTypeDataset:
 		return createDataset(ctx, request.Resource, client, false)
+	case models.ResourceTypeExternalTable:
+		return createExternalTable(ctx, request.Resource, client, false)
 	}
 	return fmt.Errorf("unsupported resource type %s", request.Resource.Type)
 }
@@ -87,6 +90,8 @@ func (b *BigQuery) UpdateResource(ctx context.Context, request models.UpdateReso
 		return createStandardView(ctx, request.Resource, client, true)
 	case models.ResourceTypeDataset:
 		return createDataset(ctx, request.Resource, client, true)
+	case models.ResourceTypeExternalTable:
+		return createExternalTable(ctx, request.Resource, client, true)
 	}
 	return fmt.Errorf("unsupported resource type %s", request.Resource.Type)
 }

ext/datastore/bigquery/external_table.go

@@ -0,0 +1,81 @@
+package bigquery
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	bqapi "cloud.google.com/go/bigquery"
+
+	"google.golang.org/api/googleapi"
+
+	"github.com/googleapis/google-cloud-go-testing/bigquery/bqiface"
+	"github.com/odpf/optimus/models"
+	"github.com/pkg/errors"
+)
+
+func createExternalTable(ctx context.Context, spec models.ResourceSpec, client bqiface.Client, upsert bool) error {
+	bqResource, ok := spec.Spec.(BQTable)
+
+	if !ok {
+		return errors.New("failed to read table spec for bigquery")
+	}
+
+	// inherit from base
+	bqResource.Metadata.Labels = spec.Labels
+
+	dataset := client.DatasetInProject(bqResource.Project, bqResource.Dataset)
+	if err := ensureDataset(ctx, dataset, BQDataset{
+		Project:  bqResource.Project,
+		Dataset:  bqResource.Dataset,
+		Metadata: BQDatasetMetadata{},
+	}, false); err != nil {
+		return err
+	}
+	table := dataset.Table(bqResource.Table)
+	return ensureExternalTable(ctx, table, bqResource, upsert)
+}
+
+func ensureExternalTable(ctx context.Context, tableHandle bqiface.Table, t BQTable, upsert bool) error {
+	meta, err := tableHandle.Metadata(ctx)
+	if err != nil {
+		if metaErr, ok := err.(*googleapi.Error); !ok || metaErr.Code != http.StatusNotFound {
+			return err
+		}
+		meta, err := bqCreateTableMetaAdapter(t)
+		if err != nil {
+			return err
+		}
+
+		if t.Metadata.ExpirationTime != "" {
+			expiryTime, err := time.Parse(time.RFC3339, t.Metadata.ExpirationTime)
+			if err != nil {
+				return errors.Wrapf(err, "unable to parse timestamp %s", t.Metadata.ExpirationTime)
+			}
+			meta.ExpirationTime = expiryTime
+		}
+		return tableHandle.Create(ctx, meta)
+	}
+	if !upsert {
+		return nil
+	}
+
+	// update if already exists
+	m := bqapi.TableMetadataToUpdate{
+		Description: t.Metadata.Description,
+	}
+	if t.Metadata.ExpirationTime != "" {
+		expiryTime, err := time.Parse(time.RFC3339, t.Metadata.ExpirationTime)
+		if err != nil {
+			return errors.Wrapf(err, "unable to parse timestamp %s", t.Metadata.ExpirationTime)
+		}
+		m.ExpirationTime = expiryTime
+	}
+	for k, v := range t.Metadata.Labels {
+		m.SetLabel(k, v)
+	}
+	if _, err := tableHandle.Update(ctx, m, meta.ETag); err != nil {
+		return err
+	}
+	return nil
+}

ext/datastore/bigquery/external_table_spec.go

@@ -0,0 +1,72 @@
+package bigquery
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/odpf/optimus/models"
+	"google.golang.org/protobuf/types/known/structpb"
+)
+
+const (
+	ExternalTableTypeGoogleSheets ExternalTableType = "GOOGLE_SHEETS"
+)
+
+type ExternalTableType string
+
+// BQExternalSource specifies table source information for external data source
+type BQExternalSource struct {
+	SourceType string `yaml:"type,omitempty" structs:"type"`
+
+	// External Table URI string for the referenced spreadsheets
+	SourceURIs []string `yaml:"uris,omitempty" structs:"uris,omitempty"`
+
+	// Additional configs for CSV, GoogleSheets, Bigtable, and Parquet formats.
+	Config map[string]interface{} `yaml:"config,omitempty" structs:"config"`
+}
+
+type externalTableSpec struct{}
+
+func (s externalTableSpec) Adapter() models.DatastoreSpecAdapter {
+	return &tableSpecHandler{}
+}
+
+func (s externalTableSpec) Validator() models.DatastoreSpecValidator {
+	return func(spec models.ResourceSpec) error {
+		if !tableNameParseRegex.MatchString(spec.Name) {
+			return fmt.Errorf("for example 'project_name.dataset_name.table_name'")
+		}
+		parsedNames := tableNameParseRegex.FindStringSubmatch(spec.Name)
+		if len(parsedNames) < 3 || len(parsedNames[1]) == 0 || len(parsedNames[2]) == 0 || len(parsedNames[3]) == 0 {
+			return fmt.Errorf("for example 'project_name.dataset_name.table_name'")
+		}
+		return nil
+	}
+}
+
+func (s externalTableSpec) DefaultAssets() map[string]string {
+	return map[string]string{}
+}
+
+func extractTableSourceFromProtoStruct(protoVal *structpb.Value) *BQExternalSource {
+	sInfo := &BQExternalSource{}
+	if protoVal.GetStructValue() == nil {
+		return sInfo
+	}
+	if f, ok := protoVal.GetStructValue().Fields["type"]; ok {
+		sInfo.SourceType = strings.ToUpper(f.GetStringValue())
+	}
+	if f, ok := protoVal.GetStructValue().Fields["uris"]; ok {
+		var sourceURIs []string
+		if f.GetListValue() != nil {
+			for _, value := range f.GetListValue().GetValues() {
+				sourceURIs = append(sourceURIs, value.GetStringValue())
+			}
+		}
+		sInfo.SourceURIs = sourceURIs
+	}
+	if f, ok := protoVal.GetStructValue().Fields["config"]; ok {
+		sInfo.Config = f.GetStructValue().AsMap()
+	}
+	return sInfo
+}