store.go

/*
Copyright 2026 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package ipam

import (
	"database/sql"
	_ "embed"
	"fmt"
	"os"
	"path/filepath"

	"github.com/go-logr/logr"
	_ "github.com/mattn/go-sqlite3" // SQLite driver
)

//go:embed schema.sql
var schemaSQL string

const (
	// dbSchemaVersion tracks the SQLite schema version to allow safe local
	// migrations and prevent state corruption across daemon restarts.
	dbSchemaVersion = 1
	maxOpenConns    = 10
	maxIdleConns    = 10
)

// Store manages database operations for IPAM.
type Store struct {
	db  *sql.DB
	log logr.Logger
}

// NewStore creates a new Store instance and initializes the database.
func NewStore(log logr.Logger, dbPath string) (*Store, error) {
	if dbPath == "" {
		return nil, fmt.Errorf("dbPath cannot be empty: an absolute path must be explicitly provided")
	}

	log.Info("Opening or creating database", "path", dbPath)

	if err := os.MkdirAll(filepath.Dir(dbPath), 0755); err != nil {
		return nil, fmt.Errorf("failed to create db directory: %w", err)
	}

	// SQLite is configured directly through the DSN string. This approach
	// guarantees every new connection spawned by the sql.DB pool inherits these
	// exact configurations natively.
	dsn := dbPath +
		// Enables Write-Ahead Logging (WAL) mode. This significantly improves
		// concurrency by allowing multiple readers to access the database
		// simultaneously without blocking a writer, which is critical for burst
		// requests.
		// See: https://www.sqlite.org/pragma.html#pragma_journal_mode
		"?_journal_mode=WAL" +
		// Enforces foreign key constraints. SQLite ignores these by default.
		// This is required to ensure ON DELETE CASCADE functions correctly on the
		// ip_addresses table when a draining CIDR block is officially removed.
		// See: https://www.sqlite.org/pragma.html#pragma_foreign_keys
		"&_foreign_keys=on" +
		// Sets the busy timeout to 5000 milliseconds. If the database is locked
		// by another transaction, this tells the SQLite driver to wait for up
		// to 5 seconds before giving up and returning a locked error.
		// See: https://www.sqlite.org/pragma.html#pragma_busy_timeout
		"&_busy_timeout=5000" +
		// Instructs the Go driver to send "BEGIN IMMEDIATE" instead of standard
		// "BEGIN" when starting a transaction. This grabs a write lock instantly,
		// preventing deadlocks when concurrent requests try to upgrade their
		// read locks to write locks simultaneously. Note: This is a go-sqlite3
		// driver feature, not a native SQLite PRAGMA.
		// See: https://github.com/mattn/go-sqlite3#connection-string
		"&_txlock=immediate" +
		// Maps to PRAGMA synchronous = NORMAL. In WAL mode, this is the optimal
		// setting for high-concurrency daemons. It prevents database corruption
		// during power loss or hard crashes while offering much faster write
		// performance than FULL mode, sacrificing only a few milliseconds of
		// un-checkpointed durability.
		// See: https://www.sqlite.org/pragma.html#pragma_synchronous
		"&_synchronous=1"

	db, err := sql.Open("sqlite3", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	db.SetMaxOpenConns(maxOpenConns)
	db.SetMaxIdleConns(maxIdleConns)
	// Sets the maximum amount of time a connection may be reused to infinity
	// (0). This guarantees the single connection never expires.
	db.SetConnMaxLifetime(0)

	store := &Store{
		db:  db,
		log: log,
	}

	// Only a single process enters this execution block at a time.
	if err := store.initSchema(); err != nil {
		db.Close()
		return nil, fmt.Errorf("failed to initialize schema: %w", err)
	}

	log.Info("Initialized or updated database schema", "path", dbPath)

	return store, nil
}

// initSchema creates the necessary tables if they don't exist.
func (s *Store) initSchema() error {
	var currentVersion int
	err := s.db.QueryRow("PRAGMA user_version").Scan(&currentVersion)
	if err != nil {
		return fmt.Errorf("failed to check schema version: %w", err)
	}

	if currentVersion == dbSchemaVersion {
		s.log.V(4).Info("Database schema already initialized", "version", currentVersion)
		return nil
	}

	s.log.Info("Initializing DB schema", "currentVersion", currentVersion, "expectedVersion", dbSchemaVersion)

	// 1. Begin an atomic transaction
	tx, err := s.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	// Safe to defer; Rollback does nothing if Commit() is successful
	defer tx.Rollback()

	// 2. Execute the embedded schema.sql file
	if _, err := tx.Exec(schemaSQL); err != nil {
		return fmt.Errorf("failed to execute schema.sql: %w", err)
	}

	// 3. Set User Version
	setVersion := fmt.Sprintf("PRAGMA user_version = %d;", dbSchemaVersion)
	if _, err := tx.Exec(setVersion); err != nil {
		return fmt.Errorf("failed to set user_version: %w", err)
	}

	// 4. Commit everything atomically
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit schema transaction: %w", err)
	}

	s.log.Info("Database schema initialized or updated successfully")
	return nil
}

// Close safely closes the database connection and releases any file locks.
// This should be called during the daemon's graceful shutdown sequence.
func (s *Store) Close() error {
	s.log.Info("Closing IPAM database connection")

	if err := s.db.Close(); err != nil {
		return fmt.Errorf("failed to close database connection: %w", err)
	}

	return nil
}