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 }