在Go中使用gocql高效处理Cassandra Set数据类型

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

INSERT INTO mykeyspace.category (category_id, product_list) VALUES ('electronics', {'laptop', 'smartphone', 'tablet'});

package main

import (
    "fmt"
    "github.com/gocql/gocql"
    "log"
)

func main() {
    // 配置Cassandra集群
    cluster := gocql.NewCluster("127.0.0.1") // 替换为你的Cassandra节点地址
    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 productIdList []string // 使用 []string 来接收 set<text> 类型的数据

    // 执行查询
    err = session.Query("SELECT product_list FROM category WHERE category_id=?", key).Scan(&productIdList)
    if err != nil {
        if err == gocql.ErrNotFound {
            fmt.Printf("未找到类别ID: %s\n", key)
        } else {
            log.Fatalf("查询失败: %v", err)
        }
    } else {
        fmt.Printf("从Cassandra获取的Product ID列表 (%s): %v\n", key, productIdList)
        fmt.Printf("数据类型: %T\n", productIdList)

        // 遍历Set中的元素
        fmt.Println("产品列表中的元素:")
        for _, productID := range productIdList {
            fmt.Printf(" - %s\n", productID)
        }
    }
}

type Unmarshaller interface {
    UnmarshalCQL(info TypeInfo, data []byte) error
}

type Marshaller interface {
    MarshalCQL(info TypeInfo) ([]byte, error)
}

package main

import (
    "fmt"
    "github.com/gocql/gocql"
    "log"
    "strings" // 用于模拟解析
)

// CustomSet 是一个自定义类型，用于表示Cassandra的Set，这里以map[string]bool为例
type CustomSet map[string]bool

// UnmarshalCQL 实现 gocql.Unmarshaller 接口
// 将Cassandra的Set数据（[]byte）解析到 CustomSet 类型中
func (s *CustomSet) UnmarshalCQL(info gocql.TypeInfo, data []byte) error {
    if data == nil {
        *s = make(CustomSet) // 处理NULL值，初始化为空map
        return nil
    }

    // 注意：实际的UnmarshalCQL实现需要解析gocql传递的Cassandra二进制数据格式。
    // 这是一个复杂的过程，通常需要深入了解Cassandra的内部序列化协议。
    // 以下代码仅为演示目的，模拟解析一个简单的文本表示的Set，实际生产环境不可用。
    // gocql不会直接传递 "{item1, item2}" 这样的字符串，而是二进制数据。

    strData := string(data) // 仅用于演示，实际应解析二进制
    log.Printf("UnmarshalCQL: 模拟接收到数据: %s\n", strData)

    // 假设我们能从某种格式（比如字符串"{item1, item2}"）中解析出元素
    // 实际情况需要更复杂的二进制解析逻辑
    if *s == nil {
        *s = make(CustomSet)
    }

    // 模拟解析逻辑：移除大括号，按逗号分割
    trimmedData := strings.Trim(strData, "{} ")
    if trimmedData == "" {
        return nil // 空集合
    }
    elements := strings.Split(trimmedData, ",")
    for _, elem := range elements {
        (*s)[strings.TrimSpace(elem)] = true
    }
    return nil
}

// MarshalCQL 实现 gocql.Marshaller 接口
// 将 CustomSet 类型的数据序列化为 []byte，以便写入Cassandra
func (s CustomSet) MarshalCQL(info gocql.TypeInfo) ([]byte, error) {
    if s == nil || len(s) == 0 {
        return nil, nil // 处理空Set
    }

    // 注意：实际的MarshalCQL实现需要将Go类型转换为Cassandra的二进制数据格式。
    // 这是一个复杂的过程，通常需要深入了解Cassandra的内部序列化协议。
    // 以下代码仅为演示目的，模拟生成一个简单的文本表示的Set，实际生产环境不可用。

    var elements []string
    for k := range s {
        elements = append(elements, k)
    }
    // 模拟生成CQL Set的文本表示，实际应生成二进制
    cqlSetString := "{" + strings.Join(elements, ", ") + "}"
    log.Printf("MarshalCQL: 模拟序列化数据为: %s\n", cqlSetString)
    return []byte(cqlSetString), nil
}

func main() {
    // 配置Cassandra集群 (与上一个例子相同)
    cluster := gocql.NewCluster("127.0.0.1")
    cluster.Keyspace = "mykeyspace"
    cluster.Consistency = gocql.Quorum

    session, err := cluster.CreateSession()
    if err != nil {
        log.Fatalf("创建Cassandra会话失败: %v", err)
    }
    defer session.Close()

    // --- 演示读取 (UnmarshalCQL) ---
    var customProductSet CustomSet
    key := "electronics" // 假设 'electronics' 存在

    // 实际应用中，gocql会调用 customProductSet.UnmarshalCQL
    // err = session.Query("SELECT product_list FROM category WHERE category_id=?", key).Scan(&customProductSet)
    // if err != nil {
    //  if err == gocql.ErrNotFound {
    //      fmt.Printf("未找到类别ID: %s\n", key)
    //  } else {
    //      log.Fatalf("查询失败: %v", err)
    //  }
    // } else {
    //  fmt.Printf("从Cassandra获取的自定义Product Set (%s): %v\n", key, customProductSet)
    //  fmt.Printf("数据类型: %T\n", customProductSet)
    // }

    // 模拟 UnmarshalCQL 被调用后的结果 (这里我们手动填充数据以模拟)
    // 注意：实际的UnmarshalCQL需要解析gocql提供的二进制数据
    // 这里为了演示，直接构造一个模拟的二进制数据（实际上是字符串）
    mockBinaryData := []byte("{laptop, smartphone, tablet}")
    err = (&customProductSet).UnmarshalCQL(gocql.TypeInfo{}, mockBinaryData) // 模拟调用
    if err != nil {
        log.Fatalf("UnmarshalCQL 模拟失败: %v", err)
    }
    fmt.Printf("从Cassandra获取的自定义Product Set (模拟读取): %v\n", customProductSet)
    fmt.Printf("数据类型: %T\n", customProductSet)
    fmt.Println("是否包含 'laptop'?", customProductSet["laptop"])

    // --- 演示写入 (MarshalCQL) ---
    newProductSet := CustomSet{
        "keyboard": true,
        "mouse":    true,
    }
    newKey := "peripherals"

    // 实际应用中，gocql会调用 newProductSet.MarshalCQL
    // err = session.Query("INSERT INTO category (category_id, product_list) VALUES (?, ?)", newKey, newProductSet).Exec()
    // if err != nil {
    //  log.Fatalf("插入数据失败: %v", err)
    // }
    // fmt.Printf("成功插入新的类别 '%s' 及其产品列表。\n", newKey)

    // 模拟 MarshalCQL 被调用后的结果
    marshaledData, err := newProductSet.MarshalCQL(gocql.TypeInfo{})
    if err != nil {
        log.Fatalf("MarshalCQL 模拟失败: %v", err)
    }
    fmt.Printf("自定义Product Set序列化后的数据 (模拟写入): %s\n", string(marshaledData))
}