connection.go

// Copyright (c) 2025 ADBC Drivers Contributors
//
// 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 mysql

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync/atomic"

	"github.com/adbc-drivers/driverbase-go/driverbase"
	sqlwrapper "github.com/adbc-drivers/driverbase-go/sqlwrapper"
	"github.com/apache/arrow-adbc/go/adbc"
	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	gomysql "github.com/go-sql-driver/mysql"
)

var loadReaderCounter atomic.Uint64

const (
	// Default num of rows per batch for batched INSERT
	MySQLDefaultIngestBatchSize = 1000
	// MySQL's maximum number of placeholders in a prepared statement
	MySQLMaxPlaceholders = 65535
)

// GetCurrentCatalog implements driverbase.CurrentNamespacer.
func (c *mysqlConnectionImpl) GetCurrentCatalog() (string, error) {
	var database string
	err := c.Db.QueryRowContext(context.Background(), "SELECT DATABASE()").Scan(&database)
	if err != nil {
		return "", c.ErrorHelper.WrapIO(err, "failed to get current database")
	}
	if database == "" {
		return "", c.ErrorHelper.InvalidState("no current database set")
	}
	return database, nil
}

// GetCurrentDbSchema implements driverbase.CurrentNamespacer.
func (c *mysqlConnectionImpl) GetCurrentDbSchema() (string, error) {
	return "", nil
}

// SetCurrentCatalog implements driverbase.CurrentNamespacer.
func (c *mysqlConnectionImpl) SetCurrentCatalog(catalog string) error {
	_, err := c.Db.ExecContext(context.Background(), "USE "+quoteIdentifier(catalog))
	return err
}

// SetCurrentDbSchema implements driverbase.CurrentNamespacer.
func (c *mysqlConnectionImpl) SetCurrentDbSchema(schema string) error {
	if schema != "" {
		return c.ErrorHelper.InvalidArgument("cannot set schema in MySQL: schemas are not supported")
	}
	return nil
}

func (c *mysqlConnectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc.InfoCode) error {
	if c.version == "" {
		var version, comment string
		if err := c.Conn.QueryRowContext(ctx, "SELECT @@version, @@version_comment").Scan(&version, &comment); err != nil {
			return c.ErrorHelper.WrapIO(err, "failed to get version")
		}
		c.version = fmt.Sprintf("%s (%s)", version, comment)
	}
	return c.DriverInfo.RegisterInfoCode(adbc.InfoVendorVersion, c.version)
}

// GetTableSchema returns the Arrow schema for a MySQL table
func (c *mysqlConnectionImpl) GetTableSchema(ctx context.Context, catalog *string, dbSchema *string, tableName string) (schema *arrow.Schema, err error) {
	// Struct to capture MySQL column information
	type tableColumn struct {
		OrdinalPosition        int32
		ColumnName             string
		DataType               string
		IsNullable             string
		CharacterMaximumLength sql.NullInt64
		NumericPrecision       sql.NullInt64
		NumericScale           sql.NullInt64
	}

	query := `SELECT
		ORDINAL_POSITION,
		COLUMN_NAME,
		DATA_TYPE,
		IS_NULLABLE,
		CHARACTER_MAXIMUM_LENGTH,
		NUMERIC_PRECISION,
		NUMERIC_SCALE
	FROM INFORMATION_SCHEMA.COLUMNS
	WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?
	ORDER BY ORDINAL_POSITION`

	var args []any
	if catalog != nil && *catalog != "" {
		// Use specified catalog (database)
		args = []any{*catalog, tableName}
	} else {
		// Use current database
		currentDB, err := c.GetCurrentCatalog()
		if err != nil {
			return nil, err
		}
		args = []any{currentDB, tableName}
	}

	// Execute query to get column information
	rows, err := c.Conn.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, c.ErrorHelper.WrapIO(err, "failed to query table schema")
	}
	defer func() {
		err = errors.Join(err, rows.Close())
	}()

	var columns []tableColumn
	for rows.Next() {
		var col tableColumn
		err := rows.Scan(
			&col.OrdinalPosition,
			&col.ColumnName,
			&col.DataType,
			&col.IsNullable,
			&col.CharacterMaximumLength,
			&col.NumericPrecision,
			&col.NumericScale,
		)
		if err != nil {
			return nil, c.ErrorHelper.WrapIO(err, "failed to scan column information")
		}
		columns = append(columns, col)
	}

	if err := rows.Err(); err != nil {
		return nil, c.ErrorHelper.WrapIO(err, "rows error")
	}

	if len(columns) == 0 {
		return nil, c.ErrorHelper.NotFound("table not found: %s", tableName)
	}

	// Build Arrow schema from column information using type converter
	fields := make([]arrow.Field, len(columns))
	for i, col := range columns {
		// Create ColumnType struct for the type converter
		var length, precision, scale *int64
		if col.CharacterMaximumLength.Valid {
			length = &col.CharacterMaximumLength.Int64
		}
		if col.NumericPrecision.Valid {
			precision = &col.NumericPrecision.Int64
		}
		if col.NumericScale.Valid {
			scale = &col.NumericScale.Int64
		}

		colType := sqlwrapper.ColumnType{
			Name:             col.ColumnName,
			DatabaseTypeName: col.DataType,
			Nullable:         col.IsNullable == "YES",
			Length:           length,
			Precision:        precision,
			Scale:            scale,
		}

		arrowType, nullable, metadata, err := c.TypeConverter.ConvertRawColumnType(colType)
		if err != nil {
			return nil, c.ErrorHelper.WrapIO(err, "failed to convert column type for %s", col.ColumnName)
		}

		fields[i] = arrow.Field{
			Name:     col.ColumnName,
			Type:     arrowType,
			Nullable: nullable,
			Metadata: metadata,
		}
	}

	return arrow.NewSchema(fields, nil), nil
}

// ListTableTypes implements driverbase.TableTypeLister interface
func (c *mysqlConnectionImpl) ListTableTypes(ctx context.Context) ([]string, error) {
	// MySQL supports these standard table types
	return []string{
		"BASE TABLE",  // Regular tables
		"VIEW",        // Views
		"SYSTEM VIEW", // System/information schema views
	}, nil
}

// QuoteIdentifier implements BulkIngester
func (c *mysqlConnectionImpl) QuoteIdentifier(name string) string {
	return quoteIdentifier(name)
}

// GetPlaceholder implements BulkIngester
func (c *mysqlConnectionImpl) GetPlaceholder(field *arrow.Field, index int) string {
	return "?"
}

// Ensure mysqlConnectionImpl implements BulkIngester
var _ sqlwrapper.BulkIngester = (*mysqlConnectionImpl)(nil)

// ExecuteBulkIngest performs MySQL bulk ingest using LOAD DATA LOCAL INFILE with a fallback to batched INSERTs.
func (c *mysqlConnectionImpl) ExecuteBulkIngest(ctx context.Context, conn *sqlwrapper.LoggingConn, options *driverbase.BulkIngestOptions, stream array.RecordReader) (rowCount int64, err error) {
	schema := stream.Schema()
	if err := c.createTableIfNeeded(ctx, conn, options.TableName, schema, options); err != nil {
		return -1, c.ErrorHelper.WrapIO(err, "failed to create table")
	}

	if c.isLoadDataEnabled(ctx, conn) {
		return c.executeLoadDataIngest(ctx, conn, options, stream)
	}

	// Validate MySQL-specific options
	if options.MaxQuerySizeBytes > 0 {
		return -1, c.ErrorHelper.InvalidArgument(
			"MySQL driver does not support '%s'. "+
				"Use '%s' instead to control the number of rows per INSERT statement.",
			driverbase.OptionKeyIngestMaxQuerySizeBytes,
			driverbase.OptionKeyIngestBatchSize)
	}

	// Set MySQL-specific default batch size if user hasn't overridden,
	// capping to stay within MySQL's 65,535 placeholder limit.
	numCols := len(schema.Fields())
	maxBatchSize := MySQLMaxPlaceholders / numCols
	if options.IngestBatchSize == 0 {
		options.IngestBatchSize = min(MySQLDefaultIngestBatchSize, maxBatchSize)
	} else if options.IngestBatchSize > maxBatchSize {
		options.IngestBatchSize = maxBatchSize
	}

	return sqlwrapper.ExecuteBatchedBulkIngest(
		ctx, conn, options, stream,
		c.TypeConverter, c, &c.Base().ErrorHelper,
	)
}

// isLoadDataEnabled checks if LOAD DATA LOCAL INFILE is enabled on the server.
func (c *mysqlConnectionImpl) isLoadDataEnabled(ctx context.Context, conn *sqlwrapper.LoggingConn) bool {
	var localInfile int
	err := conn.QueryRowContext(ctx, "SELECT @@local_infile").Scan(&localInfile)
	return err == nil && localInfile == 1
}

// executeLoadDataIngest performs bulk ingestion using the LOAD DATA LOCAL INFILE command.
func (c *mysqlConnectionImpl) executeLoadDataIngest(ctx context.Context, conn *sqlwrapper.LoggingConn, options *driverbase.BulkIngestOptions, stream array.RecordReader) (int64, error) {
	r, w := io.Pipe()
	readerId := loadReaderCounter.Add(1)
	readerName := fmt.Sprintf("adbc_ingest_%s_%d", options.TableName, readerId)

	gomysql.RegisterReaderHandler(readerName, func() io.Reader {
		return r
	})
	defer gomysql.DeregisterReaderHandler(readerName)
	batchSize := options.IngestBatchSize
	if batchSize <= 0 {
		batchSize = 10000 // Default batch size for streaming chunks
	}

	it, err := sqlwrapper.NewRowBufferIterator(stream, batchSize, c.TypeConverter)
	if err != nil {
		return -1, c.ErrorHelper.WrapIO(err, "failed to create row buffer iterator")
	}

	numCols := len(stream.Schema().Fields())
	go func() {
		config := CSVConfig{
			FieldDelimiter:  '\t',
			LineTerminator:  '\n',
			NullValue:       "\\N",
			EscapeBackslash: true,
		}
		err := arrowToCSV(ctx, w, it, numCols, config)
		if err != nil {
			_ = w.CloseWithError(err)
		} else {
			_ = w.Close()
		}
	}()

	var colNames []string
	for _, field := range stream.Schema().Fields() {
		colNames = append(colNames, quoteIdentifier(field.Name))
	}
	colsList := strings.Join(colNames, ", ")

	query := fmt.Sprintf(
		"LOAD DATA LOCAL INFILE 'Reader::%s' INTO TABLE %s CHARACTER SET utf8mb4 FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' LINES TERMINATED BY '\\n' (%s)",
		readerName, c.QuoteIdentifier(options.TableName), colsList,
	)

	res, err := conn.ExecContext(ctx, query)
	if err != nil {
		return -1, c.ErrorHelper.WrapIO(err, "failed to execute LOAD DATA statement")
	}

	rowCount, err := res.RowsAffected()
	if err != nil {
		return -1, c.ErrorHelper.WrapIO(err, "failed to get rows affected")
	}

	return rowCount, nil
}

// createTableIfNeeded creates the table based on the ingest mode
func (c *mysqlConnectionImpl) createTableIfNeeded(ctx context.Context, conn *sqlwrapper.LoggingConn, tableName string, schema *arrow.Schema, options *driverbase.BulkIngestOptions) error {
	switch options.Mode {
	case adbc.OptionValueIngestModeCreate:
		// Create the table (fail if exists)
		return c.createTable(ctx, conn, tableName, schema, false, options.Temporary)
	case adbc.OptionValueIngestModeCreateAppend:
		// Create the table if it doesn't exist
		return c.createTable(ctx, conn, tableName, schema, true, options.Temporary)
	case adbc.OptionValueIngestModeReplace:
		// Drop and recreate the table
		if err := c.dropTable(ctx, conn, tableName, options.Temporary); err != nil {
			return err
		}
		return c.createTable(ctx, conn, tableName, schema, false, options.Temporary)
	case adbc.OptionValueIngestModeAppend:
		// Table should already exist, do nothing
		return nil
	default:
		return c.ErrorHelper.InvalidArgument("unsupported ingest mode: %s", options.Mode)
	}
}

// createTable creates a MySQL table from Arrow schema
func (c *mysqlConnectionImpl) createTable(ctx context.Context, conn *sqlwrapper.LoggingConn, tableName string, schema *arrow.Schema, ifNotExists bool, temporary bool) error {
	var queryBuilder strings.Builder
	if temporary {
		queryBuilder.WriteString("CREATE TEMPORARY TABLE ")
	} else {
		queryBuilder.WriteString("CREATE TABLE ")
	}
	if ifNotExists {
		queryBuilder.WriteString("IF NOT EXISTS ")
	}
	queryBuilder.WriteString(quoteIdentifier(tableName))
	queryBuilder.WriteString(" (")

	for i, field := range schema.Fields() {
		if i > 0 {
			queryBuilder.WriteString(", ")
		}

		queryBuilder.WriteString(quoteIdentifier(field.Name))
		queryBuilder.WriteString(" ")

		// Convert Arrow type to MySQL type
		mysqlType := c.arrowToMySQLType(field.Type, field.Nullable)
		queryBuilder.WriteString(mysqlType)
	}

	queryBuilder.WriteString(")")

	_, err := conn.ExecContext(ctx, queryBuilder.String())
	return err
}

// dropTable drops a MySQL table
func (c *mysqlConnectionImpl) dropTable(ctx context.Context, conn *sqlwrapper.LoggingConn, tableName string, temporary bool) error {
	keyword := "TABLE"
	if temporary {
		keyword = "TEMPORARY TABLE"
	}
	dropSQL := fmt.Sprintf("DROP %s IF EXISTS %s", keyword, quoteIdentifier(tableName))
	_, err := conn.ExecContext(ctx, dropSQL)
	return err
}

// arrowToMySQLType converts Arrow data type to MySQL column type
func (c *mysqlConnectionImpl) arrowToMySQLType(arrowType arrow.DataType, nullable bool) string {
	var mysqlType string

	switch arrowType := arrowType.(type) {
	case *arrow.BooleanType:
		mysqlType = "BOOLEAN"
	case *arrow.Int8Type:
		mysqlType = "TINYINT"
	case *arrow.Int16Type:
		mysqlType = "SMALLINT"
	case *arrow.Int32Type:
		mysqlType = "INT"
	case *arrow.Int64Type:
		mysqlType = "BIGINT"
	case *arrow.Float32Type:
		mysqlType = "FLOAT"
	case *arrow.Float64Type:
		mysqlType = "DOUBLE"
	case *arrow.StringType:
		mysqlType = "TEXT"
	case *arrow.BinaryType, *arrow.FixedSizeBinaryType, *arrow.BinaryViewType:
		mysqlType = "BLOB"
	case *arrow.LargeBinaryType:
		mysqlType = "LONGBLOB"
	case *arrow.Date32Type:
		mysqlType = "DATE"
	case *arrow.TimestampType:

		// Determine precision based on Arrow timestamp unit
		var precision string
		switch arrowType.Unit {
		case arrow.Second:
			precision = ""
		case arrow.Millisecond:
			precision = "(3)"
		case arrow.Microsecond:
			precision = "(6)"
		case arrow.Nanosecond:
			precision = "(6)" // MySQL max is 6 digits
		default:
			// should never happen, but panic here for defensive programming
			panic(fmt.Sprintf("unexpected Arrow timestamp unit: %v", arrowType.Unit))
		}

		// Use DATETIME for timezone-naive timestamps, TIMESTAMP for timezone-aware
		if arrowType.TimeZone != "" {
			// Timezone-aware (timestamptz) -> TIMESTAMP
			mysqlType = "TIMESTAMP" + precision
		} else {
			// Timezone-naive (timestamp) -> DATETIME
			mysqlType = "DATETIME" + precision
		}
	case *arrow.Time32Type:
		// Determine precision based on Arrow time unit
		switch arrowType.Unit {
		case arrow.Second:
			mysqlType = "TIME"
		case arrow.Millisecond:
			mysqlType = "TIME(3)"
		default:
			// should never happen, but panic here for defensive programming
			panic(fmt.Sprintf("unexpected Time32 unit: %v", arrowType.Unit))
		}

	case *arrow.Time64Type:
		// Determine precision based on Arrow time unit
		switch arrowType.Unit {
		case arrow.Microsecond:
			mysqlType = "TIME(6)"
		case arrow.Nanosecond:
			mysqlType = "TIME(6)" // MySQL max is 6 digits
		default:
			// should never happen, but panic here for defensive programming
			panic(fmt.Sprintf("unexpected Time64 unit: %v", arrowType.Unit))
		}
	case arrow.DecimalType:
		mysqlType = fmt.Sprintf("DECIMAL(%d,%d)", arrowType.GetPrecision(), arrowType.GetScale())
	default:
		// Default to TEXT for unknown types
		mysqlType = "TEXT"
	}

	if nullable {
		mysqlType += " NULL"
	} else {
		mysqlType += " NOT NULL"
	}

	return mysqlType
}

func quoteIdentifier(name string) string {
	return "`" + strings.ReplaceAll(name, "`", "``") + "`"
}