parameters.go

package dbsql

import (
	"database/sql/driver"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/databricks/databricks-sql-go/internal/cli_service"
)

type Parameter struct {
	Name  string
	Type  SqlType
	Value any
}

type SqlType int

const (
	SqlUnkown SqlType = iota
	SqlString
	SqlDate
	SqlTimestamp
	SqlFloat
	SqlDecimal
	SqlDouble
	SqlInteger
	SqlBigInt
	SqlSmallInt
	SqlTinyInt
	SqlBoolean
	SqlIntervalMonth
	SqlIntervalDay
)

func (s SqlType) String() string {
	switch s {
	case SqlString:
		return "STRING"
	case SqlDate:
		return "DATE"
	case SqlTimestamp:
		return "TIMESTAMP"
	case SqlFloat:
		return "FLOAT"
	case SqlDecimal:
		return "DECIMAL"
	case SqlDouble:
		return "DOUBLE"
	case SqlInteger:
		return "INTEGER"
	case SqlBigInt:
		return "BIGINT"
	case SqlSmallInt:
		return "SMALLINT"
	case SqlTinyInt:
		return "TINYINT"
	case SqlBoolean:
		return "BOOLEAN"
	case SqlIntervalMonth:
		return "INTERVAL MONTH"
	case SqlIntervalDay:
		return "INTERVAL DAY"
	}
	return "unknown"
}

func valuesToParameters(namedValues []driver.NamedValue) []Parameter {
	var params []Parameter
	for i := range namedValues {
		newParam := *new(Parameter)
		namedValue := namedValues[i]
		param, ok := namedValue.Value.(Parameter)
		if ok {
			newParam.Name = param.Name
			newParam.Value = param.Value
			newParam.Type = param.Type
		} else {
			newParam.Name = namedValue.Name
			newParam.Value = namedValue.Value
		}
		params = append(params, newParam)
	}
	return params
}

func inferTypes(params []Parameter) {
	for i := range params {
		param := &params[i]
		if param.Type == SqlUnkown {
			switch value := param.Value.(type) {
			case bool:
				param.Value = strconv.FormatBool(value)
				param.Type = SqlBoolean
			case string:
				param.Value = value
				param.Type = SqlString
			case int:
				param.Value = strconv.Itoa(value)
				param.Type = SqlInteger
			case uint:
				param.Value = strconv.FormatUint(uint64(value), 10)
				param.Type = SqlInteger
			case int8:
				param.Value = strconv.Itoa(int(value))
				param.Type = SqlInteger
			case uint8:
				param.Value = strconv.FormatUint(uint64(value), 10)
				param.Type = SqlInteger
			case int16:
				param.Value = strconv.Itoa(int(value))
				param.Type = SqlInteger
			case uint16:
				param.Value = strconv.FormatUint(uint64(value), 10)
				param.Type = SqlInteger
			case int32:
				param.Value = strconv.Itoa(int(value))
				param.Type = SqlInteger
			case uint32:
				param.Value = strconv.FormatUint(uint64(value), 10)
				param.Type = SqlInteger
			case int64:
				param.Value = strconv.Itoa(int(value))
				param.Type = SqlInteger
			case uint64:
				param.Value = strconv.FormatUint(uint64(value), 10)
				param.Type = SqlInteger
			case float32:
				param.Value = strconv.FormatFloat(float64(value), 'f', -1, 32)
				param.Type = SqlFloat
			case time.Time:
				param.Value = value.String()
				param.Type = SqlTimestamp
			default:
				s := fmt.Sprintf("%s", param.Value)
				param.Value = s
				param.Type = SqlString
			}
		}
	}
}
func convertNamedValuesToSparkParams(values []driver.NamedValue) []*cli_service.TSparkParameter {
	var sparkParams []*cli_service.TSparkParameter

	sqlParams := valuesToParameters(values)
	inferTypes(sqlParams)
	for i := range sqlParams {
		sqlParam := sqlParams[i]
		sparkParamValue := sqlParam.Value.(string)
		var sparkParamType string
		if sqlParam.Type == SqlDecimal {
			sparkParamType = inferDecimalType(sparkParamValue)
		} else {
			sparkParamType = sqlParam.Type.String()
		}
		sparkParam := cli_service.TSparkParameter{Name: &sqlParam.Name, Type: &sparkParamType, Value: &cli_service.TSparkParameterValue{StringValue: &sparkParamValue}}
		sparkParams = append(sparkParams, &sparkParam)
	}
	return sparkParams
}

func inferDecimalType(d string) (t string) {
	var overall int
	var after int
	if strings.HasPrefix(d, "0.") {
		// Less than one
		overall = len(d) - 2
		after = len(d) - 2
	} else if !strings.Contains(d, ".") {
		// Less than one
		overall = len(d)
		after = 0
	} else {
		components := strings.Split(d, ".")
		overall, after = len(components[0])+len(components[1]), len(components[1])
	}

	return fmt.Sprintf("DECIMAL(%d,%d)", overall, after)
}