Go语言中处理Cassandra Set数据类型：使用gocql的实践指南

CREATE TABLE category (
    category_id text PRIMARY KEY,
    product_list set<text>
);

INSERT INTO category (category_id, product_list) VALUES ('electronics', {'laptop', 'mouse', 'keyboard'});

package main

import (
    "fmt"
    "log"

    "github.com/gocql/gocql"
)

func main() {
    // 连接Cassandra集群
    cluster := gocql.NewCluster("127.0.0.1") // 替换为你的Cassandra节点IP
    cluster.Keyspace = "mykeyspace"          // 替换为你的Keyspace
    cluster.Consistency = gocql.Quorum
    session, err := cluster.CreateSession()
    if err != nil {
        log.Fatalf("无法连接Cassandra: %v", err)
    }
    defer session.Close()

    key := "electronics"
    var productIdList []string // product_list是set<text>，所以Go中对应[]string

    // 执行查询并扫描结果
    err = session.Query("SELECT product_list FROM category WHERE category_id=?", key).Scan(&productIdList)
    if err != nil {
        log.Fatalf("查询失败: %v", err)
    }

    fmt.Printf("产品列表: %v\n", productIdList)
    // 预期输出: 产品列表: [laptop mouse keyboard] (顺序可能不同，因为SET是无序的)
}

// Marshaller is an interface that can be implemented by custom types to allow them
// to be marshalled into CQL values.
type Marshaller interface {
    MarshalCQL(info TypeInfo) ([]byte, error)
}

// Unmarshaller is an interface that can be implemented by custom types to allow them
// to be unmarshalled from CQL values.
type Unmarshaller interface {
    UnmarshalCQL(info TypeInfo, data []byte) error
}

package main

import (
    "bytes"
    "encoding/binary"
    "fmt"
    "log"

    "github.com/gocql/gocql"
)

// CustomSet 定义一个Go类型来表示Cassandra的SET
type CustomSet map[string]bool

// UnmarshalCQL 实现 gocql.Unmarshaller 接口，将CQL的SET数据解析到CustomSet
func (s *CustomSet) UnmarshalCQL(info gocql.TypeInfo, data []byte) error {
    if data == nil {
        *s = make(CustomSet)
        return nil
    }

    if info.Type() != gocql.TypeSet {
        return fmt.Errorf("UnmarshalCQL: expected SET type, got %s", info.Type())
    }

    if *s == nil {
        *s = make(CustomSet)
    }

    buf := bytes.NewBuffer(data)
    // 读取SET的元素数量 (4字节整数)
    var count int32
    if err := binary.Read(buf, binary.BigEndian, &count); err != nil {
        return fmt.Errorf("UnmarshalCQL: failed to read set element count: %w", err)
    }

    // 遍历并读取每个元素
    for i := 0; i < int(count); i++ {
        // 读取每个元素的长度 (4字节整数)
        var elemLen int32
        if err := binary.Read(buf, binary.BigEndian, &elemLen); err != nil {
            return fmt.Errorf("UnmarshalCQL: failed to read element length: %w", err)
        }
        // 读取元素数据
        elemBytes := make([]byte, elemLen)
        if _, err := buf.Read(elemBytes); err != nil {
            return fmt.Errorf("UnmarshalCQL: failed to read element data: %w", err)
        }
        // 假设SET的元素是text，直接转换为string
        (*s)[string(elemBytes)] = true
    }
    return nil
}

// MarshalCQL 实现 gocql.Marshaller 接口，将CustomSet数据序列化为CQL的SET数据
func (s CustomSet) MarshalCQL(info gocql.TypeInfo) ([]byte, error) {
    if s == nil {
        return nil, nil
    }

    if info.Type() != gocql.TypeSet {
        return nil, fmt.Errorf("MarshalCQL: expected SET type, got %s", info.Type())
    }

    buf := new(bytes.Buffer)
    // 写入SET的元素数量
    if err := binary.Write(buf, binary.BigEndian, int32(len(s))); err != nil {
        return nil, fmt.Errorf("MarshalCQL: failed to write set element count: %w", err)
    }

    // 遍历并写入每个元素
    // 注意：Cassandra SET是无序的，这里遍历map的顺序不影响最终结果
    for elem := range s {
        elemBytes := []byte(elem)
        // 写入每个元素的长度
        if err := binary.Write(buf, binary.BigEndian, int32(len(elemBytes))); err != nil {
            return nil, fmt.Errorf("MarshalCQL: failed to write element length: %w", err)
        }
        // 写入元素数据
        if _, err := buf.Write(elemBytes); err != nil {
            return nil, fmt.Errorf("MarshalCQL: failed to write element data: %w", err)
        }
    }
    return buf.Bytes(), nil
}

func main() {
    cluster := gocql.NewCluster("127.00.1")
    cluster.Keyspace = "mykeyspace"
    cluster.Consistency = gocql.Quorum
    session, err := cluster.CreateSession()
    if err != nil {
        log.Fatalf("无法连接Cassandra: %v", err)
    }
    defer session.Close()

    key := "electronics"
    var customProductList CustomSet // 使用自定义类型

    // 查询数据
    err = session.Query("SELECT product_list FROM category WHERE category_id=?", key).Scan(&customProductList)
    if err != nil {
        log.Fatalf("查询失败: %v", err)
    }
    fmt.Printf("自定义产品列表 (查询): %v\n", customProductList)
    // 预期输出: 自定义产品列表 (查询): map[keyboard:true laptop:true mouse:true] (顺序可能不同)

    // 写入数据示例
    newProductList := CustomSet{
        "monitor": true,
        "webcam":  true,
    }
    err = session.Query("UPDATE category SET product_list = ? WHERE category_id = ?", newProductList, "electronics").Exec()
    if err != nil {
        log.Fatalf("更新失败: %v", err)
    }
    fmt.Println("产品列表更新成功。")

    // 再次查询以验证更新
    var updatedProductList CustomSet
    err = session.Query("SELECT product_list FROM category WHERE category_id=?", key).Scan(&updatedProductList)
    if err != nil {
        log.Fatalf("再次查询失败: %v", err)
    }
    fmt.Printf("更新后的自定义产品列表: %v\n", updatedProductList)
    // 预期输出: 更新后的自定义产品列表: map[keyboard:true laptop:true monitor:true mouse:true webcam:true]
}