package cassandra import ( "fmt" "log" "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 // 是否使用帳號密碼驗證 RetryMin time.Duration // 重試間隔最小值 RetryMax time.Duration // 重試間隔最大值 ReconnectInitial time.Duration // 重連初始間隔 ReconnectMax time.Duration // 重連最大間隔 CQLVersion string // 執行連線的CQL 版本號 } // CassandraDB 是封裝了 Cassandra 資料庫 session 的結構 type CassandraDB struct { session gocqlx.Session SaiSupported bool // 是否支援 sai Version string // 資料庫版本 } // 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, 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.NumConns 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) log.Printf("[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, } } log.Printf("[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) } log.Printf("[CassandraDB] success init Cassandra cluster") db := &CassandraDB{ 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 *CassandraDB) Close() { db.session.Close() } // GetSession 返回目前使用的 Cassandra Session func (db *CassandraDB) GetSession() gocqlx.Session { return db.session } // 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 == 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 }