blockchain/internal/lib/cassandra/cassandra.go

package cassandra

import (
	"fmt"
	"github.com/gocql/gocql"
	"github.com/scylladb/gocqlx/v3"
	"github.com/zeromicro/go-zero/core/logx"
	"strconv"
	"strings"
	"time"
)

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

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

// NewDB 初始化並建立 Cassandra 資料庫連線，使用預設設定並可透過Option修改
func NewDB(hosts []string, opts ...Option) (*DB, error) {
	config := &conf{
		Hosts:             hosts,
		Port:              defaultPort,
		Consistency:       defaultConsistency,
		ConnectTimeoutSec: defaultTimeoutSec,
		NumConnect:        defaultNumConnections,
		MaxRetries:        defaultMaxRetries,
		RetryMin:          defaultRetryMin,
		RetryMax:          defaultRetryMax,
		ReconnectInitial:  defaultReconnectInitial,
		ReconnectMax:      defaultReconnectMax,
		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.NumConnect
	cluster.RetryPolicy = &gocql.ExponentialBackoffRetryPolicy{
		NumRetries: config.MaxRetries,
		Min:        config.RetryMin,
		Max:        config.RetryMax,
	}

	cluster.ReconnectionPolicy = &gocql.ExponentialReconnectionPolicy{
		MaxRetries:      config.MaxRetries,
		InitialInterval: config.ReconnectInitial,
		MaxInterval:     config.ReconnectMax,
	}
	var session *gocql.Session
	var err error
	for i := 1; i <= config.MaxRetries; i++ {
		if session != nil {
			session.Close()
		}
		session, err = cluster.CreateSession()
		if err == nil {
			break
		}

		waitInterval := cluster.ReconnectionPolicy.GetInterval(i)
		logx.Errorf("[CassandraDB] Retry attempt #%d, waiting %s...", i, waitInterval)
		time.Sleep(waitInterval)
	}

	if session == nil {
		panic("failed to connect ....")
	}

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

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

	logx.Infof("[CassandraDB] try to connect to Cassandra cluster %v, port: %d", config.Hosts, config.Port)
	// 建立 Session
	s, err := gocqlx.WrapSession(session, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to Cassandra cluster: %s", err)
	}
	logx.Infof("[CassandraDB] success init Cassandra cluster")

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

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

	return db, nil
}

// NewDBFromSession 用現成的 gocql.Session 封裝
func NewDBFromSession(session *gocql.Session) (*DB, error) {
	s, err := gocqlx.WrapSession(session, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to wrap gocql session: %w", err)
	}

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

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

	return db, nil
}

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

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

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

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

func (db *DB) 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 == 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
}