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 }