snapshot of https://github.com/golang-migrate/migrate.git (master, 2026-09-15; blobs > 131072B omitted)

database/rqlite/rqlite.go

@@ -0,0 +1,328 @@
+package rqlite
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	nurl "net/url"
+	"strconv"
+	"strings"
+	"sync/atomic"
+
+	"github.com/golang-migrate/migrate/v4"
+	"github.com/golang-migrate/migrate/v4/database"
+	"github.com/rqlite/gorqlite"
+)
+
+func init() {
+	database.Register("rqlite", &Rqlite{})
+}
+
+const (
+	// DefaultMigrationsTable defines the default rqlite migrations table
+	DefaultMigrationsTable = "schema_migrations"
+
+	// DefaultConnectInsecure defines the default setting for connect insecure
+	DefaultConnectInsecure = false
+)
+
+// ErrNilConfig is returned if no configuration was passed to WithInstance
+var ErrNilConfig = fmt.Errorf("no config")
+
+// ErrBadConfig is returned if configuration was invalid
+var ErrBadConfig = fmt.Errorf("bad parameter")
+
+// Config defines the driver configuration
+type Config struct {
+	// ConnectInsecure sets whether the connection uses TLS. Ineffectual when using WithInstance
+	ConnectInsecure bool
+	// MigrationsTable configures the migrations table name
+	MigrationsTable string
+}
+
+type Rqlite struct {
+	db       *gorqlite.Connection
+	isLocked atomic.Bool
+
+	config *Config
+}
+
+// WithInstance creates a rqlite database driver with an existing gorqlite database connection
+// and a Config struct
+func WithInstance(instance *gorqlite.Connection, config *Config) (database.Driver, error) {
+	if config == nil {
+		return nil, ErrNilConfig
+	}
+
+	// we use the consistency level check as a database ping
+	if _, err := instance.ConsistencyLevel(); err != nil {
+		return nil, err
+	}
+
+	if len(config.MigrationsTable) == 0 {
+		config.MigrationsTable = DefaultMigrationsTable
+	}
+
+	driver := &Rqlite{
+		db:     instance,
+		config: config,
+	}
+
+	if err := driver.ensureVersionTable(); err != nil {
+		return nil, err
+	}
+
+	return driver, nil
+}
+
+// OpenURL creates a rqlite database driver from a connect URL
+func OpenURL(url string) (database.Driver, error) {
+	d := &Rqlite{}
+	return d.Open(url)
+}
+
+func (r *Rqlite) ensureVersionTable() (err error) {
+	if err = r.Lock(); err != nil {
+		return err
+	}
+
+	defer func() {
+		if e := r.Unlock(); e != nil {
+			err = errors.Join(err, e)
+		}
+	}()
+
+	stmts := []string{
+		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (version uint64, dirty bool)`, r.config.MigrationsTable),
+		fmt.Sprintf(`CREATE UNIQUE INDEX IF NOT EXISTS version_unique ON %s (version)`, r.config.MigrationsTable),
+	}
+
+	if _, err := r.db.Write(stmts); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Open returns a new driver instance configured with parameters
+// coming from the URL string. Migrate will call this function
+// only once per instance.
+func (r *Rqlite) Open(url string) (database.Driver, error) {
+	dburl, config, err := parseUrl(url)
+	if err != nil {
+		return nil, err
+	}
+	r.config = config
+
+	r.db, err = gorqlite.Open(dburl.String())
+	if err != nil {
+		return nil, err
+	}
+
+	if err := r.ensureVersionTable(); err != nil {
+		return nil, err
+	}
+
+	return r, nil
+}
+
+// Close closes the underlying database instance managed by the driver.
+// Migrate will call this function only once per instance.
+func (r *Rqlite) Close() error {
+	r.db.Close()
+	return nil
+}
+
+// Lock should acquire a database lock so that only one migration process
+// can run at a time. Migrate will call this function before Run is called.
+// If the implementation can't provide this functionality, return nil.
+// Return database.ErrLocked if database is already locked.
+func (r *Rqlite) Lock() error {
+	if !r.isLocked.CompareAndSwap(false, true) {
+		return database.ErrLocked
+	}
+	return nil
+}
+
+// Unlock should release the lock. Migrate will call this function after
+// all migrations have been run.
+func (r *Rqlite) Unlock() error {
+	if !r.isLocked.CompareAndSwap(true, false) {
+		return database.ErrNotLocked
+	}
+	return nil
+}
+
+// Run applies a migration to the database. migration is guaranteed to be not nil.
+func (r *Rqlite) Run(migration io.Reader) error {
+	migr, err := io.ReadAll(migration)
+	if err != nil {
+		return err
+	}
+
+	query := string(migr[:])
+	if _, err := r.db.WriteOne(query); err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+
+	return nil
+}
+
+// SetVersion saves version and dirty state.
+// Migrate will call this function before and after each call to Run.
+// version must be >= -1. -1 means NilVersion.
+func (r *Rqlite) SetVersion(version int, dirty bool) error {
+	deleteQuery := fmt.Sprintf(`DELETE FROM %s`, r.config.MigrationsTable)
+	statements := []gorqlite.ParameterizedStatement{
+		{
+			Query: deleteQuery,
+		},
+	}
+
+	// Also re-write the schema version for nil dirty versions to prevent
+	// empty schema version for failed down migration on the first migration
+	// See: https://github.com/golang-migrate/migrate/issues/330
+	insertQuery := fmt.Sprintf(`INSERT INTO %s (version, dirty) VALUES (?, ?)`, r.config.MigrationsTable)
+	if version >= 0 || (version == database.NilVersion && dirty) {
+		statements = append(statements, gorqlite.ParameterizedStatement{
+			Query: insertQuery,
+			Arguments: []interface{}{
+				version,
+				dirty,
+			},
+		})
+	}
+
+	wr, err := r.db.WriteParameterized(statements)
+	if err != nil {
+		for i, res := range wr {
+			if res.Err != nil {
+				return &database.Error{OrigErr: err, Query: []byte(statements[i].Query)}
+			}
+		}
+
+		// if somehow we're still here, return the original error with combined queries
+		return &database.Error{OrigErr: err, Query: []byte(deleteQuery + "\n" + insertQuery)}
+	}
+
+	return nil
+}
+
+// Version returns the currently active version and if the database is dirty.
+// When no migration has been applied, it must return version -1.
+// Dirty means, a previous migration failed and user interaction is required.
+func (r *Rqlite) Version() (version int, dirty bool, err error) {
+	query := "SELECT version, dirty FROM " + r.config.MigrationsTable + " LIMIT 1"
+
+	qr, err := r.db.QueryOne(query)
+	if err != nil {
+		return database.NilVersion, false, nil
+	}
+
+	if !qr.Next() {
+		return database.NilVersion, false, nil
+	}
+
+	if err := qr.Scan(&version, &dirty); err != nil {
+		return database.NilVersion, false, &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+
+	return version, dirty, nil
+}
+
+// Drop deletes everything in the database.
+// Note that this is a breaking action, a new call to Open() is necessary to
+// ensure subsequent calls work as expected.
+func (r *Rqlite) Drop() error {
+	query := `SELECT name FROM sqlite_master WHERE type = 'table'`
+
+	tables, err := r.db.QueryOne(query)
+	if err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+
+	statements := make([]string, 0)
+	for tables.Next() {
+		var tableName string
+		if err := tables.Scan(&tableName); err != nil {
+			return err
+		}
+
+		if len(tableName) > 0 {
+			statement := fmt.Sprintf(`DROP TABLE %s`, tableName)
+			statements = append(statements, statement)
+		}
+	}
+
+	// return if nothing to do
+	if len(statements) <= 0 {
+		return nil
+	}
+
+	wr, err := r.db.Write(statements)
+	if err != nil {
+		for i, res := range wr {
+			if res.Err != nil {
+				return &database.Error{OrigErr: err, Query: []byte(statements[i])}
+			}
+		}
+
+		// if somehow we're still here, return the original error with combined queries
+		return &database.Error{OrigErr: err, Query: []byte(strings.Join(statements, "\n"))}
+	}
+
+	return nil
+}
+
+func parseUrl(url string) (*nurl.URL, *Config, error) {
+	parsedUrl, err := nurl.Parse(url)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	config, err := parseConfigFromQuery(parsedUrl.Query())
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if parsedUrl.Scheme != "rqlite" {
+		return nil, nil, fmt.Errorf("bad scheme: %w", ErrBadConfig)
+	}
+
+	// adapt from rqlite to http/https schemes
+	if config.ConnectInsecure {
+		parsedUrl.Scheme = "http"
+	} else {
+		parsedUrl.Scheme = "https"
+	}
+
+	filteredUrl := migrate.FilterCustomQuery(parsedUrl)
+
+	return filteredUrl, config, nil
+}
+
+func parseConfigFromQuery(queryVals nurl.Values) (*Config, error) {
+	c := Config{
+		ConnectInsecure: DefaultConnectInsecure,
+		MigrationsTable: DefaultMigrationsTable,
+	}
+
+	migrationsTable := queryVals.Get("x-migrations-table")
+	if migrationsTable != "" {
+		if strings.HasPrefix(migrationsTable, "sqlite_") {
+			return nil, fmt.Errorf("invalid value for x-migrations-table: %w", ErrBadConfig)
+		}
+		c.MigrationsTable = migrationsTable
+	}
+
+	connectInsecureStr := queryVals.Get("x-connect-insecure")
+	if connectInsecureStr != "" {
+		connectInsecure, err := strconv.ParseBool(connectInsecureStr)
+		if err != nil {
+			return nil, fmt.Errorf("invalid value for x-connect-insecure: %w", ErrBadConfig)
+		}
+		c.ConnectInsecure = connectInsecure
+	}
+
+	return &c, nil
+}