sqlite.go

package sqlite

import (
	"context"
	"database/sql"
	"errors"
	"net/url"
	"os"
	"strconv"
	"strings"

	"gorm.io/gorm"
	"gorm.io/gorm/callbacks"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/migrator"
	"gorm.io/gorm/schema"
	moderncsqlite "modernc.org/sqlite"
	moderncsqlitelib "modernc.org/sqlite/lib"
)

// DriverName is the default driver name for SQLite.
const DriverName = "sqlite"

// init turns on modernc.org/sqlite's process-wide StrictPragmas, OFDLocking
// and Fdatasync switches, which are off upstream. An application can opt out
// by calling any of them with false before opening its first database.
func init() {
	moderncsqlite.StrictPragmas(true)
	// respect an explicit MODERNC_SQLITE_OFD_LOCK; errors only mean the
	// platform lacks the feature (OFD locks are Linux-only, fdatasync is
	// unavailable on macOS and Windows), and the switch then stays off
	if _, ok := os.LookupEnv("MODERNC_SQLITE_OFD_LOCK"); !ok {
		_, _ = moderncsqlite.OFDLocking(true)
	}
	_, _ = moderncsqlite.Fdatasync(true)
}

type Dialector struct {
	DriverName string
	DSN        string
	Conn       gorm.ConnPool
}

type Config struct {
	DriverName string
	DSN        string
	Conn       gorm.ConnPool
}

func Open(dsn string) gorm.Dialector {
	return &Dialector{DSN: dsn}
}

func New(config Config) gorm.Dialector {
	return &Dialector{DSN: config.DSN, DriverName: config.DriverName, Conn: config.Conn}
}

func (d Dialector) Name() string {
	return "sqlite"
}

func (d Dialector) Initialize(db *gorm.DB) (err error) {
	if d.DriverName == "" {
		d.DriverName = DriverName
	}

	if d.Conn != nil {
		db.ConnPool = d.Conn
	} else {
		conn, err := sql.Open(d.DriverName, injectDSNParams(d.DSN))
		if err != nil {
			return err
		}
		db.ConnPool = conn
	}

	var version string
	if err := db.ConnPool.QueryRowContext(context.Background(), "select sqlite_version()").Scan(&version); err != nil {
		return err
	}
	// https://www.sqlite.org/releaselog/3_35_0.html
	if compareVersion(version, "3.35.0") >= 0 {
		callbacks.RegisterDefaultCallbacks(db, &callbacks.Config{
			CreateClauses:        []string{"INSERT", "VALUES", "ON CONFLICT", "RETURNING"},
			UpdateClauses:        []string{"UPDATE", "SET", "FROM", "WHERE", "RETURNING"},
			DeleteClauses:        []string{"DELETE", "FROM", "WHERE", "RETURNING"},
			LastInsertIDReversed: true,
		})
	} else {
		callbacks.RegisterDefaultCallbacks(db, &callbacks.Config{
			LastInsertIDReversed: true,
		})
	}

	for k, v := range d.ClauseBuilders() {
		if _, ok := db.ClauseBuilders[k]; !ok {
			db.ClauseBuilders[k] = v
		}
	}
	return
}

func (d Dialector) ClauseBuilders() map[string]clause.ClauseBuilder {
	return map[string]clause.ClauseBuilder{
		"INSERT": func(c clause.Clause, builder clause.Builder) {
			if insert, ok := c.Expression.(clause.Insert); ok {
				if stmt, ok := builder.(*gorm.Statement); ok {
					_, _ = stmt.WriteString("INSERT ")
					if insert.Modifier != "" {
						_, _ = stmt.WriteString(insert.Modifier)
						_ = stmt.WriteByte(' ')
					}

					_, _ = stmt.WriteString("INTO ")
					if insert.Table.Name == "" {
						stmt.WriteQuoted(stmt.Table)
					} else {
						stmt.WriteQuoted(insert.Table)
					}
					return
				}
			}

			c.Build(builder)
		},
		"LIMIT": func(c clause.Clause, builder clause.Builder) {
			if limit, ok := c.Expression.(clause.Limit); ok {
				var lmt = -1
				if limit.Limit != nil && *limit.Limit >= 0 {
					lmt = *limit.Limit
				}
				if lmt >= 0 || limit.Offset > 0 {
					_, _ = builder.WriteString("LIMIT ")
					_, _ = builder.WriteString(strconv.Itoa(lmt))
				}
				if limit.Offset > 0 {
					_, _ = builder.WriteString(" OFFSET ")
					_, _ = builder.WriteString(strconv.Itoa(limit.Offset))
				}
			}
		},
		"FOR": func(c clause.Clause, builder clause.Builder) {
			if _, ok := c.Expression.(clause.Locking); ok {
				// SQLite3 does not support row-level locking.
				return
			}
			c.Build(builder)
		},
	}
}

func (d Dialector) DefaultValueOf(field *schema.Field) clause.Expression {
	if field.AutoIncrement {
		return clause.Expr{SQL: "NULL"}
	}

	// doesn't work, will raise error
	return clause.Expr{SQL: "DEFAULT"}
}

func (d Dialector) Migrator(db *gorm.DB) gorm.Migrator {
	return &Migrator{migrator.Migrator{Config: migrator.Config{
		DB:                          db,
		Dialector:                   d,
		CreateIndexAfterCreateTable: true,
	}}}
}

func (d Dialector) BindVarTo(writer clause.Writer, stmt *gorm.Statement, v any) {
	_ = writer.WriteByte('?')
}

func (d Dialector) QuoteTo(writer clause.Writer, str string) {
	var (
		underQuoted, selfQuoted bool
		continuousBacktick      int8
		shiftDelimiter          int8
	)

	for _, v := range []byte(str) {
		switch v {
		case '`':
			continuousBacktick++
			if continuousBacktick == 2 {
				_, _ = writer.WriteString("``")
				continuousBacktick = 0
			}
		case '.':
			if continuousBacktick > 0 || !selfQuoted {
				shiftDelimiter = 0
				underQuoted = false
				continuousBacktick = 0
				_, _ = writer.WriteString("`")
			}
			_ = writer.WriteByte(v)
			continue
		default:
			if shiftDelimiter-continuousBacktick <= 0 && !underQuoted {
				_, _ = writer.WriteString("`")
				underQuoted = true
				if selfQuoted = continuousBacktick > 0; selfQuoted {
					continuousBacktick -= 1
				}
			}

			for ; continuousBacktick > 0; continuousBacktick -= 1 {
				_, _ = writer.WriteString("``")
			}

			_ = writer.WriteByte(v)
		}
		shiftDelimiter++
	}

	if continuousBacktick > 0 && !selfQuoted {
		_, _ = writer.WriteString("``")
	}
	_, _ = writer.WriteString("`")
}

func (d Dialector) Explain(sql string, vars ...any) string {
	// single quotes: double quotes are identifier quoting in SQLite
	return logger.ExplainSQL(sql, nil, `'`, vars...)
}

func (d Dialector) DataTypeOf(field *schema.Field) string {
	// Computed (generated) column. SQLite 3.31+ supports STORED generated
	// columns; the expression is carried by the `generated` tag, separate from
	// the column type. https://www.sqlite.org/gencol.html
	if expr, ok := generatedColumnExpr(field); ok {
		return d.dataTypeOf(field) + " GENERATED ALWAYS AS (" + expr + ") STORED"
	}

	return d.dataTypeOf(field)
}

func (d Dialector) dataTypeOf(field *schema.Field) string {
	switch field.DataType {
	case schema.Bool:
		return "numeric"
	case schema.Int, schema.Uint:
		if field.AutoIncrement && !isCompositePrimaryKey(field) {
			// doesn't check `PrimaryKey`, to keep backward compatibility
			// https://www.sqlite.org/autoinc.html
			return "integer PRIMARY KEY AUTOINCREMENT"
		}

		return "integer"
	case schema.Float:
		return "real"
	case schema.String:
		return "text"
	case schema.Time:
		// Distinguish between schema.Time and tag time
		if val, ok := field.TagSettings["TYPE"]; ok {
			return val
		}

		return "datetime"
	case schema.Bytes:
		return "blob"
	}

	return string(field.DataType)
}

func (d Dialector) SavePoint(tx *gorm.DB, name string) error {
	tx.Exec("SAVEPOINT " + name)
	return nil
}

func (d Dialector) RollbackTo(tx *gorm.DB, name string) error {
	tx.Exec("ROLLBACK TO SAVEPOINT " + name)
	return nil
}

func (d Dialector) Translate(err error) error {
	var sqliteErr *moderncsqlite.Error
	if !errors.As(err, &sqliteErr) || sqliteErr == nil {
		return err
	}

	switch sqliteErr.Code() {
	case moderncsqlitelib.SQLITE_CONSTRAINT_UNIQUE,
		moderncsqlitelib.SQLITE_CONSTRAINT_PRIMARYKEY,
		moderncsqlitelib.SQLITE_CONSTRAINT_ROWID:
		return gorm.ErrDuplicatedKey
	case moderncsqlitelib.SQLITE_CONSTRAINT_FOREIGNKEY:
		return gorm.ErrForeignKeyViolated
	case moderncsqlitelib.SQLITE_CONSTRAINT_CHECK:
		return gorm.ErrCheckConstraintViolated
	}
	return err
}

func compareVersion(version1, version2 string) int {
	n, m := len(version1), len(version2)
	i, j := 0, 0
	for i < n || j < m {
		x := 0
		for ; i < n && version1[i] != '.'; i++ {
			x = x*10 + int(version1[i]-'0')
		}
		i++
		y := 0
		for ; j < m && version2[j] != '.'; j++ {
			y = y*10 + int(version2[j]-'0')
		}
		j++
		if x > y {
			return 1
		}
		if x < y {
			return -1
		}
	}
	return 0
}

// injectDSNParams injects defaults into the DSN to align modernc.org/sqlite's
// behavior with mattn/go-sqlite3:
//   - _texttotime=1 / _inttotime=1: return time.Time for datetime columns
//     (https://gitlab.com/cznic/sqlite/-/issues/245).
//   - _time_format=sqlite: format time.Time as "2006-01-02 15:04:05.999999999-07:00"
//     (https://github.com/libtnb/sqlite/issues/15).
//   - _busy_timeout=5000: default 5s busy timeout, just like
//     mattn/go-sqlite3 does, to avoid "database is locked" errors.
//
// Existing values in the DSN are preserved. The busy timeout is injected via
// the _busy_timeout shorthand rather than _pragma because modernc.org/sqlite
// applies _busy_timeout before any _pragma value, so whichever form the user
// sets it with (_busy_timeout, its _timeout alias or _pragma) overrides the
// default.
func injectDSNParams(dsn string) string {
	path, rawQuery := splitDSN(dsn)
	if path == "" {
		return dsn
	}
	values, err := url.ParseQuery(rawQuery)
	if err != nil {
		return dsn
	}

	var additions []string
	addIfMissing := func(key, value string) {
		if _, ok := values[key]; !ok {
			additions = append(additions, key+"="+value)
		}
	}
	addIfMissing("_texttotime", "1")
	addIfMissing("_inttotime", "1")
	addIfMissing("_time_format", "sqlite")
	if !hasBusyTimeout(values) {
		additions = append(additions, "_busy_timeout=5000")
	}

	if len(additions) == 0 {
		return dsn
	}
	if rawQuery == "" {
		return path + "?" + strings.Join(additions, "&")
	}
	return path + "?" + rawQuery + "&" + strings.Join(additions, "&")
}

// splitDSN splits a DSN into its path and raw-query parts at the first '?'.
func splitDSN(dsn string) (path, query string) {
	if i := strings.IndexRune(dsn, '?'); i >= 0 {
		return dsn[:i], dsn[i+1:]
	}
	return dsn, ""
}

// hasBusyTimeout reports whether the DSN query already sets the busy timeout,
// through the _busy_timeout/_timeout shorthand or a _pragma value.
func hasBusyTimeout(values url.Values) bool {
	if values.Has("_busy_timeout") || values.Has("_timeout") {
		return true
	}
	for _, p := range values["_pragma"] {
		if strings.Contains(strings.ToLower(p), "busy_timeout") {
			return true
		}
	}
	return false
}

// isCompositePrimaryKey reports whether field is part of a multi-column
// primary key. AUTOINCREMENT only works on a single-column INTEGER PRIMARY
// KEY; emitting it for a composite key would silently reduce the primary key
// to that one column (GORM skips the table-level PRIMARY KEY clause when a
// field type already contains PRIMARY KEY).
func isCompositePrimaryKey(field *schema.Field) bool {
	return field.PrimaryKey && field.Schema != nil && len(field.Schema.PrimaryFields) > 1
}

// generatedColumnExpr returns the expression of a computed (generated) column
// declared via the `generated` tag, if any. The `identity` keyword is reserved
// for identity columns (rendered through the dialect's native auto-increment)
// and is not a computed-column expression.
func generatedColumnExpr(field *schema.Field) (string, bool) {
	value, ok := field.TagSettings["GENERATED"]
	if !ok {
		return "", false
	}
	// Ignore an empty value or a bare `generated` tag, which the tag parser
	// stores as the upper-cased key, rather than treating it as an expression.
	if value = strings.TrimSpace(value); value == "" || value == "GENERATED" {
		return "", false
	}
	if isIdentityKeyword(value) {
		return "", false
	}
	return value, true
}

// isIdentityKeyword reports whether value is the `identity` keyword, optionally
// combined with the generation mode `always` / `by default`. Any other token
// means value is a computed-column expression.
func isIdentityKeyword(value string) bool {
	identity := false
	for _, token := range strings.Fields(strings.ToLower(value)) {
		switch token {
		case "identity":
			identity = true
		case "always", "by", "default":
		default:
			return false
		}
	}
	return identity
}