connection.go

package database

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"
	"reflect"
	"strings"
	"sync/atomic"
	"time"

	"github.com/pluto-org-co/sqlite"
	"gorm.io/gorm"
	"modernc.org/libc"
	sqlite3 "modernc.org/sqlite/lib"
)

type DriverName string

const DriverNameSQLite DriverName = "sqlite"

type Connection struct {
	Driver             DriverName
	TransactionCounter atomic.Int64
	Logger             *slog.Logger
	*gorm.DB
}

func CreateBatchSize[T any](c *Connection) (size int, err error) {
	var dst T
	columns, err := c.Migrator().ColumnTypes(&dst)
	if err != nil {
		return 0, fmt.Errorf("failed to retrieve table type: %w", err)
	}

	if c.Driver == DriverNameSQLite {
		return sqlite3.SQLITE_MAX_VARIABLE_NUMBER / len(columns), nil
	}

	return 1_000, nil
}

func (conn *Connection) Close() (err error) {
	db, err := conn.DB.DB()
	if err == nil {
		err = db.Close()
	}

	return
}

type Config struct {
	Logger           *slog.Logger
	DSN              string      `yaml:"dsn"`
	ConnectionConfig gorm.Config `yaml:"-"`
}

// Easier local testing Database. Should not be used in production.
func SQLite(config Config) (conn *Connection, err error) {
	logger := config.Logger.With("action", "Configuring SQLite")
	const MaxBatchSize = 1_000

	const SetupScript = `
	PRAGMA journal_mode=WAL;
	-- - PRAGMA locking_mode=EXCLUSIVE;
	PRAGMA foreign_keys = ON;
	PRAGMA defer_foreign_keys = ON;
	PRAGMA wal_checkpoint(RESTART);
	PRAGMA auto_vacuum = INCREMENTAL;
	`

	config.ConnectionConfig.PrepareStmt = true
	config.ConnectionConfig.CreateBatchSize = MaxBatchSize
	db, err := gorm.Open(sqlite.Open(config.DSN), &config.ConnectionConfig)
	if err != nil {
		err = fmt.Errorf("failed to connect to sqlite database: %s: %w", config.DSN, err)
		return
	}
	defer func() {
		if err == nil {
			return
		}

		sqlDB, err := db.DB()
		if err != nil {
			logger.Error("failed to retrieve SQLite driver", "error-msg", err)
			return
		}
		sqlDB.Close()
	}()

	sqlDB, err := db.DB()
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve conn: %w", err)
	}
	dbConn, err := sqlDB.Conn(context.TODO())
	if err != nil {
		return nil, fmt.Errorf("failed to get db conn: %w", err)
	}
	err = dbConn.Raw(func(driverConn any) error {
		type conn struct {
			db  uintptr // *sqlite3.Xsqlite3
			tls *libc.TLS
		}

		connValue := (*conn)(reflect.ValueOf(driverConn).UnsafePointer())

		out := sqlite3.Xsqlite3_limit(connValue.tls, connValue.db, sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER, sqlite3.SQLITE_MAX_VARIABLE_NUMBER)
		if out == sqlite3.SQLITE_MAX_VARIABLE_NUMBER {
			logger.Debug("Set SQLite variable limit", "new-value", out)
			return nil
		}
		return fmt.Errorf("Xsqlite3_limit error: %d", out)
	})
	if err != nil {
		return nil, fmt.Errorf("failed to apply new limit: %w", err)
	}

	err = db.
		Exec(SetupScript).
		Error
	if err != nil {
		err = fmt.Errorf("failed to execute setup script: %w", err)
		return
	}

	conn = &Connection{
		Driver: DriverNameSQLite,
		Logger: config.Logger,
		DB:     db,
	}
	if conn.Logger == nil {
		conn.Logger = slog.Default()
	}
	return
}

const MaxAttempts = 1_000_000

var ErrMaxAttemptsExceeded = errors.New("max numbers of attempts exceeded")

func (c *Connection) Transaction(ctx context.Context, fc func(tx *gorm.DB) (err error), opts ...*sql.TxOptions) (err error) {
	var attempts int
	txN := c.TransactionCounter.Add(1)
	now := time.Now()

	logger := c.Logger.With("#", txN, "start-time", now)

	logger.Debug("Transaction created")
	defer func() {
		if err != nil {
			logger.Error("Transaction failed", "duration", time.Since(now), "total-attempts", attempts, "error-msg", err)
		} else {
			logger.Debug("Transaction succeed", "duration", time.Since(now), "total-attempts", attempts)
		}
	}()

	for ; attempts < MaxAttempts; attempts++ {
		logger := logger.With("attempt", attempts)

		select {
		case <-ctx.Done():
			logger.Debug("Context trigger")
			err = ctx.Err()
			if err != nil {
				return fmt.Errorf("context error: %w", err)
			}
			return nil
		default:
			logger.Debug("Running transaction")
			err = c.DB.Transaction(func(tx *gorm.DB) (err error) {
				tx = tx.WithContext(ctx)

				err = fc(tx)
				if err != nil {
					return fmt.Errorf("failed to handle function: %w", err)
				}
				return nil
			}, opts...)
			if err == nil {
				return nil
			}
			// This usually happens with SQLite when multiple concurrent writers try to insert/update at the same time
			if strings.Contains(err.Error(), "database is locked") {
				logger.Warn("Database locked")
				continue
			}

			return fmt.Errorf("failed to execute transaction: %w", err)
		}
	}
	return ErrMaxAttemptsExceeded
}