baby-nas/internal/db/db.go

package db

import (
	"database/sql"
	"embed"
	"encoding/json"
	"fmt"
	"io/fs"
	"sort"
	"strings"

	"github.com/darroyo/nasctl/internal/system"
	_ "modernc.org/sqlite"
)

//go:embed migrations/*.sql
var migrationsFS embed.FS

type DB struct {
	conn *sql.DB
}

func Open(path string) (*DB, error) {
	conn, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	if _, err := conn.Exec(`PRAGMA foreign_keys = ON`); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("enable foreign keys: %w", err)
	}
	return &DB{conn: conn}, nil
}

func (d *DB) Close() error {
	return d.conn.Close()
}

func (d *DB) Conn() *sql.DB {
	return d.conn
}

func (d *DB) Migrate() error {
	if err := d.createMigrationsTable(); err != nil {
		return fmt.Errorf("create migrations table: %w", err)
	}

	entries, err := fs.ReadDir(migrationsFS, "migrations")
	if err != nil {
		return fmt.Errorf("read migrations: %w", err)
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Name() < entries[j].Name()
	})

	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sql") {
			continue
		}
		name := entry.Name()

		applied, err := d.isMigrationApplied(name)
		if err != nil {
			return fmt.Errorf("check migration %s: %w", name, err)
		}
		if applied {
			continue
		}

		content, err := migrationsFS.ReadFile("migrations/" + name)
		if err != nil {
			return fmt.Errorf("read migration %s: %w", name, err)
		}

		if err := d.runMigration(name, string(content)); err != nil {
			return fmt.Errorf("apply migration %s: %w", name, err)
		}

		if name == "0002_nfs_structured.sql" {
			if err := d.PopulateLegacyFSIDs(); err != nil {
				return fmt.Errorf("populate legacy fsids: %w", err)
			}
		}
		if name == "0005_nfs_per_host_options.sql" {
			if err := d.MigrateNFSClients(); err != nil {
				return fmt.Errorf("migrate nfs clients: %w", err)
			}
		}
		if name == "0006_nfs_fix_zero_fsid.sql" {
			if err := d.FixZeroFSIDs(); err != nil {
				return fmt.Errorf("fix zero fsids: %w", err)
			}
		}
	}
	return nil
}

func (d *DB) createMigrationsTable() error {
	_, err := d.conn.Exec(`
		CREATE TABLE IF NOT EXISTS schema_migrations (
			name TEXT PRIMARY KEY,
			applied_at TEXT NOT NULL DEFAULT (datetime('now'))
		)
	`)
	return err
}

func (d *DB) isMigrationApplied(name string) (bool, error) {
	var count int
	err := d.conn.QueryRow(
		"SELECT COUNT(*) FROM schema_migrations WHERE name = ?",
		name,
	).Scan(&count)
	if err != nil {
		return false, err
	}
	return count > 0, nil
}

func (d *DB) runMigration(name, content string) error {
	if _, err := d.conn.Exec(content); err != nil {
		if isIgnorableMigrationError(err) {
			// Schema is already in target state (e.g. columns already added
			// by a previous partial run). Continue to record the migration.
		} else {
			return fmt.Errorf("execute: %w", err)
		}
	}

	if _, err := d.conn.Exec(
		"INSERT INTO schema_migrations (name) VALUES (?)",
		name,
	); err != nil {
		return fmt.Errorf("record migration: %w", err)
	}
	return nil
}

func isIgnorableMigrationError(err error) bool {
	if err == nil {
		return false
	}
	msg := strings.ToLower(err.Error())
	return strings.Contains(msg, "duplicate column name") ||
		strings.Contains(msg, "already exists")
}

func (d *DB) MigrateNFSClients() error {
	rows, err := d.conn.Query(`
		SELECT id, clients, read_only, async_, root_squash, subtree_check, fsid, advanced
		FROM nfs_exports`)
	if err != nil {
		return fmt.Errorf("query nfs exports: %w", err)
	}
	defer rows.Close()

	type rowData struct {
		ID           int64
		Clients      string
		ReadOnly     bool
		Async        bool
		RootSquash   bool
		SubtreeCheck bool
		FSID         int64
		Advanced     string
	}

	var toUpdate []rowData
	for rows.Next() {
		var r rowData
		if err := rows.Scan(&r.ID, &r.Clients, &r.ReadOnly, &r.Async, &r.RootSquash, &r.SubtreeCheck, &r.FSID, &r.Advanced); err != nil {
			return fmt.Errorf("scan row: %w", err)
		}

		var raw []any
		if err := json.Unmarshal([]byte(r.Clients), &raw); err != nil {
			continue
		}
		if len(raw) == 0 {
			continue
		}

		if _, ok := raw[0].(string); !ok {
			continue
		}

		var adv NFSAdvanced
		if r.Advanced != "" && r.Advanced != "{}" {
			_ = json.Unmarshal([]byte(r.Advanced), &adv)
		}

		clients := make([]NFSClient, len(raw))
		for i, elem := range raw {
			host, ok := elem.(string)
			if !ok {
				host = ""
			}
			clients[i] = NFSClient{
				Host:         host,
				ReadOnly:     r.ReadOnly,
				Async:        r.Async,
				RootSquash:   r.RootSquash,
				SubtreeCheck: r.SubtreeCheck,
				Advanced:     adv,
			}
		}

		clientsJSON, err := json.Marshal(clients)
		if err != nil {
			return fmt.Errorf("marshal clients for id %d: %w", r.ID, err)
		}

		r.Clients = string(clientsJSON)
		toUpdate = append(toUpdate, r)
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("rows iteration: %w", err)
	}

	for _, r := range toUpdate {
		fsid := r.FSID
		if fsid == 0 {
			n, err := system.GenerateRandomFSID()
			if err != nil {
				return fmt.Errorf("generate fsid for id %d: %w", r.ID, err)
			}
			fsid = int64(n)
		}
		if _, err := d.conn.Exec(`UPDATE nfs_exports SET clients = ?, fsid = ? WHERE id = ?`, r.Clients, fsid, r.ID); err != nil {
			return fmt.Errorf("update nfs export id %d: %w", r.ID, err)
		}
	}
	return nil
}

func (d *DB) FixZeroFSIDs() error {
	rows, err := d.conn.Query(`SELECT id, fsid FROM nfs_exports WHERE fsid = 0`)
	if err != nil {
		return fmt.Errorf("query zero fsid exports: %w", err)
	}
	defer rows.Close()

	type fsidRow struct {
		ID   int64
		FSID int64
	}
	var toFix []fsidRow
	for rows.Next() {
		var r fsidRow
		if err := rows.Scan(&r.ID, &r.FSID); err != nil {
			return fmt.Errorf("scan fsid row: %w", err)
		}
		toFix = append(toFix, r)
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("rows iteration: %w", err)
	}

	for _, r := range toFix {
		n, err := system.GenerateRandomFSID()
		if err != nil {
			return fmt.Errorf("generate fsid for id %d: %w", r.ID, err)
		}
		if _, err := d.conn.Exec(`UPDATE nfs_exports SET fsid = ? WHERE id = ?`, int64(n), r.ID); err != nil {
			return fmt.Errorf("update fsid for id %d: %w", r.ID, err)
		}
	}
	return nil
}