at-container-registry/pkg/hold/db/libsql_open.go

package db

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"log/slog"
	"strings"

	_ "github.com/tursodatabase/go-libsql"
)

// DefaultBusyTimeoutMs is how long a connection waits for a lock before giving
// up with SQLITE_BUSY. Matches the appview and labeler databases.
const DefaultBusyTimeoutMs = 5000

// localDSN normalizes a filesystem path into a libsql DSN.
func localDSN(path string) string {
	if strings.HasPrefix(path, ":memory:") || strings.HasPrefix(path, "file:") {
		return path
	}
	return "file:" + path
}

// OpenLocalDB opens a local (non-replica) libsql database with WAL journal mode
// and PRAGMA busy_timeout applied to every pooled connection.
//
// Every subsystem that opens a hold database file directly must go through
// this. database/sql opens connections lazily and on demand, and SQLite's
// busy_timeout is per-connection, so running the PRAGMA once against the pool
// only configures whichever connection happened to serve it. Every other
// connection commits with busy_timeout=0 and fails immediately the moment
// another writer holds the lock.
//
// ":memory:" databases are returned unwrapped: libsql gives each connection its
// own private in-memory database, so there is no cross-connection lock to wait
// on and WAL is meaningless.
func OpenLocalDB(path string) (*sql.DB, error) {
	dsn := localDSN(path)

	if strings.HasPrefix(dsn, ":memory:") {
		return sql.Open("libsql", dsn)
	}

	base, err := openLibsqlLocalConnector(dsn)
	if err != nil {
		return nil, err
	}
	db := sql.OpenDB(&busyTimeoutConnector{base: base, timeoutMs: DefaultBusyTimeoutMs})

	if err := applyWAL(db, dsn); err != nil {
		_ = db.Close()
		return nil, err
	}
	return db, nil
}

// applyWAL puts the database file into WAL journal mode. Unlike busy_timeout,
// journal mode is a property of the file itself, so one call covers every
// connection and every other process that opens the same file.
func applyWAL(db *sql.DB, dsn string) error {
	var journalMode string
	// libsql treats PRAGMA assignments as queries that return a row.
	if err := db.QueryRow("PRAGMA journal_mode = WAL").Scan(&journalMode); err != nil {
		return fmt.Errorf("failed to set journal mode: %w", err)
	}
	if !strings.EqualFold(journalMode, "wal") {
		// Some filesystems (notably network mounts) cannot support WAL. Carry
		// on rather than refusing to boot: busy_timeout still applies, and the
		// rollback journal is correct, just less concurrent.
		slog.Warn("Database did not accept WAL journal mode", "dsn", dsn, "journal_mode", journalMode)
	}
	return nil
}

// openLibsqlLocalConnector returns a driver.Connector for a local libsql DSN.
// go-libsql exports NewEmbeddedReplicaConnector for replica mode but no public
// constructor for local files, so we obtain the driver via a probe sql.Open
// (which is lazy and opens no connection) and ask it for a Connector.
func openLibsqlLocalConnector(dsn string) (driver.Connector, error) {
	probe, err := sql.Open("libsql", dsn)
	if err != nil {
		return nil, fmt.Errorf("probe libsql driver: %w", err)
	}
	drv := probe.Driver()
	_ = probe.Close()

	dctx, ok := drv.(driver.DriverContext)
	if !ok {
		return nil, fmt.Errorf("libsql driver does not implement driver.DriverContext")
	}
	return dctx.OpenConnector(dsn)
}

// busyTimeoutConnector wraps a driver.Connector and runs PRAGMA busy_timeout on
// every newly opened connection. SQLite's busy_timeout is per-connection, so
// this is the only way to ensure every conn in the pool waits on lock
// contention instead of returning SQLITE_BUSY immediately.
type busyTimeoutConnector struct {
	base      driver.Connector
	timeoutMs int
}

func (c *busyTimeoutConnector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.base.Connect(ctx)
	if err != nil {
		return nil, err
	}

	// libsql treats PRAGMA assignments as queries that return a row, so we
	// must use QueryerContext rather than ExecerContext.
	queryer, ok := conn.(driver.QueryerContext)
	if !ok {
		_ = conn.Close()
		return nil, fmt.Errorf("libsql conn does not support QueryerContext")
	}

	rows, err := queryer.QueryContext(ctx, fmt.Sprintf("PRAGMA busy_timeout = %d", c.timeoutMs), nil)
	if err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("set busy_timeout on new conn: %w", err)
	}
	_ = rows.Close()

	return conn, nil
}

func (c *busyTimeoutConnector) Driver() driver.Driver {
	return c.base.Driver()
}