orama/pkg/rqlite/adapter.go

package rqlite

import (
	"database/sql"
	"fmt"
	"time"

	_ "github.com/rqlite/gorqlite/stdlib" // Import the database/sql driver
)

// RQLiteAdapter adapts RQLite to the sql.DB interface
type RQLiteAdapter struct {
	manager *RQLiteManager
	db      *sql.DB
}

// NewRQLiteAdapter creates a new adapter that provides sql.DB interface for RQLite
func NewRQLiteAdapter(manager *RQLiteManager) (*RQLiteAdapter, error) {
	// Build DSN with optional basic auth credentials
	dsn := fmt.Sprintf("http://localhost:%d?disableClusterDiscovery=true&level=none", manager.config.RQLitePort)
	if manager.config.RQLiteUsername != "" && manager.config.RQLitePassword != "" {
		dsn = fmt.Sprintf("http://%s:%s@localhost:%d?disableClusterDiscovery=true&level=none",
			manager.config.RQLiteUsername, manager.config.RQLitePassword, manager.config.RQLitePort)
	}
	db, err := sql.Open("rqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to open RQLite SQL connection: %w", err)
	}

	// Configure connection pool with proper timeouts and limits
	// Optimized for concurrent operations and fast bad connection eviction
	db.SetMaxOpenConns(100)                 // Allow more concurrent connections to prevent queuing
	db.SetMaxIdleConns(10)                  // Keep fewer idle connections to force fresh reconnects
	db.SetConnMaxLifetime(30 * time.Second) // Short lifetime ensures bad connections die quickly
	db.SetConnMaxIdleTime(10 * time.Second) // Kill idle connections quickly to prevent stale state

	return &RQLiteAdapter{
		manager: manager,
		db:      db,
	}, nil
}

// GetSQLDB returns the sql.DB interface for compatibility with existing storage service
func (a *RQLiteAdapter) GetSQLDB() *sql.DB {
	return a.db
}

// GetManager returns the underlying RQLite manager for advanced operations
func (a *RQLiteAdapter) GetManager() *RQLiteManager {
	return a.manager
}

// Close closes the adapter connections
func (a *RQLiteAdapter) Close() error {
	if a.db != nil {
		a.db.Close()
	}
	return a.manager.Stop()
}