<commit_before>package main

import (
	"fmt"
	"sort"
	"time"

	bigquery "google.golang.org/api/bigquery/v2"
	"google.golang.org/api/googleapi"
)

// A BigQueryStreamer is a BigQuery streamer, queuing rows and streaming to
// BigQuery in bulk by calling InsertAll().
type BigQueryStreamer struct {
	// BigQuery client connection.
	service *bigquery.Service

	// Upon invokeing Start() streamer will fetch rows from this channel,
	// and queue it into an internal rows queue.
	rowChannel chan *row

	// Internal list to queue rows for stream insert.
	rows []*row

	// Rows index to queue next row into.
	// TODO using a row index is probably not the best way.
	// Maybe we should instead create a slice with len = 0, cap = 500 and use
	// len() instead.
	rowIndex int // 0

	// Max delay between flushes to BigQuery.
	maxDelay time.Duration

	// Sleep delay after a rejected insert and before retry.
	sleepBeforeRetry time.Duration

	// Maximum retry insert attempts for non-rejected row insert errors.
	// e.g. GoogleAPI HTTP errors, generic HTTP errors, etc.
	maxRetryInsert int

	// Shutdown channel to stop Start() execution.
	stopChannel chan bool

	// Errors are reported to this channel.
	Errors chan error

	// The following functions can be overriden for unit testing.

	// Start read-queue-stream loop function.
	Start func()

	// Stop read-queue-stream loop function.
	Stop func()

	// Flush to BigQuery function.
	flush func()

	// Insert all tables to bigquery function.
	insertAll func()

	// Insert table to BigQuery function.
	insertTable func(projectId, datasetId, tableId string, t table) (r *bigquery.TableDataInsertAllResponse, err error)
}

// NewBigQueryStreamer returns a new BigQueryStreamer.
func NewBigQueryStreamer(
	service *bigquery.Service,
	maxRows int,
	maxDelay time.Duration,
	sleepBeforeRetry time.Duration,
	maxRetryInsert int) (b *BigQueryStreamer, err error) {
	// TODO add testing for nil bigquery.Service (this will intervene with tests though,
	// maybe find a way to mock this type somehow?)

	// TODO maybe return error if maxRows > 500?

	if maxRows <= 0 {
		err = fmt.Errorf("maxRows must be positive int > 0")
		return
	}

	if maxDelay <= 0*time.Nanosecond {
		err = fmt.Errorf("maxDelay must be positive time.Duration >= 1 * time.Nanosecond")
		return
	}

	if sleepBeforeRetry <= 0*time.Nanosecond {
		err = fmt.Errorf("sleepBeforeRetry must be positive time.Duration >= 1 * time.Nanosecond")
		return
	}

	if maxRetryInsert < 0 {
		err = fmt.Errorf("maxRetryInsert must be non-negative int >= 0")
		return
	}

	b = &BigQueryStreamer{
		service:          service,
		rowChannel:       make(chan *row, maxRows),
		rows:             make([]*row, maxRows),
		rowIndex:         0,
		maxDelay:         maxDelay,
		sleepBeforeRetry: sleepBeforeRetry,
		maxRetryInsert:   maxRetryInsert,
		stopChannel:      make(chan bool),
		Errors:           make(chan error, errorBufferSize),
	}

	// Assign function defaults.
	b.Start = b.start
	b.Stop = b.stop
	b.flush = b.flushToBigQuery
	b.insertAll = b.insertAllToBigQuery
	b.insertTable = b.insertTableToBigQuery

	return
}

// startStreamer infinitely reads rows from rowChannel and queues them internaly.
// It flushes to BigQuery when queue is filled (according to maxRows) or timer has expired (according to maxDelay).
//
// This function is assigned to BigQueryStreamer.Start member.
// It is overridable so we can test BigQueryMultiStreamer without actually
// starting the streamers.
//
// Note the read-insert-flush loop will never stop, so this function should be
// executed in a goroutine, and stopped via calling Stop().
func (b *BigQueryStreamer) start() {
	t := time.NewTimer(b.maxDelay)
	toStop := false
	for {
		// Flush and reset timer when one of the following signals (channels) fire:
		select {
		case toStop = <-b.stopChannel:
		case <-t.C:
		case r := <-b.rowChannel:
			// Insert row to queue.
			b.rows[b.rowIndex] = r
			b.rowIndex++

			// Don't flush if rows queue isn't full.
			if b.rowIndex < len(b.rows) {
				continue
			}
		}

		b.flush()

		if !toStop {
			t.Reset(b.maxDelay)
		} else {
			t.Stop()
			return
		}
	}
}

// stopStreamer sends a stop message to stop channel, causing Start() infinite loop to stop.
func (b *BigQueryStreamer) stop() {
	b.stopChannel <- true
}

// flushAll streams all queued rows to BigQuery and resets rows queue by
// creating a new queue.
//
// This function is assigned to BigQueryStreamer.flush member.
// It is overridable so we can test BigQueryStreamer without actually flushing
// to BigQuery.
//
// TODO Consider making this public. If so, we should use a mutex to lock the object,
// otherwise if the object is running in another goroutine it can call this in parallel.
func (b *BigQueryStreamer) flushToBigQuery() {
	b.insertAll()

	// Init (reset) a new rows queue - clear old one and re-allocate.
	b.rows = nil
	b.rows = make([]*row, len(b.rows))
	b.rowIndex = 0
}

// QueueRow sends a single row to the row channel, which will be queued and inserted in bulk with other queued rows.
func (b *BigQueryStreamer) QueueRow(projectID, datasetID, tableID string, jsonRow map[string]bigquery.JsonValue) {
	b.rowChannel <- &row{projectID, datasetID, tableID, jsonRow}
}

// BigQuery data types.
type project map[string]dataset<commit_msg>fixed index out of range bug with last PR<commit_after>package main

import (
	"fmt"
	"sort"
	"time"

	bigquery "google.golang.org/api/bigquery/v2"
	"google.golang.org/api/googleapi"
)

// A BigQueryStreamer is a BigQuery streamer, queuing rows and streaming to
// BigQuery in bulk by calling InsertAll().
type BigQueryStreamer struct {
	// BigQuery client connection.
	service *bigquery.Service

	// Upon invokeing Start() streamer will fetch rows from this channel,
	// and queue it into an internal rows queue.
	rowChannel chan *row

	// Internal list to queue rows for stream insert.
	rows []*row

	// Rows index to queue next row into.
	// TODO using a row index is probably not the best way.
	// Maybe we should instead create a slice with len = 0, cap = 500 and use
	// len() instead.
	rowIndex int // 0

	// Max delay between flushes to BigQuery.
	maxDelay time.Duration

	// Sleep delay after a rejected insert and before retry.
	sleepBeforeRetry time.Duration

	// Maximum retry insert attempts for non-rejected row insert errors.
	// e.g. GoogleAPI HTTP errors, generic HTTP errors, etc.
	maxRetryInsert int

	// Shutdown channel to stop Start() execution.
	stopChannel chan bool

	// Errors are reported to this channel.
	Errors chan error

	// The following functions can be overriden for unit testing.

	// Start read-queue-stream loop function.
	Start func()

	// Stop read-queue-stream loop function.
	Stop func()

	// Flush to BigQuery function.
	flush func()

	// Insert all tables to bigquery function.
	insertAll func()

	// Insert table to BigQuery function.
	insertTable func(projectId, datasetId, tableId string, t table) (r *bigquery.TableDataInsertAllResponse, err error)
}

// NewBigQueryStreamer returns a new BigQueryStreamer.
func NewBigQueryStreamer(
	service *bigquery.Service,
	maxRows int,
	maxDelay time.Duration,
	sleepBeforeRetry time.Duration,
	maxRetryInsert int) (b *BigQueryStreamer, err error) {
	// TODO add testing for nil bigquery.Service (this will intervene with tests though,
	// maybe find a way to mock this type somehow?)

	// TODO maybe return error if maxRows > 500?

	if maxRows <= 0 {
		err = fmt.Errorf("maxRows must be positive int > 0")
		return
	}

	if maxDelay <= 0*time.Nanosecond {
		err = fmt.Errorf("maxDelay must be positive time.Duration >= 1 * time.Nanosecond")
		return
	}

	if sleepBeforeRetry <= 0*time.Nanosecond {
		err = fmt.Errorf("sleepBeforeRetry must be positive time.Duration >= 1 * time.Nanosecond")
		return
	}

	if maxRetryInsert < 0 {
		err = fmt.Errorf("maxRetryInsert must be non-negative int >= 0")
		return
	}

	b = &BigQueryStreamer{
		service:          service,
		rowChannel:       make(chan *row, maxRows),
		rows:             make([]*row, maxRows),
		rowIndex:         0,
		maxDelay:         maxDelay,
		sleepBeforeRetry: sleepBeforeRetry,
		maxRetryInsert:   maxRetryInsert,
		stopChannel:      make(chan bool),
		Errors:           make(chan error, errorBufferSize),
	}

	// Assign function defaults.
	b.Start = b.start
	b.Stop = b.stop
	b.flush = b.flushToBigQuery
	b.insertAll = b.insertAllToBigQuery
	b.insertTable = b.insertTableToBigQuery

	return
}

// startStreamer infinitely reads rows from rowChannel and queues them internaly.
// It flushes to BigQuery when queue is filled (according to maxRows) or timer has expired (according to maxDelay).
//
// This function is assigned to BigQueryStreamer.Start member.
// It is overridable so we can test BigQueryMultiStreamer without actually
// starting the streamers.
//
// Note the read-insert-flush loop will never stop, so this function should be
// executed in a goroutine, and stopped via calling Stop().
func (b *BigQueryStreamer) start() {
	t := time.NewTimer(b.maxDelay)
	toStop := false
	for {
		// Flush and reset timer when one of the following signals (channels) fire:
		select {
		case toStop = <-b.stopChannel:
		case <-t.C:
		case r := <-b.rowChannel:
			// Insert row to queue.
			b.rows[b.rowIndex] = r
			b.rowIndex++

			// Don't flush if rows queue isn't full.
			if b.rowIndex < len(b.rows) {
				continue
			}
		}

		b.flush()

		if !toStop {
			t.Reset(b.maxDelay)
		} else {
			t.Stop()
			return
		}
	}
}

// stopStreamer sends a stop message to stop channel, causing Start() infinite loop to stop.
func (b *BigQueryStreamer) stop() {
	b.stopChannel <- true
}

// flushAll streams all queued rows to BigQuery and resets rows queue by
// creating a new queue.
//
// This function is assigned to BigQueryStreamer.flush member.
// It is overridable so we can test BigQueryStreamer without actually flushing
// to BigQuery.
//
// TODO Consider making this public. If so, we should use a mutex to lock the object,
// otherwise if the object is running in another goroutine it can call this in parallel.
func (b *BigQueryStreamer) flushToBigQuery() {
	b.insertAll()

	// Init (reset) a new rows queue - clear old one and re-allocate.
	b.rows = make([]*row, len(b.rows))
	b.rowIndex = 0
}

// QueueRow sends a single row to the row channel, which will be queued and inserted in bulk with other queued rows.
func (b *BigQueryStreamer) QueueRow(projectID, datasetID, tableID string, jsonRow map[string]bigquery.JsonValue) {
	b.rowChannel <- &row{projectID, datasetID, tableID, jsonRow}
}

// BigQuery data types.
type project map[string]dataset
{'fraction_non_alphanumeric': 0.0940154804606381, 'fraction_numerical': 0.0035869360015102886, 'mean_word_length': 3.3529170090386198, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2929937', 'n_tokens_mistral': 3446, 'n_tokens_neox': 3227, 'n_words': 1243}