github.com/rohankumardubey/aresdb@v0.0.2-0.20190517170215-e54e3ca06b9c/client/connector.go

//  Copyright (c) 2017-2018 Uber Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"bytes"
	"fmt"
	"net/http"
	"sync"
	"time"

	"strconv"
	"strings"
	"unsafe"

	"github.com/uber-go/tally"
	memCom "github.com/uber/aresdb/memstore/common"
	metaCom "github.com/uber/aresdb/metastore/common"
	"github.com/uber/aresdb/utils"
	"go.uber.org/zap"
)

const (
	// default request time out in seconds
	defaultRequestTimeout = 5
	// default schema refresh interval in seconds
	defaultSchemaRefreshInterval = 600
	dataIngestionHeader          = "application/upsert-data"
	applicationJSONHeader        = "application/json"
)

// Row represents a row of insert data.
type Row []interface{}

// Connector is the connector interface for ares.
type Connector interface {
	// Insert inserts rows to ares
	// returns number of rows inserted and error.
	// updateModes are optional, if ignored for all columns, no need to set
	// if set, then all columns needs to be set
	Insert(tableName string, columnNames []string, rows []Row, updateModes ...memCom.ColumnUpdateMode) (int, error)
}

// UpsertBatchBuilder is an interface of upsertBatch on client side
type UpsertBatchBuilder interface {
	PrepareUpsertBatch(tableName string, columnNames []string, updateModes []memCom.ColumnUpdateMode, rows []Row) ([]byte, int, error)
}

// enumCasesWrapper is a response/request body which wraps enum cases
type enumCasesWrapper struct {
	EnumCases []string
}

type TableSchema struct {
	Table *metaCom.Table
	// maps from column name to columnID for convenience
	ColumnDict map[string]int
}

// enumDict maps from enum value to enumID
type enumDict map[string]int

// UpsertBatchBuilderImpl implements interface UpsertBatchBuilder
type UpsertBatchBuilderImpl struct {
	sync.RWMutex

	logger        *zap.SugaredLogger
	metricScope   tally.Scope
	schemaHandler *CachedSchemaHandler
}

// connector is the ares connector implementation
type connector struct {
	cfg                ConnectorConfig
	httpClient         http.Client
	upsertBatchBuilder UpsertBatchBuilder
	schemaHandler      *CachedSchemaHandler
}

// ConnectorConfig holds the configurations for ares Connector.
type ConnectorConfig struct {
	// Address is in the format of host:port
	Address string `yaml:"address"`
	// DeviceChoosingTimeout value is the request timeout in seconds for http calls
	// if <= 0, will use default
	Timeout int `yaml:"timeout"`
	// SchemaRefreshInterval is the interval in seconds for the connector to
	// fetch and refresh schema from ares
	// if <= 0, will use default
	SchemaRefreshInterval int `yaml:"schemaRefreshInterval"`
}

func NewUpsertBatchBuilderImpl(logger *zap.SugaredLogger, scope tally.Scope, schemaHandler *CachedSchemaHandler) UpsertBatchBuilder {
	return &UpsertBatchBuilderImpl{
		logger:        logger,
		metricScope:   scope,
		schemaHandler: schemaHandler,
	}
}

// NewConnector returns a new ares Connector
func (cfg ConnectorConfig) NewConnector(logger *zap.SugaredLogger, metricScope tally.Scope) (Connector, error) {
	if cfg.SchemaRefreshInterval <= 0 {
		cfg.SchemaRefreshInterval = defaultSchemaRefreshInterval
	}

	if cfg.Timeout <= 0 {
		cfg.Timeout = defaultRequestTimeout
	}

	httpClient := http.Client{
		Timeout: time.Duration(cfg.Timeout) * time.Second,
	}

	httpSchemaFetcher := NewHttpSchemaFetcher(httpClient, cfg.Address, metricScope)
	cachedSchemaHandler := NewCachedSchemaHandler(logger, metricScope, httpSchemaFetcher)
	err := cachedSchemaHandler.Start(cfg.SchemaRefreshInterval)
	if err != nil {
		return nil, err
	}

	connector := &connector{
		cfg:        cfg,
		httpClient: httpClient,
		upsertBatchBuilder: &UpsertBatchBuilderImpl{
			logger:        logger,
			metricScope:   metricScope,
			schemaHandler: cachedSchemaHandler,
		},
		schemaHandler: cachedSchemaHandler,
	}
	return connector, nil
}

// Insert inserts a batch of rows into ares
func (c *connector) Insert(tableName string, columnNames []string, rows []Row, updateModes ...memCom.ColumnUpdateMode) (int, error) {
	if len(columnNames) == 0 {
		return 0, utils.StackError(nil, "No column names specified")
	}

	// if no update modes at all, use default
	if len(updateModes) == 0 {
		updateModes = make([]memCom.ColumnUpdateMode, len(columnNames))
	}

	if len(updateModes) != len(columnNames) {
		return 0, utils.StackError(nil, "length of column update modes %d does not equal to number of columns %d", len(updateModes), len(columnNames))
	}

	if len(rows) == 0 {
		// Do nothing when there is no row to insert
		return 0, nil
	}

	for _, row := range rows {
		if len(row) != len(columnNames) {
			return 0, utils.StackError(nil,
				"Length of column names should match length of a single row, length of column names :%d, length of row: %d",
				len(columnNames),
				len(row),
			)
		}
	}

	upsertBatchBytes, numRows, err := c.prepareUpsertBatch(tableName, columnNames, updateModes, rows)
	if err != nil {
		return numRows, err
	}

	//TODO: currently always use shard zero for single instance version
	resp, err := c.httpClient.Post(c.dataPath(tableName, 0), dataIngestionHeader, bytes.NewReader(upsertBatchBytes))
	if err != nil || resp.StatusCode != http.StatusOK {
		//TODO: break status code check and error check into two parts for more specific handling like retrying on 5xx
		return 0, utils.StackError(err, "Failed to post upsert batch, table: %s, shard: %d", tableName, 0)
	}

	return numRows, nil
}

// computeHLLValue populate hyperloglog value
func computeHLLValue(dataType memCom.DataType, value interface{}) (uint32, error) {
	var ok bool
	var hashed uint64
	switch dataType {
	case memCom.UUID:
		var v [2]uint64
		v, ok = memCom.ConvertToUUID(value)
		hashed = v[0] ^ v[1]
	case memCom.Uint32:
		var v uint32
		v, ok = memCom.ConvertToUint32(value)
		hashed = utils.Murmur3Sum64(unsafe.Pointer(&v), memCom.DataTypeBytes(dataType), 0)
	case memCom.Int32:
		var v int32
		v, ok = memCom.ConvertToInt32(value)
		hashed = utils.Murmur3Sum64(unsafe.Pointer(&v), memCom.DataTypeBytes(dataType), 0)
	case memCom.Int64:
		var v int64
		v, ok = memCom.ConvertToInt64(value)
		hashed = utils.Murmur3Sum64(unsafe.Pointer(&v), memCom.DataTypeBytes(dataType), 0)
	default:
		return 0, utils.StackError(nil, "invalid type %s for fast hll value", memCom.DataTypeName[dataType])
	}
	if !ok {
		return 0, utils.StackError(nil, "invalid data value %v for data type %s", value, memCom.DataTypeName[dataType])
	}
	return utils.ComputeHLLValue(hashed), nil
}

// prepareUpsertBatch prepares the upsert batch for upsert,
// returns upsertBatch byte array, number of rows in upsert batch and error.
func (c *connector) prepareUpsertBatch(tableName string, columnNames []string, updateModes []memCom.ColumnUpdateMode, rows []Row) ([]byte, int, error) {
	schema, err := c.schemaHandler.FetchSchema(tableName)
	if err != nil {
		return nil, 0, err
	}

	// return error if primary key is missing
	if err = checkPrimaryKeys(schema, columnNames); err != nil {
		return nil, 0, err
	}

	// return error if time column is missing
	if err = checkTimeColumnExistence(schema, columnNames); err != nil {
		return nil, 0, err
	}

	return c.upsertBatchBuilder.PrepareUpsertBatch(tableName, columnNames, updateModes, rows)
}

// checkPrimaryKeys checks whether primary key is missing
func checkPrimaryKeys(schema *TableSchema, columnNames []string) error {
	for _, columnID := range schema.Table.PrimaryKeyColumns {
		pkColumn := schema.Table.Columns[columnID]
		index := utils.IndexOfStr(columnNames, pkColumn.Name)
		if index < 0 {
			return utils.StackError(nil, "Missing primary key column")
		}
	}
	return nil
}

// checkTimeColumnExistence checks if time column is missing for fact table
func checkTimeColumnExistence(schema *TableSchema, columnNames []string) error {
	if !schema.Table.IsFactTable || schema.Table.Config.AllowMissingEventTime {
		return nil
	}

	for _, columnName := range columnNames {
		columnID, exist := schema.ColumnDict[columnName]
		if !exist {
			continue
		}

		if columnID == 0 {
			return nil
		}
	}
	return utils.StackError(nil, "Missing time column")
}

func (c *connector) dataPath(tableName string, shard int) string {
	return fmt.Sprintf("http://%s/data/%s/%d", c.cfg.Address, tableName, shard)
}

func (u *UpsertBatchBuilderImpl) prepareEnumCases(tableName, columnName string, colIndex, columnID int, rows []Row, abandonRows map[int]struct{}, caseInsensitive bool, disableAutoExpand bool) error {
	enumCaseSet := make(map[string]struct{})
	for rowIndex, row := range rows {
		if _, exist := abandonRows[rowIndex]; exist {
			continue
		}
		value := row[colIndex]

		if value == nil {
			continue
		}

		if enumCase, ok := value.(string); ok {
			if caseInsensitive {
				enumCase = strings.ToLower(enumCase)
			}
			enumCaseSet[enumCase] = struct{}{}
		} else {
			u.logger.With(
				"name", "prepareEnumCases",
				"error", "Enum value should be string",
				"table", tableName,
				"columnID", columnID,
				"value", value).Debug("Enum value is not string")
			u.metricScope.Tagged(map[string]string{"table": tableName, "columnID": strconv.Itoa(columnID)}).
				Counter("abandoned_rows").Inc(1)
			abandonRows[rowIndex] = struct{}{}
		}
	}

	if len(enumCaseSet) > 0 {
		enumCases := make([]string, 0, len(enumCaseSet))
		for enumCase := range enumCaseSet {
			enumCases = append(enumCases, enumCase)
		}
		err := u.schemaHandler.PrepareEnumCases(tableName, columnName, enumCases)
		if err != nil {
			return err
		}
	}
	return nil
}

// PrepareUpsertBatch prepares the upsert batch for upsert,
// returns upsertBatch byte array, number of rows in upsert batch and error.
func (u *UpsertBatchBuilderImpl) PrepareUpsertBatch(tableName string, columnNames []string,
	updateModes []memCom.ColumnUpdateMode, rows []Row) ([]byte, int, error) {
	var err error
	upsertBatchBuilder := memCom.NewUpsertBatchBuilder()

	schema, err := u.schemaHandler.FetchSchema(tableName)
	if err != nil {
		return nil, 0, err
	}

	// use abandonRows to record abandoned row index due to invalid data
	abandonRows := make(map[int]struct{})

	for colIndex, columnName := range columnNames {
		columnID, exist := schema.ColumnDict[columnName]
		if !exist {
			continue
		}
		column := schema.Table.Columns[columnID]

		// following conditions only overwrite is supported:
		// 1. dimension table (TODO: might support min/max in the future if needed)
		// 2. primary key column
		// 3. archiving sort column
		// 4. data type not in uint8, int8, uint16, int16, uint32, int32, float32
		if (!schema.Table.IsFactTable ||
			utils.IndexOfInt(schema.Table.PrimaryKeyColumns, columnID) >= 0 ||
			utils.IndexOfInt(schema.Table.ArchivingSortColumns, columnID) >= 0 ||
			schema.Table.Columns[columnID].IsOverwriteOnlyDataType()) &&
			updateModes[colIndex] > memCom.UpdateForceOverwrite {
			return nil, 0, utils.StackError(nil, "column %s only supports overwrite", columnName)
		}

		dataType := memCom.DataTypeForColumn(column)
		if err = upsertBatchBuilder.AddColumnWithUpdateMode(columnID, dataType, updateModes[colIndex]); err != nil {
			return nil, 0, err
		}

		if column.IsEnumColumn() {
			if err = u.prepareEnumCases(tableName, columnName, colIndex, columnID, rows, abandonRows, column.CaseInsensitive, column.DisableAutoExpand); err != nil {
				return nil, 0, err
			}
		}
	}

	for rowIndex, row := range rows {
		if _, exist := abandonRows[rowIndex]; exist {
			continue
		}
		upsertBatchBuilder.AddRow()

		upsertBatchColumnIndex := 0
		for inputColIndex, columnName := range columnNames {
			columnID, exist := schema.ColumnDict[columnName]
			if !exist {
				continue
			}
			column := schema.Table.Columns[columnID]

			value := row[inputColIndex]

			// prevent primary key being nil
			if value == nil && utils.IndexOfInt(schema.Table.PrimaryKeyColumns, columnID) >= 0 {
				upsertBatchBuilder.RemoveRow()
				u.logger.With(
					"name", "PrepareUpsertBatch",
					"table", tableName,
					"columnID", columnID,
					"value", value).Error("PrimaryKey column is nil")
				break
			}

			// skip rows if time column is nil for fact table
			if value == nil && schema.Table.IsFactTable && !schema.Table.Config.AllowMissingEventTime && columnID == 0 {
				upsertBatchBuilder.RemoveRow()
				u.logger.With(
					"name", "PrepareUpsertBatch",
					"table", tableName,
					"columnID", columnID,
					"value", value).Error("Time column is nil")
				break
			}

			if column.IsEnumColumn() {
				value, err = u.schemaHandler.TranslateEnum(tableName, columnID, value, column.CaseInsensitive)
				if err != nil {
					upsertBatchBuilder.RemoveRow()
					u.logger.With(
						"name", "prepareUpsertBatch",
						"error", err.Error(),
						"table", tableName,
						"columnID", columnID,
						"value", value).Error("Failed to translate enum")
					break
				}

				// If enum value is not found from predefined enum cases and default value is not set, we set it to nil.
				if value == -1 {
					value = nil
				}
			}

			// Set value to the last row.
			// compute hll value to insert
			if column.HLLConfig.IsHLLColumn {
				// here use original column data type to compute hll value
				value, err = computeHLLValue(memCom.DataTypeFromString(column.Type), value)
				if err != nil {
					upsertBatchBuilder.RemoveRow()
					u.logger.With("name", "PrepareUpsertBatch", "error", err.Error(), "table", tableName, "columnID", columnID, "value", value).Error("Failed to set value")
					break
				}
				if err = upsertBatchBuilder.SetValue(upsertBatchBuilder.NumRows-1, upsertBatchColumnIndex, value); err != nil {
					upsertBatchBuilder.RemoveRow()
					u.logger.With("name", "PrepareUpsertBatch", "error", err.Error(), "table", tableName, "columnID", columnID, "value", value).Error("Failed to set value")
					break
				}
			} else {
				// directly insert value
				if err = upsertBatchBuilder.SetValue(upsertBatchBuilder.NumRows-1, upsertBatchColumnIndex, value); err != nil {
					upsertBatchBuilder.RemoveRow()
					u.logger.With("name", "PrepareUpsertBatch", "error", err.Error(), "table", tableName, "columnID", columnID, "value", value).Error("Failed to set value")
					break
				}
			}
			upsertBatchColumnIndex++
		}
	}

	batchBytes, err := upsertBatchBuilder.ToByteArray()
	return batchBytes, upsertBatchBuilder.NumRows, err
}