backend/pkg/library/cassandra/client.go

package cassandra

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/gocql/gocql"
	"github.com/scylladb/gocqlx/v3"
)

// cassandraConf 是初始化 CassandraDB 所需的內部設定（私有）
type cassandraConf struct {
	Hosts                    []string          // Cassandra 主機列表
	Port                     int               // 連線埠
	Keyspace                 string            // 預設使用的 Keyspace
	Username                 string            // 認證用戶名
	Password                 string            // 認證密碼
	Consistency              gocql.Consistency // 一致性級別
	ConnectTimeoutSec        int               // 連線逾時秒數
	NumConns                 int               // 每個節點連線數
	MaxRetries               int               // 重試次數
	UseAuth                  bool              // 是否使用帳號密碼驗證
	RetryMinInterval         time.Duration     // 重試間隔最小值
	RetryMaxInterval         time.Duration     // 重試間隔最大值
	ReconnectInitialInterval time.Duration     // 重連初始間隔
	ReconnectMaxInterval     time.Duration     // 重連最大間隔
	CQLVersion               string            // 執行連線的CQL 版本號
}

// CassandraDB 是封裝了 Cassandra 資料庫 session 的結構
type CassandraDB struct {
	session         gocqlx.Session
	SaiSupported    bool   // 是否支援 sai
	Version         string // 資料庫版本
	defaultKeyspace string // 預設 keyspace
}

// NewCassandraDB 初始化並建立 Cassandra 資料庫連線，使用預設設定並可透過Option修改
func NewCassandraDB(hosts []string, opts ...Option) (*CassandraDB, error) {
	config := &cassandraConf{
		Hosts:                    hosts,
		Port:                     defaultPort,
		Consistency:              defaultConsistency,
		ConnectTimeoutSec:        defaultTimeoutSec,
		NumConns:                 defaultNumConns,
		MaxRetries:               defaultMaxRetries,
		RetryMinInterval:         defaultRetryMinInterval,
		RetryMaxInterval:         defaultRetryMaxInterval,
		ReconnectInitialInterval: defaultReconnectInitialInterval,
		ReconnectMaxInterval:     defaultReconnectMaxInterval,
		CQLVersion:               defaultCqlVersion,
	}

	// 套用Option設定選項
	for _, opt := range opts {
		opt(config)
	}

	// 建立連線設定
	cluster := gocql.NewCluster(config.Hosts...)
	cluster.Port = config.Port
	cluster.Consistency = config.Consistency
	cluster.Timeout = time.Duration(config.ConnectTimeoutSec) * time.Second
	cluster.NumConns = config.NumConns
	cluster.RetryPolicy = &gocql.ExponentialBackoffRetryPolicy{
		NumRetries: config.MaxRetries,
		Min:        config.RetryMinInterval,
		Max:        config.RetryMaxInterval,
	}

	cluster.ReconnectionPolicy = &gocql.ExponentialReconnectionPolicy{
		MaxRetries:      config.MaxRetries,
		InitialInterval: config.ReconnectInitialInterval,
		MaxInterval:     config.ReconnectMaxInterval,
	}

	// 若有提供 Keyspace 則指定
	if config.Keyspace != "" {
		cluster.Keyspace = config.Keyspace
	}

	// 若啟用驗證則設定帳號密碼
	if config.UseAuth {
		cluster.Authenticator = gocql.PasswordAuthenticator{
			Username: config.Username,
			Password: config.Password,
		}
	}

	// 建立 Session
	s, err := gocqlx.WrapSession(cluster.CreateSession())
	if err != nil {
		return nil, fmt.Errorf("failed to connect to Cassandra cluster (hosts: %v, port: %d): %w", config.Hosts, config.Port, err)
	}

	db := &CassandraDB{
		session:         s,
		defaultKeyspace: config.Keyspace,
	}

	version, err := db.getReleaseVersion()
	if err != nil {
		return nil, fmt.Errorf("failed to get DB version: %w", err)
	}

	db.Version = version
	db.SaiSupported = isSAISupported(version)

	return db, nil
}

// Close 關閉 Cassandra 資料庫連線
func (db *CassandraDB) Close() {
	db.session.Close()
}

// GetSession 返回目前使用的 Cassandra Session
func (db *CassandraDB) GetSession() gocqlx.Session {
	return db.session
}

// GetDefaultKeyspace 返回預設的 keyspace
func (db *CassandraDB) GetDefaultKeyspace() string {
	return db.defaultKeyspace
}

// WithKeyspace 返回一個帶有指定 keyspace 的查詢構建器
// 如果 keyspace 為空，則使用預設 keyspace
func (db *CassandraDB) WithKeyspace(keyspace string) *KeyspaceDB {
	if keyspace == "" {
		keyspace = db.defaultKeyspace
	}
	return &KeyspaceDB{
		db:       db,
		keyspace: keyspace,
	}
}

// KeyspaceDB 是帶有 keyspace 的資料庫包裝器
type KeyspaceDB struct {
	db       *CassandraDB
	keyspace string
}

// GetSession 返回 session
func (kdb *KeyspaceDB) GetSession() gocqlx.Session {
	return kdb.db.GetSession()
}

// GetKeyspace 返回 keyspace
func (kdb *KeyspaceDB) GetKeyspace() string {
	return kdb.keyspace
}

// EnsureTable 確認並建立資料表
func (db *CassandraDB) EnsureTable(schema string) error {
	return db.session.ExecStmt(schema)
}

func (db *CassandraDB) InitVersionSupport() error {
	version, err := db.getReleaseVersion()
	if err != nil {
		return err
	}
	db.Version = version
	db.SaiSupported = isSAISupported(version)
	return nil
}

func (db *CassandraDB) getReleaseVersion() (string, error) {
	var version string
	stmt := "SELECT release_version FROM system.local"
	err := db.GetSession().Query(stmt, []string{"release_version"}).Consistency(gocql.One).Scan(&version)
	return version, err
}

func isSAISupported(version string) bool {
	// 只要 major >=5 就支援
	// 4.0.9+ 才有 SAI，但不穩，強烈建議 5.0+
	parts := strings.Split(version, ".")
	if len(parts) < 2 {
		return false
	}
	major, _ := strconv.Atoi(parts[0])
	minor, _ := strconv.Atoi(parts[1])

	if major >= 5 {
		return true
	}

	if major == 4 {
		if minor > 0 { // 4.1.x、4.2.x 直接支援
			return true
		}
		if minor == 0 {
			patch := 0
			if len(parts) >= 3 {
				patch, _ = strconv.Atoi(parts[2])
			}
			if patch >= 9 {
				return true
			}
		}
	}

	return false
}