streamer.go

package bqworker

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"cloud.google.com/go/bigquery"
	"github.com/gjbae1212/go-bqworker/util"
	"github.com/rotisserie/eris"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/option"
)

type (
	Row interface {
		bigquery.ValueSaver

		ProjectId() string      // GCP ProjectId
		Schema() *TableSchema   // BigQuery Table Schema
		PublishedAt() time.Time // PublishedAt is time, such as sends row to Bigquery.
		InsertId() string       // InsertId is unique id in BigQuery table.
	}

	Streamer interface {
		AddRow(ctx context.Context, row Row) error     // Add bigquery row to dispatcher
		AddRowSync(ctx context.Context, row Row) error // Add bigquery row to dispatcher sync.
		Start() error                                  // Start Streamer
		Stop() error                                   // Stop Streamer
	}

	streamer struct {
		sync.Mutex
		clientMap      map[string]*bigquery.Client // map[gcp project id]client
		tableSchemaMap map[string][]*TableSchema   // map[gcp project id]client
		dispatcher     *workerDispatcher
		loopDone       chan struct{}
		cfg            *config
	}
)

var (
	defaultOptions = []Option{
		WithQueueSize(defaultQueueSize),
		WithWorkerSize(defaultWorkerSize),
		WithWorkerStack(defaultWorkerStack),
		WithWorkerWaitDuration(defaultWorkerWaitDuration),
		WithDispatcherLoopWaitDuration(defaultDispatcherLoopWaitDuration),
		WithMaxRetry(defaultMaxRetry),
		WithErrorHandler(defaultErrorHandler),
	}
)

// AddRow adds row to dispatcher.
func (st *streamer) AddRow(ctx context.Context, row Row) error {
	if row == nil || row.PublishedAt().IsZero() || row.ProjectId() == "" || row.InsertId() == "" {
		return eris.Wrap(ErrInvalidParams, "")
	}
	schema := row.Schema()
	if schema == nil {
		return eris.Wrap(ErrNotFoundBigQueryTableSchema, "")
	}

	if _, ok := st.clientMap[row.ProjectId()]; !ok {
		return eris.Wrap(ErrNotFoundBigQueryClient, "")
	}

	if err := st.dispatcher.addQueue(ctx, &message{
		projectId: row.ProjectId(),
		datasetId: schema.datasetId,
		tableId:   st.getTableId(schema, row.PublishedAt()),
		row:       row,
	}); err != nil {
		return eris.Wrap(err, "")
	}
	return nil
}

// AddRowSync adds row to dispatcher based on sync.
// It's to wait until it is completed.
func (st *streamer) AddRowSync(ctx context.Context, row Row) error {
	if row == nil || row.PublishedAt().IsZero() || row.ProjectId() == "" || row.InsertId() == "" {
		return eris.Wrap(ErrInvalidParams, "")
	}
	schema := row.Schema()
	if schema == nil {
		return eris.Wrap(ErrNotFoundBigQueryTableSchema, "")
	}

	client, ok := st.clientMap[row.ProjectId()]
	if !ok {
		return eris.Wrap(ErrNotFoundBigQueryClient, "")
	}

	tableId := st.getTableId(schema, row.PublishedAt())
	if err := putRowsToBigQuery(ctx, client, schema.datasetId, tableId, []*bqRow{&bqRow{row: row}}); err != nil {
		return eris.Wrap(err, "")
	}
	return nil
}

// Start starts streamer.
func (st *streamer) Start() error {
	st.Lock()
	defer st.Unlock()
	st.dispatcher.start()
	st.runLoop()
	return nil
}

// Stop stops streamer.
func (st *streamer) Stop() error {
	st.Lock()
	defer st.Unlock()
	st.loopDone <- struct{}{}
	return nil
}

func (st *streamer) runLoop() error {
	go func() {
		// distributes time tick.
		seed := time.Duration(rand.Int63()%600) * time.Second
		ticker := time.NewTicker(st.cfg.dispatcherLoopWaitDuration + seed)
		if err := st.createOrUpdateTable(); err != nil {
			st.cfg.errFunc(eris.Wrap(err, ""))
		}
		for {
			select {
			case <-st.loopDone:
				return
			case <-ticker.C:
				if err := st.createOrUpdateTable(); err != nil {
					st.cfg.errFunc(eris.Wrap(err, ""))
				}
			}
		}
	}()
	return nil
}

// createOrUpdateTable creates or updates bigquery table.
func (st *streamer) createOrUpdateTable() error {
	for projectId, schemaList := range st.tableSchemaMap {
		for _, schema := range schemaList {
			// get bigquery client
			client, ok := st.clientMap[projectId]
			if !ok {
				st.cfg.errFunc(eris.Wrap(ErrNotFoundBigQueryClient, ""))
				continue
			}

			// create or update table.
			// today and tomorrow table will be created, if create table action is executed.
			for _, t := range []time.Time{time.Now(), time.Now().Add(24 * time.Hour)} {
				func() {
					ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
					defer cancel()

					tableId := st.getTableId(schema, t)
					table := client.Dataset(schema.datasetId).Table(tableId)
					md, err := table.Metadata(ctx)
					if err != nil || md == nil { // table not exists
						if suberr := table.Create(ctx, schema.meta); suberr != nil {
							st.cfg.errFunc(eris.Wrap(suberr, ""))
							return
						}
						fmt.Printf("[%s][bq-table][%s] create %s\n", util.GetHostname(), util.TimeToString(time.Now()), tableId)
						return
					} else {
						// Add new columns to table, if a table exists,
						fields := make(map[string]*bigquery.FieldSchema)
						for _, field := range schema.meta.Schema {
							fields[field.Name] = field
						}
						for _, field := range md.Schema {
							delete(fields, field.Name)
						}
						if len(fields) > 0 {
							newSchemaList := md.Schema
							for _, field := range fields {
								newSchemaList = append(newSchemaList, field)
							}
							update := bigquery.TableMetadataToUpdate{Schema: newSchemaList}
							if _, suberr := table.Update(ctx, update, ""); suberr != nil {
								st.cfg.errFunc(eris.Wrap(suberr, ""))
								return
							}
							fmt.Printf("[%s][bq-table][%s] update %s (%v)\n", util.GetHostname(), util.TimeToString(time.Now()), tableId, fields)
						}
						return
					}
				}()
			}
		}
	}
	return nil
}

// getTableId get bigquery table id.
func (st *streamer) getTableId(schema *TableSchema, t time.Time) string {
	switch schema.period {
	case NotExist:
		return schema.prefix
	case Daily:
		return schema.prefix + util.TimeToDailyStringFormat(t)
	case Monthly:
		return schema.prefix + util.TimeToMonthlyStringFormat(t)
	case Yearly:
		return schema.prefix + util.TimeToYearlyStringFormat(t)
	}
	return ""
}

// NewStreamer returns bigquery streamer which inserts data with bulk parallel.
func NewStreamer(tableSchemaGroup []*TableSchemaGroup, opts ...Option) (Streamer, error) {
	if len(tableSchemaGroup) == 0 {
		return nil, eris.Wrap(ErrInvalidParams, "")
	}

	cfg := &config{credentialMap: map[string]string{}}
	for _, opt := range defaultOptions {
		opt.apply(cfg)
	}
	for _, opt := range opts {
		opt.apply(cfg)
	}

	st := &streamer{cfg: cfg, clientMap: map[string]*bigquery.Client{},
		tableSchemaMap: map[string][]*TableSchema{},
		loopDone:       make(chan struct{}),
	}

	for _, group := range tableSchemaGroup {
		projectId := group.projectId
		credential := group.credential
		if projectId == "" || credential == "" {
			return nil, eris.Wrap(ErrInvalidParams, "")
		}
		tableSchemaMap := make([]*TableSchema, 0, len(group.tableSchemaList))
		for _, schema := range group.tableSchemaList {
			if !schema.valid() {
				return nil, eris.Wrap(ErrInvalidParams, "")
			}
			tableSchemaMap = append(tableSchemaMap, schema)
		}

		// get google jwt obj from jwt bytes.
		jwtConfig, err := google.JWTConfigFromJSON([]byte(credential), bigquery.Scope)
		if err != nil {
			return nil, eris.Wrap(err, "")
		}

		// create client for BigQuery.
		client, err := bigquery.NewClient(context.Background(), projectId,
			option.WithTokenSource(jwtConfig.TokenSource(context.Background())))
		if err != nil {
			return nil, eris.Wrap(err, "")
		}
		st.cfg.credentialMap[projectId] = credential
		st.clientMap[projectId] = client
		st.tableSchemaMap[projectId] = tableSchemaMap
	}

	// create worker dispatcher.
	dispatcher, err := newWorkerDispatcher(st.cfg)
	if err != nil {
		return nil, eris.Wrap(err, "")
	}
	st.dispatcher = dispatcher
	return st, nil
}

// putRowsToBigQuery puts rows to bigquery.
// it must not wrap eris error.
func putRowsToBigQuery(ctx context.Context, client *bigquery.Client, datasetId, tableId string, rows []*bqRow) error {
	if ctx == nil || client == nil || datasetId == "" || tableId == "" {
		return ErrInvalidParams
	}
	if len(rows) == 0 {
		return nil
	}
	inserter := client.Dataset(datasetId).Table(tableId).Inserter()
	inserter.SkipInvalidRows = true
	inserter.IgnoreUnknownValues = true
	return inserter.Put(ctx, rows)
}