feat(cassandra): add native Go agent

This commit is contained in:
t8y2 2026-08-04 00:39:05 +08:00
parent a6f1a9a871
commit 5c2dfde334
No known key found for this signature in database
27 changed files with 2654 additions and 29 deletions

View File

@ -45,13 +45,14 @@ function fileContainsCommonDependency(path, moduleExists, readModuleFile) {
}
const nativeDriverDirectories = {
cassandra: "cassandra-go",
duckdb: "duckdb",
oracle: "oracle-go",
kingbase: "kingbase-go",
vastbase: "vastbase-go",
rabbitmq: "rabbitmq",
};
const nativeDriverModules = new Set(["duckdb", "oracle", "xugu", "kingbase", "vastbase", "rabbitmq"]);
const nativeDriverModules = new Set(["cassandra", "duckdb", "oracle", "xugu", "kingbase", "vastbase", "rabbitmq"]);
function resolveAgentModule(moduleName, { legacyStandaloneModules, moduleExists, readModuleFile }) {
let checkDir = null;

View File

@ -57,6 +57,18 @@ test("bumps the native Vastbase agent from its independent Go directory", () =>
assert.deepEqual(result.nativeModules, ["vastbase"]);
});
test("bumps Cassandra from its native Go source directory", () => {
const result = evaluateAgentVersionBump({
versions: { cassandra: "0.1.37" },
changedFiles: ["agents/drivers/cassandra-go/main.go"],
moduleExists: (path) => path === "agents/drivers/cassandra-go",
readModuleFile: () => "",
});
assert.equal(result.versions.cassandra, "0.1.38");
assert.deepEqual(result.nativeModules, ["cassandra"]);
});
test("builds a manually versioned module even without runtime file changes", () => {
const result = evaluateAgentVersionBump({
versions: { duckdb: "0.1.1" },

View File

@ -47,6 +47,7 @@ const LABEL_PALETTE = [
];
const DRIVER_DATABASE_ALIASES = {
"cassandra-go": "cassandra",
gbase8a: "gbase",
gbase8s: "gbase",
"h2-legacy": "h2",

View File

@ -14,6 +14,7 @@ import {
const knownDatabaseTypes = new Set([
"access",
"cassandra",
"doris",
"jdbc",
"mongodb",
@ -55,12 +56,13 @@ test("maps agent and dialect paths to existing database types", () => {
assert.deepEqual(
inferDatabaseTypes([
"agents/drivers/oracle-go/go.mod",
"agents/drivers/cassandra-go/go.mod",
"agents/drivers/vastbase-go/go.mod",
"agents/drivers/kafka/build.gradle",
"plugins/dialects/postgresql.yaml",
"plugins/dialects/oceanbase.yaml",
], knownDatabaseTypes),
["mq", "oceanbase-oracle", "oracle", "postgres", "vastbase"],
["cassandra", "mq", "oceanbase-oracle", "oracle", "postgres", "vastbase"],
);
});

View File

@ -276,6 +276,46 @@ jobs:
name: rabbitmq-native
path: "release-native/dbx-agent-rabbitmq-*"
build-cassandra-native:
needs: [bump-versions]
if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'cassandra') }}
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-go@v5
with:
go-version: "1.22.x"
- name: Test Cassandra native agent
working-directory: agents/drivers/cassandra-go
run: go test ./...
- name: Cross-compile Cassandra native agent
shell: bash
run: |
mkdir -p release-native
cd agents/drivers/cassandra-go
declare -A TARGETS=(
["macos-aarch64"]="darwin/arm64"
["macos-x64"]="darwin/amd64"
["linux-aarch64"]="linux/arm64"
["linux-x64"]="linux/amd64"
["windows-aarch64"]="windows/arm64"
["windows-x64"]="windows/amd64"
)
for platform in "${!TARGETS[@]}"; do
IFS=/ read -r goos goarch <<< "${TARGETS[$platform]}"
output="../../../release-native/dbx-agent-cassandra-${platform}"
if [[ "$goos" == "windows" ]]; then
output="${output}.exe"
fi
echo "Building $platform ($goos/$goarch)"
CGO_ENABLED=0 GOOS="$goos" GOARCH="$goarch" go build -trimpath -ldflags="-s -w" -o "$output" .
done
ls -lh ../../../release-native
- uses: actions/upload-artifact@v4
with:
name: cassandra-native
path: "release-native/dbx-agent-cassandra-*"
build-kingbase-native:
needs: [bump-versions]
if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'kingbase') }}
@ -571,7 +611,7 @@ jobs:
retention-days: 1
release:
needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-kingbase-native, build-vastbase-native, build-duckdb-native, build-jre, reuse-previous-assets]
needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-cassandra-native, build-kingbase-native, build-vastbase-native, build-duckdb-native, build-jre, reuse-previous-assets]
if: ${{ always() && !contains(needs.*.result, 'failure') && !contains(needs.*.result, 'cancelled') }}
runs-on: ubuntu-latest
steps:
@ -693,6 +733,7 @@ jobs:
duckdb) echo "DuckDB" ;;
xugu) echo "虚谷 XuguDB" ;;
rabbitmq) echo "RabbitMQ" ;;
cassandra) echo "Apache Cassandra" ;;
*) echo "$name" ;;
esac
}
@ -756,7 +797,7 @@ jobs:
[ -n "$DRIVERS" ] && DRIVERS="${DRIVERS},"$'\n'
DRIVERS="${DRIVERS}$(generate_jar_entry "$name" "$label" "$f" "$jre_key" "$version" "$external_driver" "$native_json")"
done
for name in oracle xugu kingbase vastbase duckdb rabbitmq; do
for name in oracle xugu kingbase vastbase duckdb rabbitmq cassandra; do
version=$(get_module_version "$name")
[ -f "release/dbx-agent-${name}-${version}.jar" ] && continue
native_json=$(generate_native_platforms "$name" "$version")
@ -826,6 +867,7 @@ jobs:
oracle) echo "Oracle" ;;
xugu) echo "虚谷 XuguDB" ;;
rabbitmq) echo "RabbitMQ" ;;
cassandra) echo "Apache Cassandra" ;;
*) echo "$name" ;;
esac
}
@ -850,6 +892,8 @@ jobs:
LOG_PATH="agents/drivers/kingbase-go/"
elif [ "$name" = "vastbase" ]; then
LOG_PATH="agents/drivers/vastbase-go/"
elif [ "$name" = "cassandra" ]; then
LOG_PATH="agents/drivers/cassandra-go/"
elif [ -d "agents/drivers/$name" ]; then
LOG_PATH="agents/drivers/$name/"
else

View File

@ -538,6 +538,10 @@ jobs:
run: go test ./...
working-directory: agents/drivers/rabbitmq
- name: Cassandra native agent tests
run: go test ./...
working-directory: agents/drivers/cassandra-go
- name: Vastbase native agent tests
run: go test ./...
working-directory: agents/drivers/vastbase-go
@ -554,6 +558,10 @@ jobs:
run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-rabbitmq-linux-x64 .
working-directory: agents/drivers/rabbitmq
- name: Cassandra native agent build
run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-cassandra-linux-x64 .
working-directory: agents/drivers/cassandra-go
- name: Vastbase native agent build
run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-vastbase-linux-x64 .
working-directory: agents/drivers/vastbase-go

View File

@ -34,7 +34,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo
| db2 | IBM DB2 | DB2 JDBC |
| informix | IBM Informix | Informix JDBC |
| neo4j | Neo4j | Neo4j JDBC |
| cassandra | Apache Cassandra | Cassandra JDBC |
| cassandra | Apache Cassandra 2.1+ | Apache cassandra-gocql-driver native agent |
| bigquery | Google BigQuery | BigQuery JDBC |
| kylin | Apache Kylin | Kylin JDBC |
| sundb | SunDB | SunDB JDBC |
@ -49,7 +49,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo
## Multi-JRE Support
Most Java agents target JRE 21. Native agents, such as `duckdb`, `oracle`, `kingbase`, `xugu`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents.
Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `oracle`, `kingbase`, `xugu`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents.
## JDBC Connection Pooling
@ -75,7 +75,7 @@ Set `DBX_AGENT_JDBC_POOL_ENABLED=false` for a runtime-level compatibility fallba
For new agents, prefer a **native (Go or Rust) driver** over a Java/JDBC agent whenever a mature, license-compatible native driver is available. Native agents ship as a single self-contained executable with no JRE, which significantly reduces memory footprint and startup time — the JVM baseline that every Java agent pays even when idle is avoided entirely.
- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/duckdb`, `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/xugu`, and `drivers/rabbitmq` (amqp091-go) as reference implementations. No JRE download or management is needed.
- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/cassandra-go` (Apache cassandra-gocql-driver), `drivers/duckdb`, `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/xugu`, and `drivers/rabbitmq` (amqp091-go) as reference implementations. No JRE download or management is needed.
- **Java/JDBC** — the default fallback when only a JDBC driver exists for the database, or when the native driver is immature or unmaintained. Most agents still fall in this category.
Native agents implement the same JSON-RPC contract and `versions.json` registration as Java agents; they ship an `agent` executable instead of `agent.jar`. If both native and Java source implementations exist for the same database, publish only the native artifact unless the Java variant has a separately registered compatibility profile, such as `oracle-legacy` / `oracle-10g`.
@ -87,13 +87,14 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed).
```bash
./gradlew shadowJar
(cd drivers/oracle-go && go build -o agent .)
(cd drivers/cassandra-go && go build -o agent .)
(cd drivers/kingbase-go && go build -o agent .)
(cd drivers/vastbase-go && go build -o agent .)
(cd drivers/xugu && go build -o agent .)
(cd drivers/rabbitmq && go build -o agent .)
```
Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/xugu`, and `drivers/rabbitmq`.
Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/xugu`, and `drivers/rabbitmq`.
### Local DBX Runtime Test
@ -107,7 +108,7 @@ cp agents/drivers/<db_type>/build/libs/*-all.jar ~/.dbx/agents/drivers/<db_type>
Restart DBX or disconnect and reconnect the database so the new agent process loads the replacement JAR.
Native agents such as `oracle`, `kingbase`, `xugu`, and `rabbitmq` use the `agent` executable in the driver directory instead of `agent.jar`.
Native agents such as `cassandra`, `oracle`, `kingbase`, `xugu`, and `rabbitmq` use the `agent` executable in the driver directory instead of `agent.jar`.
## Versioning

View File

@ -34,7 +34,7 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数
| db2 | IBM DB2 | DB2 JDBC |
| informix | IBM Informix | Informix JDBC |
| neo4j | Neo4j | Neo4j JDBC |
| cassandra | Apache Cassandra | Cassandra JDBC |
| cassandra | Apache Cassandra 2.1+ | Apache cassandra-gocql-driver 原生 Agent |
| bigquery | Google BigQuery | BigQuery JDBC |
| kylin | Apache Kylin | Kylin JDBC |
| sundb | SunDB | SunDB JDBC |
@ -49,7 +49,7 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数
## 多 JRE 支持
多数 Java agent 以 JRE 21 为目标。原生 agent`oracle`、`kingbase`、`xugu` 和 `rabbitmq`)不需要 JRE。对 Java agentDBX 会自动下载并管理 JRE 21 安装。
多数 Java agent 以 JRE 21 为目标。原生 agent`cassandra`、`oracle`、`kingbase`、`xugu` 和 `rabbitmq`)不需要 JRE。对 Java agentDBX 会自动下载并管理 JRE 21 安装。
## JDBC 连接池
@ -75,7 +75,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J
对于新 agent只要存在成熟、许可证兼容的原生驱动优先选择**原生Go 或 Rust驱动**而非 Java/JDBC agent。原生 agent 以单一自包含可执行文件发布,无需 JRE可显著降低内存占用和启动时间 —— 完全避开 Java agent 即便空闲也要付出的 JVM 基线开销。
- **原生Go/Rust** —— 存在可用原生驱动时首选。参考 `drivers/oracle-go`go-ora、`drivers/kingbase-go`gokb、`drivers/vastbase-go`openGauss connector、`drivers/xugu` 和 `drivers/rabbitmq`amqp091-go。无需 JRE 下载与管理。
- **原生Go/Rust** —— 存在可用原生驱动时首选。参考 `drivers/cassandra-go`Apache cassandra-gocql-driver、`drivers/oracle-go`go-ora、`drivers/kingbase-go`gokb、`drivers/vastbase-go`openGauss connector、`drivers/xugu` 和 `drivers/rabbitmq`amqp091-go。无需 JRE 下载与管理。
- **Java/JDBC** —— 当某数据库只有 JDBC 驱动,或原生驱动不成熟、缺乏维护时的默认兜底方案。多数 agent 仍属此类。
原生 agent 实现与 Java agent 相同的 JSON-RPC 契约和 `versions.json` 登记;它发布的是 `agent` 可执行文件而非 `agent.jar`。若同一数据库同时保留原生和 Java 源码实现,默认只发布原生产物;只有 Java 变体以独立兼容配置登记时才同时发布,例如 `oracle-legacy` / `oracle-10g`
@ -87,13 +87,14 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J
```bash
./gradlew shadowJar
(cd drivers/oracle-go && go build -o agent .)
(cd drivers/cassandra-go && go build -o agent .)
(cd drivers/kingbase-go && go build -o agent .)
(cd drivers/vastbase-go && go build -o agent .)
(cd drivers/xugu && go build -o agent .)
(cd drivers/rabbitmq && go build -o agent .)
```
产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/xugu` 和 `drivers/rabbitmq` 构建。
产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/cassandra-go`、`drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/xugu` 和 `drivers/rabbitmq` 构建。
### 本地 DBX 运行时测试
@ -107,7 +108,7 @@ cp agents/drivers/<db_type>/build/libs/*-all.jar ~/.dbx/agents/drivers/<db_type>
重启 DBX 或断开重连数据库,使新 agent 进程加载替换后的 JAR。
`oracle`、`kingbase`、`xugu` 和 `rabbitmq` 等原生 agent 使用驱动目录下的 `agent` 可执行文件而非 `agent.jar`
`cassandra`、`oracle`、`kingbase`、`xugu` 和 `rabbitmq` 等原生 agent 使用驱动目录下的 `agent` 可执行文件而非 `agent.jar`
## 版本管理

View File

@ -5,7 +5,7 @@ plugins {
def infrastructureProjects = ['common', 'test-support'] as Set
def legacyStandaloneProjects = ['mongodb', 'kafka', 'rocketmq'] as Set
def pooledJdbcProjects = [
'access', 'bigquery', 'cassandra', 'dameng', 'databend', 'databricks', 'db2', 'exasol',
'access', 'bigquery', 'dameng', 'databend', 'databricks', 'db2', 'exasol',
'firebird', 'gbase8a', 'gbase8s', 'goldendb', 'h2', 'h2-legacy', 'highgo', 'hive',
'informix', 'iotdb', 'iris', 'kylin', 'neo4j', 'oceanbase-oracle', 'oscar', 'saphana',
'snowflake', 'spark', 'sqlserver-legacy', 'sundb', 'tdengine', 'teradata', 'trino', 'uxdb',

View File

@ -0,0 +1,317 @@
package main
import (
"fmt"
"net"
"net/url"
"strconv"
"strings"
"time"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
type cassandraConfig struct {
hosts []string
port int
keyspace string
username string
password string
localDatacenter string
requestTimeout time.Duration
connectTimeout time.Duration
protocolVersion int
consistency string
serialConsistency string
numConnections int
pageSize int
cqlVersion string
ssl bool
caCertPath string
clientCertPath string
clientKeyPath string
hostVerification bool
disableInitialHostLookup bool
}
func parseCassandraConfig(cp connectParams) (cassandraConfig, error) {
config := cassandraConfig{
port: 9042,
keyspace: strings.TrimSpace(cp.Database),
username: cp.Username,
password: cp.Password,
requestTimeout: 11 * time.Second,
connectTimeout: defaultConnectTimeout,
numConnections: 2,
pageSize: 5000,
ssl: cp.SSL,
caCertPath: cp.CACertPath,
clientCertPath: cp.ClientCertPath,
clientKeyPath: cp.ClientKeyPath,
}
if cp.Port > 0 {
config.port = cp.Port
}
params := url.Values{}
if strings.TrimSpace(cp.ConnectionString) != "" {
if err := applyConnectionString(&config, params, cp.ConnectionString); err != nil {
return cassandraConfig{}, err
}
}
if len(config.hosts) == 0 {
config.hosts = splitHosts(cp.Host)
}
if len(config.hosts) == 0 {
return cassandraConfig{}, fmt.Errorf("Cassandra host is required")
}
urlParams, err := parseURLParams(cp.URLParams)
if err != nil {
return cassandraConfig{}, err
}
for key, values := range urlParams {
params[key] = values
}
if err := applyCassandraURLParams(&config, params); err != nil {
return cassandraConfig{}, err
}
if !config.disableInitialHostLookup && allLoopbackHosts(config.hosts) {
config.disableInitialHostLookup = true
}
return config, nil
}
func applyConnectionString(config *cassandraConfig, params url.Values, raw string) error {
value := strings.TrimSpace(raw)
value = strings.TrimPrefix(value, "jdbc:")
if !strings.Contains(value, "://") {
return fmt.Errorf("unsupported Cassandra connection string: %s", raw)
}
parsed, err := url.Parse(value)
if err != nil {
return fmt.Errorf("invalid Cassandra connection string: %w", err)
}
if parsed.Scheme != "cassandra" {
return fmt.Errorf("unsupported Cassandra connection scheme: %s", parsed.Scheme)
}
if parsed.User != nil {
config.username = parsed.User.Username()
if password, ok := parsed.User.Password(); ok {
config.password = password
}
}
config.hosts = splitHosts(parsed.Hostname())
if strings.Contains(parsed.Host, ",") {
hostPart := parsed.Host
if parsed.User != nil {
hostPart = strings.TrimPrefix(hostPart, parsed.User.String()+"@")
}
config.hosts = splitHosts(hostPart)
}
if port := parsed.Port(); port != "" {
parsedPort, parseErr := strconv.Atoi(port)
if parseErr != nil || parsedPort < 1 || parsedPort > 65535 {
return fmt.Errorf("invalid Cassandra port: %s", port)
}
config.port = parsedPort
}
if keyspace := strings.Trim(strings.TrimSpace(parsed.Path), "/"); keyspace != "" {
config.keyspace = keyspace
}
for key, values := range parsed.Query() {
params[key] = values
}
return nil
}
func parseURLParams(raw string) (url.Values, error) {
raw = strings.TrimPrefix(strings.TrimSpace(raw), "?")
if raw == "" {
return url.Values{}, nil
}
values, err := url.ParseQuery(raw)
if err != nil {
return nil, fmt.Errorf("invalid Cassandra URL parameters: %w", err)
}
return values, nil
}
func applyCassandraURLParams(config *cassandraConfig, params url.Values) error {
for rawKey, values := range params {
if len(values) == 0 {
continue
}
key := normalizeOptionName(rawKey)
value := strings.TrimSpace(values[len(values)-1])
switch key {
case "localdatacenter", "datacenter", "dc":
config.localDatacenter = value
case "requesttimeout", "timeout":
duration, err := parseDurationOption(value)
if err != nil {
return fmt.Errorf("invalid requesttimeout: %w", err)
}
config.requestTimeout = duration
case "connecttimeout", "logintimeout":
duration, err := parseDurationOption(value)
if err != nil {
return fmt.Errorf("invalid connecttimeout: %w", err)
}
config.connectTimeout = duration
case "protocolversion", "protoversion":
version, err := strconv.Atoi(value)
if err != nil || version < 3 || version > 5 {
return fmt.Errorf("protocolversion must be between 3 and 5")
}
config.protocolVersion = version
case "consistency":
if _, err := gocql.ParseConsistencyWrapper(value); err != nil {
return err
}
config.consistency = value
case "serialconsistency":
consistency, err := gocql.ParseConsistencyWrapper(value)
if err != nil {
return err
}
if consistency != gocql.Serial && consistency != gocql.LocalSerial {
return fmt.Errorf("serialconsistency must be SERIAL or LOCAL_SERIAL")
}
config.serialConsistency = value
case "numconns", "connectionsperhost":
count, err := strconv.Atoi(value)
if err != nil || count < 1 || count > 32 {
return fmt.Errorf("numconns must be between 1 and 32")
}
config.numConnections = count
case "pagesize", "fetchsize":
size, err := strconv.Atoi(value)
if err != nil || size < 1 {
return fmt.Errorf("pagesize must be positive")
}
config.pageSize = size
case "cqlversion":
config.cqlVersion = value
case "ssl":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid ssl option: %w", err)
}
config.ssl = enabled
case "hostverification", "verifyhostname", "sslhostnameverification":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid host verification option: %w", err)
}
config.hostVerification = enabled
case "disableinitialhostlookup":
disabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid disableinitialhostlookup option: %w", err)
}
config.disableInitialHostLookup = disabled
case "loadbalancing":
if value != "" && !strings.EqualFold(value, "DcInferringLoadBalancingPolicy") {
return fmt.Errorf("unsupported Cassandra loadbalancing policy: %s", value)
}
default:
return fmt.Errorf("unsupported Cassandra URL parameter: %s", rawKey)
}
}
return nil
}
func (config cassandraConfig) clusterConfig(keyspace string) (*gocql.ClusterConfig, error) {
cluster := gocql.NewCluster(config.hosts...)
cluster.Port = config.port
cluster.Keyspace = strings.TrimSpace(keyspace)
cluster.Timeout = config.requestTimeout
cluster.ConnectTimeout = config.connectTimeout
cluster.WriteTimeout = config.requestTimeout
cluster.NumConns = config.numConnections
cluster.PageSize = config.pageSize
cluster.DisableInitialHostLookup = config.disableInitialHostLookup
cluster.IgnorePeerAddr = config.disableInitialHostLookup
if config.protocolVersion != 0 {
cluster.ProtoVersion = config.protocolVersion
}
if config.cqlVersion != "" {
cluster.CQLVersion = config.cqlVersion
}
if config.consistency != "" {
consistency, err := gocql.ParseConsistencyWrapper(config.consistency)
if err != nil {
return nil, err
}
cluster.Consistency = consistency
}
if config.serialConsistency != "" {
consistency, err := gocql.ParseConsistencyWrapper(config.serialConsistency)
if err != nil {
return nil, err
}
cluster.SerialConsistency = consistency
}
if config.username != "" {
cluster.Authenticator = gocql.PasswordAuthenticator{Username: config.username, Password: config.password}
}
if config.ssl {
cluster.SslOpts = &gocql.SslOptions{
CaPath: config.caCertPath,
CertPath: config.clientCertPath,
KeyPath: config.clientKeyPath,
EnableHostVerification: config.hostVerification,
}
}
if config.localDatacenter != "" {
cluster.PoolConfig.HostSelectionPolicy = gocql.TokenAwareHostPolicy(
gocql.DCAwareRoundRobinPolicy(config.localDatacenter),
)
}
return cluster, nil
}
func splitHosts(raw string) []string {
parts := strings.FieldsFunc(raw, func(char rune) bool { return char == ',' || char == ';' })
hosts := make([]string, 0, len(parts))
for _, part := range parts {
host := strings.TrimSpace(part)
if host == "" {
continue
}
if parsedHost, _, err := net.SplitHostPort(host); err == nil {
host = parsedHost
}
hosts = append(hosts, strings.Trim(host, "[]"))
}
return hosts
}
func allLoopbackHosts(hosts []string) bool {
for _, host := range hosts {
if strings.EqualFold(host, "localhost") {
continue
}
ip := net.ParseIP(host)
if ip == nil || !ip.IsLoopback() {
return false
}
}
return len(hosts) > 0
}
func parseDurationOption(value string) (time.Duration, error) {
if duration, err := time.ParseDuration(value); err == nil {
return duration, nil
}
milliseconds, err := strconv.Atoi(value)
if err != nil || milliseconds < 1 {
return 0, fmt.Errorf("expected duration or positive milliseconds")
}
return time.Duration(milliseconds) * time.Millisecond, nil
}
func normalizeOptionName(value string) string {
return strings.NewReplacer("_", "", "-", "", ".", "").Replace(strings.ToLower(strings.TrimSpace(value)))
}

View File

@ -0,0 +1,82 @@
package main
import (
"testing"
"time"
)
func TestParseCassandraConfigSupportsLegacyJDBCOptions(t *testing.T) {
config, err := parseCassandraConfig(connectParams{
Host: "127.0.0.1",
Database: "app",
Username: "cassandra",
Password: "secret",
URLParams: "?localdatacenter=dc1&requesttimeout=10000&connecttimeout=5s&protocolversion=4&consistency=local_quorum&numconns=4",
})
if err != nil {
t.Fatal(err)
}
if len(config.hosts) != 1 || config.hosts[0] != "127.0.0.1" {
t.Fatalf("unexpected hosts: %#v", config.hosts)
}
if config.port != 9042 || config.keyspace != "app" {
t.Fatalf("unexpected endpoint: port=%d keyspace=%q", config.port, config.keyspace)
}
if config.localDatacenter != "dc1" || config.protocolVersion != 4 {
t.Fatalf("unexpected topology config: %#v", config)
}
if config.requestTimeout != 10*time.Second || config.connectTimeout != 5*time.Second {
t.Fatalf("unexpected timeouts: request=%s connect=%s", config.requestTimeout, config.connectTimeout)
}
if config.numConnections != 4 || !config.disableInitialHostLookup {
t.Fatalf("unexpected pool/tunnel config: %#v", config)
}
}
func TestParseCassandraConfigAcceptsConnectionString(t *testing.T) {
config, err := parseCassandraConfig(connectParams{
ConnectionString: "jdbc:cassandra://alice:secret@db.example.com:9142/catalog?protocolversion=5",
})
if err != nil {
t.Fatal(err)
}
if len(config.hosts) != 1 || config.hosts[0] != "db.example.com" || config.port != 9142 {
t.Fatalf("unexpected endpoint: %#v", config)
}
if config.keyspace != "catalog" || config.username != "alice" || config.password != "secret" {
t.Fatalf("unexpected credentials/keyspace: %#v", config)
}
if config.protocolVersion != 5 {
t.Fatalf("unexpected protocol version: %d", config.protocolVersion)
}
}
func TestParseCassandraConfigRejectsUnsupportedLoadBalancingClass(t *testing.T) {
_, err := parseCassandraConfig(connectParams{
Host: "localhost",
URLParams: "loadbalancing=example.CustomPolicy",
})
if err == nil {
t.Fatal("expected unsupported load-balancing policy error")
}
}
func TestParseCassandraConfigRejectsCassandra20Protocol(t *testing.T) {
_, err := parseCassandraConfig(connectParams{
Host: "localhost",
URLParams: "protocolversion=2",
})
if err == nil {
t.Fatal("expected native protocol v2 rejection")
}
}
func TestParseDurationOptionTreatsBareNumbersAsMilliseconds(t *testing.T) {
duration, err := parseDurationOption("1500")
if err != nil {
t.Fatal(err)
}
if duration != 1500*time.Millisecond {
t.Fatalf("unexpected duration: %s", duration)
}
}

View File

@ -0,0 +1,7 @@
module github.com/t8y2/dbx/agents/drivers/cassandra-go
go 1.22
require github.com/apache/cassandra-gocql-driver/v2 v2.1.2
require gopkg.in/inf.v0 v0.9.1 // indirect

View File

@ -0,0 +1,28 @@
github.com/apache/cassandra-gocql-driver/v2 v2.1.2 h1:lu/p0Db2av18enHJvWJQoChLssI0P+AR06STq4VdvCc=
github.com/apache/cassandra-gocql-driver/v2 v2.1.2/go.mod h1:QH/asJjB3mHvY6Dot6ZKMMpTcOrWJ8i9GhsvG1g0PK4=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/golang/snappy v0.0.3 h1:fHPg5GQYlCeLIPB9BZqMVR5nR9A+IM5zcgeTdjMYmLA=
github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/pierrec/lz4/v4 v4.1.8 h1:ieHkV+i2BRzngO4Wd/3HGowuZStgq6QkPsD1eolNAO4=
github.com/pierrec/lz4/v4 v4.1.8/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ=
go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc=
gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

View File

@ -0,0 +1,553 @@
package main
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"os"
"strings"
"sync"
"time"
)
const (
protocolVersion = 2
defaultMaxRows = 10000
defaultPageSize = 500
legacyAgentSessionID = "__legacy__"
maxAgentSessions = 256
defaultConnectTimeout = 15 * time.Second
)
type request struct {
ID json.RawMessage `json:"id"`
Method string `json:"method"`
Params map[string]json.RawMessage `json:"params"`
}
type response struct {
JSONRPC string `json:"jsonrpc,omitempty"`
ID json.RawMessage `json:"id,omitempty"`
Result any `json:"result,omitempty"`
Error *rpcError `json:"error,omitempty"`
}
type connectParams struct {
Host string `json:"host"`
Port int `json:"port"`
Database string `json:"database"`
Username string `json:"username"`
Password string `json:"password"`
URLParams string `json:"url_params"`
ConnectionString string `json:"connection_string"`
SSL bool `json:"ssl"`
CACertPath string `json:"ca_cert_path"`
ClientCertPath string `json:"client_cert_path"`
ClientKeyPath string `json:"client_key_path"`
SessionRole string `json:"sessionRole"`
}
type queryOptions struct {
SQL string `json:"sql"`
Database string `json:"database"`
Schema string `json:"schema"`
MaxRows int `json:"maxRows"`
FetchSize int `json:"fetchSize"`
TimeoutSecs int `json:"timeoutSecs"`
}
type queryResult struct {
Columns []string `json:"columns"`
ColumnTypes []string `json:"column_types"`
Rows [][]any `json:"rows"`
AffectedRows int64 `json:"affected_rows"`
ExecutionTimeMS int64 `json:"execution_time_ms"`
Truncated bool `json:"truncated"`
}
type queryPageResult struct {
Columns []string `json:"columns"`
ColumnTypes []string `json:"column_types"`
Rows [][]any `json:"rows"`
AffectedRows int64 `json:"affected_rows"`
ExecutionTimeMS int64 `json:"execution_time_ms"`
Truncated bool `json:"truncated"`
SessionID *string `json:"session_id"`
HasMore bool `json:"has_more"`
}
type querySession struct {
sql string
keyspace string
pageState []byte
remaining int
}
type server struct {
runtime *connectionRuntime
params connectParams
querySessions map[string]*querySession
nextSessionID uint64
activeMu sync.Mutex
activeCancel context.CancelFunc
}
type agentSession struct {
server *server
runtimeKey string
mu sync.Mutex
}
type runtimeServer struct {
mu sync.RWMutex
sessions map[string]*agentSession
runtimesMu sync.Mutex
runtimes map[string]*connectionRuntime
}
func main() {
runtime := newRuntimeServer()
encoder := json.NewEncoder(os.Stdout)
var encoderMu sync.Mutex
var requests sync.WaitGroup
fmt.Fprintln(os.Stdout, `{"ready":true}`)
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 512*1024*1024)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" {
continue
}
var envelope request
if json.Unmarshal([]byte(line), &envelope) == nil && envelope.Method == "shutdown" {
requests.Wait()
resp, _ := runtime.handleLine(line)
encoderMu.Lock()
_ = encoder.Encode(resp)
encoderMu.Unlock()
return
}
requests.Add(1)
go func(line string) {
defer requests.Done()
resp, _ := runtime.handleLine(line)
encoderMu.Lock()
defer encoderMu.Unlock()
if err := encoder.Encode(resp); err != nil {
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
}
}(line)
}
requests.Wait()
}
func newRuntimeServer() *runtimeServer {
return &runtimeServer{
sessions: map[string]*agentSession{},
runtimes: map[string]*connectionRuntime{},
}
}
func (r *runtimeServer) handleLine(line string) (response, bool) {
var req request
if err := json.Unmarshal([]byte(line), &req); err != nil {
return errorResponse(nil, "", "", err), false
}
if len(req.ID) == 0 {
req.ID = json.RawMessage("1")
}
result, shutdown, err := r.dispatch(req.Method, req.Params)
if err != nil {
return errorResponse(req.ID, req.Method, stringParam(req.Params, "agentSessionId"), err), false
}
return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown
}
func (r *runtimeServer) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) {
switch method {
case "handshake":
return handshakeResult(true), false, nil
case "open_session":
id := stringParam(params, "agentSessionId")
if id == "" {
return nil, false, errors.New("agentSessionId is required")
}
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
return map[string]bool{"ok": true}, false, r.openSession(id, cp)
case "close_session":
return map[string]bool{"ok": true}, false, r.closeSession(stringParam(params, "agentSessionId"))
case "validate_session":
session, err := r.session(stringParam(params, "agentSessionId"))
if err != nil {
return nil, false, err
}
session.mu.Lock()
defer session.mu.Unlock()
return map[string]bool{"ok": true}, false, session.server.validateConnection()
case "cancel_session":
session, err := r.session(stringParam(params, "agentSessionId"))
if err != nil {
return nil, false, err
}
session.server.cancelActiveQuery()
return map[string]bool{"ok": true}, false, nil
case "test_connection":
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
result, err := testConnection(cp)
return result, false, err
case "connect":
var cp connectParams
if err := decodeParams(params, &cp); err != nil {
return nil, false, err
}
_ = r.closeSession(legacyAgentSessionID)
return map[string]bool{"ok": true}, false, r.openSession(legacyAgentSessionID, cp)
case "disconnect":
return map[string]bool{"ok": true}, false, r.closeSession(legacyAgentSessionID)
case "shutdown":
return map[string]bool{"ok": true}, true, r.closeAllSessions()
default:
id := stringParam(params, "agentSessionId")
if id == "" {
id = legacyAgentSessionID
}
session, err := r.session(id)
if err != nil {
return nil, false, err
}
session.mu.Lock()
defer session.mu.Unlock()
release, err := session.server.runtime.acquire(isMetadataOperation(method))
if err != nil {
return nil, false, err
}
defer release()
return session.server.dispatch(method, params)
}
}
func (r *runtimeServer) openSession(id string, cp connectParams) error {
r.mu.Lock()
if _, exists := r.sessions[id]; exists {
r.mu.Unlock()
return fmt.Errorf("agent session already exists: %s", id)
}
if len(r.sessions) >= maxAgentSessions {
r.mu.Unlock()
return fmt.Errorf("agent session limit reached: %d", maxAgentSessions)
}
r.mu.Unlock()
runtime, key, err := r.acquireRuntime(cp)
if err != nil {
return err
}
s := newServer(runtime, cp)
if err := s.validateConnection(); err != nil {
r.releaseRuntime(key)
return err
}
r.mu.Lock()
defer r.mu.Unlock()
if _, exists := r.sessions[id]; exists {
r.releaseRuntime(key)
return fmt.Errorf("agent session already exists: %s", id)
}
r.sessions[id] = &agentSession{server: s, runtimeKey: key}
return nil
}
func (r *runtimeServer) session(id string) (*agentSession, error) {
r.mu.RLock()
session := r.sessions[id]
r.mu.RUnlock()
if session == nil {
return nil, fmt.Errorf("agent session not found: %s", id)
}
return session, nil
}
func (r *runtimeServer) closeSession(id string) error {
r.mu.Lock()
session := r.sessions[id]
delete(r.sessions, id)
r.mu.Unlock()
if session == nil {
return nil
}
session.server.cancelActiveQuery()
session.mu.Lock()
session.server.disconnect()
session.mu.Unlock()
r.releaseRuntime(session.runtimeKey)
return nil
}
func (r *runtimeServer) closeAllSessions() error {
r.mu.RLock()
ids := make([]string, 0, len(r.sessions))
for id := range r.sessions {
ids = append(ids, id)
}
r.mu.RUnlock()
for _, id := range ids {
_ = r.closeSession(id)
}
r.runtimesMu.Lock()
runtimes := r.runtimes
r.runtimes = map[string]*connectionRuntime{}
r.runtimesMu.Unlock()
for _, runtime := range runtimes {
runtime.close()
}
return nil
}
func newServer(runtime *connectionRuntime, cp connectParams) *server {
return &server{runtime: runtime, params: cp, querySessions: map[string]*querySession{}}
}
func (s *server) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) {
switch method {
case "handshake":
return handshakeResult(false), false, nil
case "validate_connection":
return map[string]bool{"ok": true}, false, s.validateConnection()
case "connection_info":
result, err := s.connectionInfo()
return result, false, err
case "list_databases":
result, err := s.listDatabases()
return result, false, err
case "list_schemas":
result, err := s.listSchemas()
return result, false, err
case "list_tables":
result, err := s.listTables(stringParam(params, "schema"), metadataListConstraintsFromParams(params))
return result, false, err
case "get_table_comment":
return nil, false, nil
case "list_objects":
result, err := s.listObjects(stringParam(params, "schema"), metadataListConstraintsFromParams(params))
return result, false, err
case "list_data_types":
return cassandraDataTypes(), false, nil
case "completion_assistant_search_v1":
var input completionAssistantRequest
if err := decodeParams(params, &input); err != nil {
return nil, false, err
}
result, err := s.completionAssistantSearch(input)
return result, false, err
case "get_columns":
result, err := s.getColumns(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_indexes":
result, err := s.listIndexes(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "list_foreign_keys":
return []foreignKeyInfo{}, false, nil
case "list_triggers":
return []triggerInfo{}, false, nil
case "get_object_source":
return nil, false, errors.New("object source is not supported by Cassandra")
case "get_table_ddl":
result, err := s.getTableDDL(stringParam(params, "schema"), stringParam(params, "table"))
return result, false, err
case "get_explain_info":
return nil, false, errors.New("execution plans are not supported by Cassandra")
case "execute_query":
result, err := s.executeQuery(queryOptionsFromParams(params))
return result, false, err
case "execute_query_page", "start_table_read":
result, err := s.executeQueryPage(queryOptionsFromParams(params), intParam(params, "pageSize"))
return result, false, err
case "fetch_query_page", "fetch_table_read_page":
result, err := s.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize"))
return result, false, err
case "close_query_session", "close_table_read_session":
return s.closeQuerySession(stringParam(params, "sessionId")), false, nil
case "execute_transaction":
result, err := s.executeStatements(params, true)
return result, false, err
case "execute_batch":
result, err := s.executeStatements(params, false)
return result, false, err
case "disconnect":
s.disconnect()
return map[string]bool{"ok": true}, false, nil
case "shutdown":
s.disconnect()
return map[string]bool{"ok": true}, true, nil
default:
return nil, false, fmt.Errorf("unknown method: %s", method)
}
}
func handshakeResult(multiSession bool) map[string]any {
capabilities := []string{
"connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "structured_error_v1",
}
if multiSession {
capabilities = append(capabilities, "multi_session")
}
return map[string]any{
"protocolVersion": protocolVersion,
"agentProtocolVersion": protocolVersion,
"capabilities": capabilities,
}
}
func (s *server) validateConnection() error {
ctx, cancel := context.WithTimeout(context.Background(), defaultConnectTimeout)
defer cancel()
session, err := s.runtime.sessionFor(s.defaultKeyspace())
if err != nil {
return err
}
var releaseVersion string
return session.Query("SELECT release_version FROM system.local").WithContext(ctx).Scan(&releaseVersion)
}
func testConnection(cp connectParams) (map[string]any, error) {
runtime, err := newConnectionRuntime(cp)
if err != nil {
return nil, err
}
defer runtime.close()
s := newServer(runtime, cp)
if err := s.validateConnection(); err != nil {
return nil, err
}
info, err := s.connectionInfo()
if err != nil {
return nil, err
}
return map[string]any{"ok": true, "info": info}, nil
}
func (s *server) disconnect() {
s.cancelActiveQuery()
s.querySessions = map[string]*querySession{}
}
func (s *server) defaultKeyspace() string {
if keyspace := strings.TrimSpace(s.params.Database); keyspace != "" {
return keyspace
}
return strings.TrimSpace(s.runtime.config.keyspace)
}
func (s *server) beginOperation(timeoutSecs int) (context.Context, context.CancelFunc) {
var ctx context.Context
var cancel context.CancelFunc
if timeoutSecs > 0 {
ctx, cancel = context.WithTimeout(context.Background(), time.Duration(timeoutSecs)*time.Second)
} else {
ctx, cancel = context.WithCancel(context.Background())
}
s.activeMu.Lock()
s.activeCancel = cancel
s.activeMu.Unlock()
return ctx, cancel
}
func (s *server) endOperation(cancel context.CancelFunc) {
cancel()
s.activeMu.Lock()
s.activeCancel = nil
s.activeMu.Unlock()
}
func (s *server) cancelActiveQuery() {
s.activeMu.Lock()
cancel := s.activeCancel
s.activeMu.Unlock()
if cancel != nil {
cancel()
}
}
func queryOptionsFromParams(params map[string]json.RawMessage) queryOptions {
return queryOptions{
SQL: stringParam(params, "sql"),
Database: stringParam(params, "database"),
Schema: stringParam(params, "schema"),
MaxRows: intParam(params, "maxRows"),
FetchSize: intParam(params, "fetchSize"),
TimeoutSecs: intParam(params, "timeoutSecs"),
}
}
func decodeParams(params map[string]json.RawMessage, target any) error {
data, err := json.Marshal(params)
if err != nil {
return err
}
return json.Unmarshal(data, target)
}
func stringParam(params map[string]json.RawMessage, key string) string {
if raw, ok := params[key]; ok {
var value string
if json.Unmarshal(raw, &value) == nil {
return value
}
}
return ""
}
func intParam(params map[string]json.RawMessage, key string) int {
if raw, ok := params[key]; ok {
var value int
if json.Unmarshal(raw, &value) == nil {
return value
}
}
return 0
}
func boolParam(params map[string]json.RawMessage, key string) bool {
if raw, ok := params[key]; ok {
var value bool
if json.Unmarshal(raw, &value) == nil {
return value
}
}
return false
}
func stringSliceParam(params map[string]json.RawMessage, key string) []string {
if raw, ok := params[key]; ok {
var value []string
if json.Unmarshal(raw, &value) == nil {
return value
}
}
return []string{}
}
func errorResponse(id json.RawMessage, method, sessionID string, err error) response {
return response{JSONRPC: "2.0", ID: id, Error: classifyRPCError(method, sessionID, err)}
}
func isMetadataOperation(method string) bool {
switch method {
case "connection_info", "list_databases", "list_schemas", "list_tables", "get_table_comment", "list_objects",
"list_data_types", "completion_assistant_search_v1", "get_columns", "list_indexes", "list_foreign_keys",
"list_triggers", "get_object_source", "get_table_ddl", "get_explain_info":
return true
default:
return false
}
}

View File

@ -0,0 +1,586 @@
package main
import (
"encoding/json"
"fmt"
"sort"
"strings"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
var cassandraTypes = []string{
"ascii", "bigint", "blob", "boolean", "counter", "date", "decimal", "double", "duration",
"float", "inet", "int", "list", "map", "set", "smallint", "text", "time", "timestamp",
"timeuuid", "tinyint", "tuple", "uuid", "varchar", "varint", "vector", "frozen",
}
type databaseInfo struct {
Name string `json:"name"`
}
type tableInfo struct {
Name string `json:"name"`
TableType string `json:"table_type"`
Comment *string `json:"comment"`
}
type objectInfo struct {
Name string `json:"name"`
ObjectType string `json:"object_type"`
Schema string `json:"schema"`
Comment *string `json:"comment"`
Valid *bool `json:"valid,omitempty"`
}
type columnInfo struct {
Name string `json:"name"`
DataType string `json:"data_type"`
IsNullable bool `json:"is_nullable"`
ColumnDefault *string `json:"column_default"`
IsPrimaryKey bool `json:"is_primary_key"`
Extra *string `json:"extra"`
Comment *string `json:"comment"`
NumericPrecision *int `json:"numeric_precision"`
NumericScale *int `json:"numeric_scale"`
CharacterMaximumLength *int `json:"character_maximum_length"`
}
type indexInfo struct {
Name string `json:"name"`
Columns []string `json:"columns"`
IsUnique bool `json:"is_unique"`
IsPrimary bool `json:"is_primary"`
Filter *string `json:"filter"`
IndexType *string `json:"index_type"`
IncludedColumns []string `json:"included_columns"`
Comment *string `json:"comment"`
}
func (i indexInfo) MarshalJSON() ([]byte, error) {
type alias indexInfo
value := alias(i)
if value.Columns == nil {
value.Columns = []string{}
}
if value.IncludedColumns == nil {
value.IncludedColumns = []string{}
}
return json.Marshal(value)
}
type foreignKeyInfo struct {
Name string `json:"name"`
Column string `json:"column"`
RefTable string `json:"ref_table"`
RefColumn string `json:"ref_column"`
}
type triggerInfo struct {
Name string `json:"name"`
Event string `json:"event"`
Timing string `json:"timing"`
}
type metadataListConstraints struct {
Filter string
Limit int
Offset int
ObjectTypes []string
}
type completionAssistantRequest struct {
ConnectionID string `json:"connection_id"`
Database string `json:"database"`
Schema string `json:"schema"`
ObjectKinds []string `json:"object_kinds"`
Mask string `json:"mask"`
CaseSensitive bool `json:"case_sensitive"`
GlobalSearch bool `json:"global_search"`
MaxResults int `json:"max_results"`
ParentSchema string `json:"parent_schema"`
ParentName string `json:"parent_name"`
MatchMode string `json:"match_mode"`
}
type completionAssistantCandidate struct {
Name string `json:"name"`
Kind string `json:"kind"`
Database *string `json:"database"`
Schema *string `json:"schema"`
ParentSchema *string `json:"parent_schema"`
ParentName *string `json:"parent_name"`
Comment *string `json:"comment"`
DataType *string `json:"data_type"`
}
type completionAssistantResponse struct {
Candidates []completionAssistantCandidate `json:"candidates"`
Incomplete bool `json:"incomplete"`
FallbackUsed bool `json:"fallback_used"`
}
func cassandraDataTypes() []string {
return append([]string(nil), cassandraTypes...)
}
func (s *server) connectionInfo() (map[string]any, error) {
session, err := s.runtime.sessionFor("")
if err != nil {
return nil, err
}
var clusterName, version, cqlVersion, dataCenter string
err = session.Query("SELECT cluster_name, release_version, cql_version, data_center FROM system.local").Scan(
&clusterName, &version, &cqlVersion, &dataCenter,
)
if err != nil {
return nil, err
}
return map[string]any{
"database": s.defaultKeyspace(),
"schema": s.defaultKeyspace(),
"username": s.params.Username,
"version": version,
"clusterName": clusterName,
"cqlVersion": cqlVersion,
"localDatacenter": dataCenter,
"identifierQuote": `"`,
"compatibilityMode": "cql",
"databaseInfo": map[string]string{
"productName": "Apache Cassandra",
"productVersion": version,
"unquotedIdentifierCase": "lower",
"quotedIdentifierCase": "mixed",
"driverName": "Apache cassandra-gocql-driver",
"driverVersion": "2.1.2",
},
}, nil
}
func (s *server) allKeyspaceMetadata() (map[string]*gocql.KeyspaceMetadata, error) {
session, err := s.runtime.sessionFor("")
if err != nil {
return nil, err
}
return session.AllKeyspaceMetadata()
}
func (s *server) keyspaceMetadata(schema string) (*gocql.KeyspaceMetadata, error) {
session, err := s.runtime.sessionFor("")
if err != nil {
return nil, err
}
metadata, err := session.KeyspaceMetadata(schema)
if err != nil {
return nil, err
}
if metadata == nil {
return nil, fmt.Errorf("Cassandra keyspace not found: %s", schema)
}
return metadata, nil
}
func (s *server) tableMetadata(schema, table string) (*gocql.TableMetadata, error) {
keyspace, err := s.keyspaceMetadata(schema)
if err != nil {
return nil, err
}
metadata := keyspace.Tables[table]
if metadata == nil {
return nil, fmt.Errorf("Cassandra table not found: %s.%s", schema, table)
}
return metadata, nil
}
func (s *server) listDatabases() ([]databaseInfo, error) {
metadata, err := s.allKeyspaceMetadata()
if err != nil {
return nil, err
}
names := sortedMapKeys(metadata)
result := make([]databaseInfo, len(names))
for index, name := range names {
result[index] = databaseInfo{Name: name}
}
return result, nil
}
func (s *server) listSchemas() ([]string, error) {
databases, err := s.listDatabases()
if err != nil {
return nil, err
}
result := make([]string, len(databases))
for index, database := range databases {
result[index] = database.Name
}
return result, nil
}
func (s *server) listTables(schema string, constraints metadataListConstraints) ([]tableInfo, error) {
metadata, err := s.keyspaceMetadata(schema)
if err != nil {
return nil, err
}
names := sortedMapKeys(metadata.Tables)
result := make([]tableInfo, 0, len(names))
for _, name := range names {
if !metadataNameMatches(name, constraints.Filter) {
continue
}
result = append(result, tableInfo{Name: name, TableType: "TABLE"})
}
return applyMetadataWindow(result, constraints.Offset, constraints.Limit), nil
}
func (s *server) listObjects(schema string, constraints metadataListConstraints) ([]objectInfo, error) {
tables, err := s.listTables(schema, metadataListConstraints{Filter: constraints.Filter})
if err != nil {
return nil, err
}
allowed := stringSet(constraints.ObjectTypes)
result := make([]objectInfo, 0, len(tables))
for _, table := range tables {
if len(allowed) > 0 && !allowed["table"] && !allowed["base_table"] {
continue
}
result = append(result, objectInfo{Name: table.Name, ObjectType: "TABLE", Schema: schema})
}
return applyMetadataWindow(result, constraints.Offset, constraints.Limit), nil
}
func (s *server) getColumns(schema, table string) ([]columnInfo, error) {
metadata, err := s.tableMetadata(schema, table)
if err != nil {
return nil, err
}
return columnsFromMetadata(metadata), nil
}
func columnsFromMetadata(metadata *gocql.TableMetadata) []columnInfo {
names := orderedColumnNames(metadata)
result := make([]columnInfo, 0, len(names))
for _, name := range names {
column := metadata.Columns[name]
if column == nil {
continue
}
primary := column.Kind == gocql.ColumnPartitionKey || column.Kind == gocql.ColumnClusteringKey
extra := column.Kind.String()
result = append(result, columnInfo{
Name: column.Name,
DataType: cqlTypeName(column.Type),
IsNullable: !primary,
IsPrimaryKey: primary,
Extra: &extra,
})
}
return result
}
func (s *server) listIndexes(schema, table string) ([]indexInfo, error) {
metadata, err := s.tableMetadata(schema, table)
if err != nil {
return nil, err
}
result := indexesFromMetadata(metadata)
queried, queryErr := s.querySystemIndexes(schema, table)
if queryErr == nil {
result = mergeIndexes(result, queried)
}
return result, nil
}
func (s *server) querySystemIndexes(schema, table string) ([]indexInfo, error) {
session, err := s.runtime.sessionFor("")
if err != nil {
return nil, err
}
iter := session.Query(
"SELECT index_name, kind, options FROM system_schema.indexes WHERE keyspace_name = ? AND table_name = ?",
schema,
table,
).Iter()
result := []indexInfo{}
var name, kind string
var options map[string]string
for iter.Scan(&name, &kind, &options) {
indexType := strings.TrimSpace(kind)
result = append(result, indexInfo{
Name: name,
Columns: targetColumns(options["target"]),
IndexType: optionalString(indexType),
IncludedColumns: []string{},
})
options = nil
}
if err := iter.Close(); err != nil {
return nil, err
}
sort.Slice(result, func(left, right int) bool { return result[left].Name < result[right].Name })
return result, nil
}
func mergeIndexes(first, second []indexInfo) []indexInfo {
byName := make(map[string]indexInfo, len(first)+len(second))
for _, index := range first {
byName[index.Name] = index
}
for _, index := range second {
if existing, ok := byName[index.Name]; ok && len(index.Columns) == 0 {
index.Columns = existing.Columns
}
byName[index.Name] = index
}
names := sortedMapKeys(byName)
result := make([]indexInfo, 0, len(names))
for _, name := range names {
result = append(result, byName[name])
}
return result
}
func targetColumns(target string) []string {
target = strings.TrimSpace(target)
for _, wrapper := range []string{"values", "keys", "entries", "full"} {
prefix := wrapper + "("
if strings.HasPrefix(strings.ToLower(target), prefix) && strings.HasSuffix(target, ")") {
target = strings.TrimSpace(target[len(prefix) : len(target)-1])
break
}
}
target = strings.Trim(target, `"'`)
if target == "" {
return []string{}
}
return []string{target}
}
func optionalString(value string) *string {
if value == "" {
return nil
}
return &value
}
func indexesFromMetadata(metadata *gocql.TableMetadata) []indexInfo {
byName := map[string]*indexInfo{}
for _, columnName := range orderedColumnNames(metadata) {
column := metadata.Columns[columnName]
if column == nil || strings.TrimSpace(column.Index.Name) == "" {
continue
}
index := byName[column.Index.Name]
if index == nil {
indexType := strings.TrimSpace(column.Index.Type)
index = &indexInfo{Name: column.Index.Name, Columns: []string{}, IncludedColumns: []string{}}
if indexType != "" {
index.IndexType = &indexType
}
byName[index.Name] = index
}
index.Columns = append(index.Columns, columnName)
}
names := sortedMapKeys(byName)
result := make([]indexInfo, 0, len(names))
for _, name := range names {
result = append(result, *byName[name])
}
return result
}
func (s *server) getTableDDL(schema, table string) (string, error) {
metadata, err := s.tableMetadata(schema, table)
if err != nil {
return "", err
}
return tableDDLFromMetadata(schema, table, metadata)
}
func tableDDLFromMetadata(schema, table string, metadata *gocql.TableMetadata) (string, error) {
definitions := make([]string, 0, len(metadata.Columns)+1)
for _, name := range orderedColumnNames(metadata) {
column := metadata.Columns[name]
if column != nil {
definitions = append(definitions, " "+quoteCQLIdentifier(column.Name)+" "+cqlTypeName(column.Type))
}
}
partitionKeys := metadataColumnNames(metadata.PartitionKey)
clusteringKeys := metadataColumnNames(metadata.ClusteringColumns)
if len(partitionKeys) == 0 {
return "", fmt.Errorf("Cassandra table has no partition key: %s.%s", schema, table)
}
primaryParts := make([]string, 0, len(clusteringKeys)+1)
if len(partitionKeys) == 1 {
primaryParts = append(primaryParts, quoteCQLIdentifier(partitionKeys[0]))
} else {
quoted := make([]string, len(partitionKeys))
for index, name := range partitionKeys {
quoted[index] = quoteCQLIdentifier(name)
}
primaryParts = append(primaryParts, "("+strings.Join(quoted, ", ")+")")
}
for _, name := range clusteringKeys {
primaryParts = append(primaryParts, quoteCQLIdentifier(name))
}
definitions = append(definitions, " PRIMARY KEY ("+strings.Join(primaryParts, ", ")+")")
ddl := "CREATE TABLE " + quoteCQLIdentifier(schema) + "." + quoteCQLIdentifier(table) + " (\n" + strings.Join(definitions, ",\n") + "\n)"
orders := make([]string, 0, len(metadata.ClusteringColumns))
for _, column := range metadata.ClusteringColumns {
if column != nil {
order := "ASC"
if column.Order == gocql.DESC {
order = "DESC"
}
orders = append(orders, quoteCQLIdentifier(column.Name)+" "+order)
}
}
if len(orders) > 0 {
ddl += " WITH CLUSTERING ORDER BY (" + strings.Join(orders, ", ") + ")"
}
return ddl + ";", nil
}
func (s *server) completionAssistantSearch(input completionAssistantRequest) (completionAssistantResponse, error) {
limit := input.MaxResults
if limit <= 0 || limit > 1000 {
limit = 100
}
candidates := make([]completionAssistantCandidate, 0, limit+1)
kinds := stringSet(input.ObjectKinds)
if kinds["column"] && input.ParentName != "" {
schema := input.ParentSchema
if schema == "" {
schema = input.Schema
}
columns, err := s.getColumns(schema, input.ParentName)
if err != nil {
return completionAssistantResponse{}, err
}
for _, column := range columns {
if !completionNameMatches(column.Name, input) {
continue
}
dataType := column.DataType
candidates = append(candidates, completionAssistantCandidate{
Name: column.Name, Kind: "COLUMN", Schema: stringPtr(schema), ParentSchema: stringPtr(schema),
ParentName: stringPtr(input.ParentName), DataType: &dataType,
})
}
} else {
schemas := []string{input.Schema}
if input.GlobalSearch || input.Schema == "" {
var err error
schemas, err = s.listSchemas()
if err != nil {
return completionAssistantResponse{}, err
}
}
for _, schema := range schemas {
objects, err := s.listObjects(schema, metadataListConstraints{ObjectTypes: input.ObjectKinds})
if err != nil {
return completionAssistantResponse{}, err
}
for _, object := range objects {
if !completionNameMatches(object.Name, input) {
continue
}
candidates = append(candidates, completionAssistantCandidate{
Name: object.Name, Kind: object.ObjectType, Schema: stringPtr(schema),
})
if len(candidates) > limit {
return completionAssistantResponse{Candidates: candidates[:limit], Incomplete: true}, nil
}
}
}
}
incomplete := len(candidates) > limit
if incomplete {
candidates = candidates[:limit]
}
return completionAssistantResponse{Candidates: candidates, Incomplete: incomplete}, nil
}
func metadataListConstraintsFromParams(params map[string]json.RawMessage) metadataListConstraints {
return metadataListConstraints{
Filter: stringParam(params, "filter"),
Limit: intParam(params, "limit"),
Offset: intParam(params, "offset"),
ObjectTypes: stringSliceParam(params, "object_types"),
}
}
func orderedColumnNames(metadata *gocql.TableMetadata) []string {
if len(metadata.OrderedColumns) > 0 {
return append([]string(nil), metadata.OrderedColumns...)
}
return sortedMapKeys(metadata.Columns)
}
func metadataColumnNames(columns []*gocql.ColumnMetadata) []string {
result := make([]string, 0, len(columns))
for _, column := range columns {
if column != nil {
result = append(result, column.Name)
}
}
return result
}
func metadataNameMatches(name, filter string) bool {
return filter == "" || strings.Contains(strings.ToLower(name), strings.ToLower(filter))
}
func applyMetadataWindow[T any](values []T, offset, limit int) []T {
if offset < 0 {
offset = 0
}
if offset >= len(values) {
return []T{}
}
values = values[offset:]
if limit > 0 && limit < len(values) {
values = values[:limit]
}
return values
}
func sortedMapKeys[T any](values map[string]T) []string {
keys := make([]string, 0, len(values))
for key := range values {
keys = append(keys, key)
}
sort.Strings(keys)
return keys
}
func stringSet(values []string) map[string]bool {
result := make(map[string]bool, len(values))
for _, value := range values {
result[strings.ToLower(strings.TrimSpace(value))] = true
}
return result
}
func completionNameMatches(name string, input completionAssistantRequest) bool {
mask := input.Mask
if mask == "" {
return true
}
if !input.CaseSensitive {
name = strings.ToLower(name)
mask = strings.ToLower(mask)
}
if strings.EqualFold(input.MatchMode, "contains") {
return strings.Contains(name, mask)
}
return strings.HasPrefix(name, mask)
}
func quoteCQLIdentifier(value string) string {
return `"` + strings.ReplaceAll(value, `"`, `""`) + `"`
}
func stringPtr(value string) *string {
return &value
}

View File

@ -0,0 +1,86 @@
package main
import (
"reflect"
"testing"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
func TestColumnsIndexesAndDDLFromMetadata(t *testing.T) {
textType := gocql.NewNativeType(4, gocql.TypeVarchar, "")
intType := gocql.NewNativeType(4, gocql.TypeInt, "")
id := &gocql.ColumnMetadata{Name: "tenant", Kind: gocql.ColumnPartitionKey, Type: textType}
bucket := &gocql.ColumnMetadata{Name: "bucket", Kind: gocql.ColumnPartitionKey, Type: intType}
created := &gocql.ColumnMetadata{Name: "created_at", Kind: gocql.ColumnClusteringKey, Type: textType, Order: gocql.DESC}
email := &gocql.ColumnMetadata{
Name: "email", Kind: gocql.ColumnRegular, Type: textType,
Index: gocql.ColumnIndexMetadata{Name: "users_email_idx", Type: "COMPOSITES"},
}
metadata := &gocql.TableMetadata{
OrderedColumns: []string{"tenant", "bucket", "created_at", "email"},
PartitionKey: []*gocql.ColumnMetadata{id, bucket},
ClusteringColumns: []*gocql.ColumnMetadata{created},
Columns: map[string]*gocql.ColumnMetadata{
"tenant": id, "bucket": bucket, "created_at": created, "email": email,
},
}
columns := columnsFromMetadata(metadata)
if len(columns) != 4 || !columns[0].IsPrimaryKey || columns[0].IsNullable || columns[3].IsPrimaryKey || !columns[3].IsNullable {
t.Fatalf("unexpected columns: %#v", columns)
}
if columns[2].Extra == nil || *columns[2].Extra != "clustering_key" {
t.Fatalf("unexpected clustering metadata: %#v", columns[2])
}
indexes := indexesFromMetadata(metadata)
if len(indexes) != 1 || indexes[0].Name != "users_email_idx" || !reflect.DeepEqual(indexes[0].Columns, []string{"email"}) {
t.Fatalf("unexpected indexes: %#v", indexes)
}
ddl, err := tableDDLFromMetadata("app", "users", metadata)
if err != nil {
t.Fatal(err)
}
want := "CREATE TABLE \"app\".\"users\" (\n" +
" \"tenant\" text,\n" +
" \"bucket\" int,\n" +
" \"created_at\" text,\n" +
" \"email\" text,\n" +
" PRIMARY KEY ((\"tenant\", \"bucket\"), \"created_at\")\n" +
") WITH CLUSTERING ORDER BY (\"created_at\" DESC);"
if ddl != want {
t.Fatalf("unexpected DDL:\n%s\nwant:\n%s", ddl, want)
}
}
func TestMetadataWindowAndFilter(t *testing.T) {
values := []string{"a", "b", "c", "d"}
if got := applyMetadataWindow(values, 1, 2); !reflect.DeepEqual(got, []string{"b", "c"}) {
t.Fatalf("unexpected window: %#v", got)
}
if !metadataNameMatches("CustomerEvents", "event") || metadataNameMatches("users", "event") {
t.Fatal("metadata filter mismatch")
}
}
func TestTargetColumnsHandlesCollectionIndexes(t *testing.T) {
for input, want := range map[string]string{
"txt": "txt",
"values(tags)": "tags",
"keys(attrs)": "attrs",
`entries("attrs")`: "attrs",
} {
got := targetColumns(input)
if !reflect.DeepEqual(got, []string{want}) {
t.Fatalf("targetColumns(%q) = %#v", input, got)
}
}
}
func TestQuoteCQLIdentifierEscapesQuotes(t *testing.T) {
if got := quoteCQLIdentifier(`a"b`); got != `"a""b"` {
t.Fatalf("unexpected quoted identifier: %s", got)
}
}

View File

@ -0,0 +1,156 @@
package main
import (
"context"
"errors"
"fmt"
"io"
"net"
"strings"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
type rpcError struct {
Code int `json:"code"`
Message string `json:"message"`
Data *rpcErrorData `json:"data,omitempty"`
}
type rpcErrorData struct {
Category string `json:"category"`
Retryable bool `json:"retryable"`
SessionDisposition string `json:"sessionDisposition"`
Stage string `json:"stage"`
ContractVersion int `json:"contractVersion"`
OperationOutcome string `json:"operationOutcome"`
SQLState string `json:"sqlState,omitempty"`
ExceptionClass string `json:"exceptionClass,omitempty"`
AgentSessionID string `json:"agentSessionId,omitempty"`
}
func classifyRPCError(method, agentSessionID string, err error) *rpcError {
stage := rpcErrorStage(method)
data := &rpcErrorData{
Category: "protocol",
Retryable: false,
SessionDisposition: "keep",
Stage: stage,
ContractVersion: 1,
OperationOutcome: rpcOperationOutcome(stage),
ExceptionClass: safeRPCDiagnostic(fmt.Sprintf("%T", err), 160),
AgentSessionID: strings.TrimSpace(agentSessionID),
}
if errors.Is(err, errOperationCapacity) {
data.Category = "resource"
data.Retryable = true
return &rpcError{Code: -1, Message: err.Error(), Data: data}
}
var requestError gocql.RequestError
if errors.As(err, &requestError) {
data.SQLState = fmt.Sprintf("0x%04x", requestError.Code())
switch requestError.Code() {
case gocql.ErrCodeUnavailable, gocql.ErrCodeOverloaded, gocql.ErrCodeBootstrapping:
data.Category = "resource"
data.Retryable = true
case gocql.ErrCodeWriteTimeout, gocql.ErrCodeReadTimeout:
data.Category = "timeout"
data.Retryable = true
case gocql.ErrCodeCredentials:
data.Category = "connection"
data.Retryable = stage == "connect" || stage == "validate"
case gocql.ErrCodeSyntax, gocql.ErrCodeUnauthorized, gocql.ErrCodeInvalid,
gocql.ErrCodeConfig, gocql.ErrCodeAlreadyExists, gocql.ErrCodeUnprepared:
data.Category = "sql"
default:
data.Category = "sql"
}
} else if errors.Is(err, context.Canceled) {
data.Category = "canceled"
data.SessionDisposition = "quarantine"
} else if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, gocql.ErrTimeoutNoResponse) || isTimeoutError(err) {
data.Category = "timeout"
data.SessionDisposition = "quarantine"
} else if isConnectionError(err) {
data.Category = "connection"
data.Retryable = stage == "connect" || stage == "validate"
if stage != "connect" {
data.SessionDisposition = "quarantine"
}
}
return &rpcError{Code: -1, Message: err.Error(), Data: data}
}
func rpcErrorStage(method string) string {
switch method {
case "connect", "open_session", "test_connection":
return "connect"
case "validate_connection", "validate_session":
return "validate"
case "cancel_session":
return "cancel"
case "close_session", "disconnect", "close_query_session", "close_table_read_session", "shutdown":
return "close"
case "fetch_query_page", "fetch_table_read_page":
return "fetch"
case "handshake", "":
return "request"
default:
return "execute"
}
}
func rpcOperationOutcome(stage string) string {
switch stage {
case "request", "connect", "validate":
return "not_started"
default:
return "unknown"
}
}
func isTimeoutError(err error) bool {
var timeout interface{ Timeout() bool }
return errors.As(err, &timeout) && timeout.Timeout()
}
func isConnectionError(err error) bool {
if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
return true
}
var networkError *net.OpError
if errors.As(err, &networkError) {
return true
}
lower := strings.ToLower(err.Error())
for _, marker := range []string{
"connection refused",
"connection reset",
"broken pipe",
"connection closed",
"connection lost",
"unexpected eof",
"no route to host",
"no hosts available",
} {
if strings.Contains(lower, marker) {
return true
}
}
return false
}
func safeRPCDiagnostic(value string, maxLength int) string {
var result strings.Builder
for _, char := range value {
if result.Len() >= maxLength {
break
}
if char >= 0x21 && char <= 0x7e {
result.WriteRune(char)
}
}
return result.String()
}

View File

@ -0,0 +1,50 @@
package main
import (
"context"
"encoding/json"
"reflect"
"testing"
)
func TestHandshakeAdvertisesMultiSessionAndStructuredErrors(t *testing.T) {
result, shutdown, err := newRuntimeServer().dispatch("handshake", nil)
if err != nil || shutdown {
t.Fatalf("unexpected handshake result: shutdown=%t err=%v", shutdown, err)
}
capabilities := result.(map[string]any)["capabilities"].([]string)
want := []string{"connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "structured_error_v1", "multi_session"}
if !reflect.DeepEqual(capabilities, want) {
t.Fatalf("unexpected capabilities: %#v", capabilities)
}
}
func TestHandleLineRejectsMissingSession(t *testing.T) {
params, _ := json.Marshal(map[string]any{"agentSessionId": "missing"})
line := `{"jsonrpc":"2.0","id":7,"method":"validate_session","params":` + string(params) + `}`
response, _ := newRuntimeServer().handleLine(line)
if response.Error == nil || response.Error.Data == nil || response.Error.Data.Stage != "validate" {
t.Fatalf("unexpected error response: %#v", response)
}
}
func TestClassifyCanceledQuery(t *testing.T) {
err := classifyRPCError("execute_query", "session-1", context.Canceled)
if err.Data.Category != "canceled" || err.Data.SessionDisposition != "quarantine" {
t.Fatalf("unexpected cancellation classification: %#v", err)
}
}
func TestRuntimeIdentityIncludesCredentials(t *testing.T) {
first := connectionRuntimeKey(connectParams{Host: "localhost", Username: "user", Password: "one"})
second := connectionRuntimeKey(connectParams{Host: "localhost", Username: "user", Password: "two"})
if first == second {
t.Fatal("runtime identities must not share sessions across credentials")
}
}
func TestTrimStatementSQL(t *testing.T) {
if got := trimStatementSQL(" SELECT * FROM t;;; \n"); got != "SELECT * FROM t" {
t.Fatalf("unexpected trimmed SQL: %q", got)
}
}

View File

@ -0,0 +1,329 @@
package main
import (
"encoding/json"
"fmt"
"reflect"
"strings"
"time"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
func (s *server) executeQuery(options queryOptions) (queryResult, error) {
start := time.Now()
maxRows := options.MaxRows
if maxRows <= 0 {
maxRows = defaultMaxRows
}
session, err := s.runtime.sessionFor(s.keyspaceForOptions(options))
if err != nil {
return queryResult{}, err
}
ctx, cancel := s.beginOperation(options.TimeoutSecs)
defer s.endOperation(cancel)
query := session.Query(trimStatementSQL(options.SQL)).WithContext(ctx)
if options.FetchSize > 0 {
query = query.PageSize(options.FetchSize)
}
iter := query.Iter()
columns := iter.Columns()
result := queryResult{
Columns: columnNames(columns),
ColumnTypes: columnTypeNames(columns),
Rows: make([][]any, 0, min(maxRows, 1024)),
}
if len(columns) == 0 {
err := iter.Close()
result.ExecutionTimeMS = time.Since(start).Milliseconds()
return result, err
}
for len(result.Rows) < maxRows {
row, ok, scanErr := scanCQLRow(iter, columns)
if scanErr != nil {
_ = iter.Close()
return queryResult{}, scanErr
}
if !ok {
break
}
result.Rows = append(result.Rows, row)
}
if len(result.Rows) == maxRows {
_, hasExtra, scanErr := scanCQLRow(iter, columns)
if scanErr != nil {
_ = iter.Close()
return queryResult{}, scanErr
}
result.Truncated = hasExtra
}
if err := iter.Close(); err != nil {
return queryResult{}, err
}
result.ExecutionTimeMS = time.Since(start).Milliseconds()
return result, nil
}
func (s *server) executeQueryPage(options queryOptions, pageSize int) (queryPageResult, error) {
if pageSize <= 0 {
pageSize = options.FetchSize
}
if pageSize <= 0 {
pageSize = defaultPageSize
}
remaining := options.MaxRows
if remaining <= 0 {
remaining = defaultMaxRows
}
result, nextState, err := s.fetchCQLPage(options.SQL, s.keyspaceForOptions(options), nil, pageSize, remaining, options.TimeoutSecs)
if err != nil {
return queryPageResult{}, err
}
remaining -= len(result.Rows)
if len(nextState) == 0 || remaining <= 0 {
result.HasMore = false
result.Truncated = len(nextState) > 0 && remaining <= 0
return result, nil
}
s.nextSessionID++
id := fmt.Sprintf("cassandra-query-%d", s.nextSessionID)
s.querySessions[id] = &querySession{
sql: trimStatementSQL(options.SQL),
keyspace: s.keyspaceForOptions(options),
pageState: append([]byte(nil), nextState...),
remaining: remaining,
}
result.SessionID = &id
result.HasMore = true
return result, nil
}
func (s *server) fetchQueryPage(id string, pageSize int) (queryPageResult, error) {
state := s.querySessions[id]
if state == nil {
return queryPageResult{}, fmt.Errorf("query session not found: %s", id)
}
if pageSize <= 0 {
pageSize = defaultPageSize
}
result, nextState, err := s.fetchCQLPage(state.sql, state.keyspace, state.pageState, pageSize, state.remaining, 0)
if err != nil {
return queryPageResult{}, err
}
state.remaining -= len(result.Rows)
if len(nextState) == 0 || state.remaining <= 0 {
delete(s.querySessions, id)
result.HasMore = false
result.Truncated = len(nextState) > 0 && state.remaining <= 0
return result, nil
}
state.pageState = append(state.pageState[:0], nextState...)
result.SessionID = &id
result.HasMore = true
return result, nil
}
func (s *server) fetchCQLPage(sql, keyspace string, pageState []byte, pageSize, remaining, timeoutSecs int) (queryPageResult, []byte, error) {
start := time.Now()
if remaining < pageSize {
pageSize = remaining
}
if pageSize <= 0 {
return queryPageResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}}, nil, nil
}
session, err := s.runtime.sessionFor(keyspace)
if err != nil {
return queryPageResult{}, nil, err
}
ctx, cancel := s.beginOperation(timeoutSecs)
defer s.endOperation(cancel)
iter := session.Query(trimStatementSQL(sql)).WithContext(ctx).PageSize(pageSize).PageState(pageState).Iter()
columns := iter.Columns()
result := queryPageResult{
Columns: columnNames(columns),
ColumnTypes: columnTypeNames(columns),
Rows: make([][]any, 0, pageSize),
}
if len(columns) == 0 {
err := iter.Close()
result.ExecutionTimeMS = time.Since(start).Milliseconds()
return result, nil, err
}
for len(result.Rows) < pageSize {
row, ok, scanErr := scanCQLRow(iter, columns)
if scanErr != nil {
_ = iter.Close()
return queryPageResult{}, nil, scanErr
}
if !ok {
break
}
result.Rows = append(result.Rows, row)
}
nextState := append([]byte(nil), iter.PageState()...)
if err := iter.Close(); err != nil {
return queryPageResult{}, nil, err
}
result.ExecutionTimeMS = time.Since(start).Milliseconds()
return result, nextState, nil
}
func (s *server) closeQuerySession(id string) bool {
if _, exists := s.querySessions[id]; !exists {
return false
}
delete(s.querySessions, id)
return true
}
func (s *server) executeStatements(params map[string]json.RawMessage, transactional bool) (queryResult, error) {
statements := stringSliceParam(params, "statements")
if len(statements) == 0 {
return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}}, nil
}
keyspace := strings.TrimSpace(stringParam(params, "schema"))
if keyspace == "" {
keyspace = strings.TrimSpace(stringParam(params, "database"))
}
if keyspace == "" {
keyspace = strings.TrimSpace(s.params.Database)
}
session, err := s.runtime.sessionFor(keyspace)
if err != nil {
return queryResult{}, err
}
batchType := gocql.UnloggedBatch
if transactional {
batchType = gocql.LoggedBatch
}
batch := session.NewBatch(batchType)
for _, statement := range statements {
statement = trimStatementSQL(statement)
if statement != "" {
batch.Query(statement)
}
}
ctx, cancel := s.beginOperation(intParam(params, "timeoutSecs"))
defer s.endOperation(cancel)
start := time.Now()
if err := session.ExecuteBatch(batch.WithContext(ctx)); err != nil {
return queryResult{}, err
}
return queryResult{
Columns: []string{},
ColumnTypes: []string{},
Rows: [][]any{},
AffectedRows: 0,
ExecutionTimeMS: time.Since(start).Milliseconds(),
}, nil
}
func (s *server) keyspaceForOptions(options queryOptions) string {
if schema := strings.TrimSpace(options.Schema); schema != "" {
return schema
}
if database := strings.TrimSpace(options.Database); database != "" {
return database
}
return s.defaultKeyspace()
}
func scanCQLRow(iter *gocql.Iter, columns []gocql.ColumnInfo) ([]any, bool, error) {
destinations := make([]any, 0, len(columns))
extractors := make([]func() any, 0, len(columns))
for _, column := range columns {
if tuple, ok := column.TypeInfo.(gocql.TupleTypeInfo); ok {
tupleDestinations := make([]*cqlDestination, 0, len(tuple.Elems))
for _, element := range tuple.Elems {
destination := newCQLDestination(element)
tupleDestinations = append(tupleDestinations, destination)
destinations = append(destinations, destination.destination)
}
extractors = append(extractors, func() any {
values := make([]any, len(tupleDestinations))
allNull := true
for index, destination := range tupleDestinations {
value, present := destination.value()
if present {
allNull = false
values[index] = value
}
}
if allNull {
return nil
}
return normalizeCQLValue(values)
})
continue
}
destination := newCQLDestination(column.TypeInfo)
destinations = append(destinations, destination.destination)
extractors = append(extractors, func() any {
value, present := destination.value()
if !present {
return nil
}
return normalizeCQLValue(value)
})
}
if !iter.Scan(destinations...) {
return nil, false, nil
}
row := make([]any, len(columns))
for index, extract := range extractors {
row[index] = extract()
}
return row, true, nil
}
type cqlDestination struct {
destination any
holder reflect.Value
fallback *any
}
func newCQLDestination(typeInfo gocql.TypeInfo) *cqlDestination {
zero := typeInfo.Zero()
valueType := reflect.TypeOf(zero)
if valueType == nil {
var fallback any
return &cqlDestination{destination: &fallback, fallback: &fallback}
}
holder := reflect.New(reflect.PointerTo(valueType))
return &cqlDestination{destination: holder.Interface(), holder: holder}
}
func (destination *cqlDestination) value() (any, bool) {
if destination.fallback != nil {
return *destination.fallback, *destination.fallback != nil
}
pointer := destination.holder.Elem()
if pointer.IsNil() {
return nil, false
}
return pointer.Elem().Interface(), true
}
func columnNames(columns []gocql.ColumnInfo) []string {
result := make([]string, len(columns))
for index, column := range columns {
result[index] = column.Name
}
return result
}
func columnTypeNames(columns []gocql.ColumnInfo) []string {
result := make([]string, len(columns))
for index, column := range columns {
result[index] = cqlTypeName(column.TypeInfo)
}
return result
}
func trimStatementSQL(sql string) string {
trimmed := strings.TrimSpace(sql)
for strings.HasSuffix(trimmed, ";") {
trimmed = strings.TrimSpace(strings.TrimSuffix(trimmed, ";"))
}
return trimmed
}

View File

@ -0,0 +1,176 @@
package main
import (
"context"
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
"os"
"strconv"
"strings"
"sync"
"time"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
const (
defaultRuntimePoolSize = 32
defaultRuntimeMetadataLimit = 8
operationPermitTimeout = 30 * time.Second
)
var errOperationCapacity = errors.New("agent operation capacity is temporarily exhausted")
type connectionRuntime struct {
mu sync.Mutex
config cassandraConfig
sessions map[string]*gocql.Session
permits chan struct{}
metadataPermits chan struct{}
references int
closed bool
}
func newConnectionRuntime(cp connectParams) (*connectionRuntime, error) {
config, err := parseCassandraConfig(cp)
if err != nil {
return nil, err
}
poolSize := runtimePoolSize()
return &connectionRuntime{
config: config,
sessions: map[string]*gocql.Session{},
permits: make(chan struct{}, poolSize),
metadataPermits: make(chan struct{}, runtimeMetadataLimit(poolSize)),
}, nil
}
func (r *connectionRuntime) sessionFor(keyspace string) (*gocql.Session, error) {
keyspace = strings.TrimSpace(keyspace)
r.mu.Lock()
defer r.mu.Unlock()
if r.closed {
return nil, errors.New("Cassandra connection runtime is closed")
}
if session := r.sessions[keyspace]; session != nil && !session.Closed() {
return session, nil
}
cluster, err := r.config.clusterConfig(keyspace)
if err != nil {
return nil, err
}
session, err := cluster.CreateSession()
if err != nil {
return nil, err
}
r.sessions[keyspace] = session
return session, nil
}
func (r *connectionRuntime) acquire(metadata bool) (func(), error) {
ctx, cancel := context.WithTimeout(context.Background(), operationPermitTimeout)
defer cancel()
metadataAcquired := false
if metadata {
select {
case r.metadataPermits <- struct{}{}:
metadataAcquired = true
case <-ctx.Done():
return nil, errOperationCapacity
}
}
select {
case r.permits <- struct{}{}:
return func() {
<-r.permits
if metadataAcquired {
<-r.metadataPermits
}
}, nil
case <-ctx.Done():
if metadataAcquired {
<-r.metadataPermits
}
return nil, errOperationCapacity
}
}
func (r *connectionRuntime) close() {
r.mu.Lock()
if r.closed {
r.mu.Unlock()
return
}
r.closed = true
sessions := r.sessions
r.sessions = map[string]*gocql.Session{}
r.mu.Unlock()
for _, session := range sessions {
session.Close()
}
}
func (r *runtimeServer) acquireRuntime(cp connectParams) (*connectionRuntime, string, error) {
key := connectionRuntimeKey(cp)
r.runtimesMu.Lock()
defer r.runtimesMu.Unlock()
runtime := r.runtimes[key]
if runtime == nil {
var err error
runtime, err = newConnectionRuntime(cp)
if err != nil {
return nil, "", err
}
r.runtimes[key] = runtime
}
runtime.references++
return runtime, key, nil
}
func (r *runtimeServer) releaseRuntime(key string) {
if key == "" {
return
}
r.runtimesMu.Lock()
runtime := r.runtimes[key]
shouldClose := false
if runtime != nil && runtime.references > 0 {
runtime.references--
}
if runtime != nil && runtime.references == 0 {
delete(r.runtimes, key)
shouldClose = true
}
r.runtimesMu.Unlock()
if shouldClose {
runtime.close()
}
}
func connectionRuntimeKey(cp connectParams) string {
data, _ := json.Marshal(cp)
digest := sha256.Sum256(data)
return fmt.Sprintf("%x", digest[:])
}
func runtimePoolSize() int {
value := defaultRuntimePoolSize
if raw := os.Getenv("DBX_AGENT_CASSANDRA_MAX_CONCURRENT_OPERATIONS"); raw != "" {
if parsed, err := strconv.Atoi(raw); err == nil && parsed >= 1 && parsed <= 128 {
value = parsed
}
}
return value
}
func runtimeMetadataLimit(poolSize int) int {
value := min(defaultRuntimeMetadataLimit, poolSize)
if raw := os.Getenv("DBX_AGENT_CASSANDRA_MAX_CONCURRENT_METADATA"); raw != "" {
if parsed, err := strconv.Atoi(raw); err == nil && parsed >= 1 && parsed <= poolSize {
value = parsed
}
}
return value
}

View File

@ -0,0 +1,119 @@
package main
import (
"encoding/hex"
"fmt"
"math/big"
"net"
"reflect"
"sort"
"strings"
"time"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
func normalizeCQLValue(value any) any {
if value == nil {
return nil
}
return cqlString(value)
}
func cqlString(value any) string {
if value == nil {
return "null"
}
switch typed := value.(type) {
case string:
return typed
case []byte:
return "0x" + hex.EncodeToString(typed)
case time.Time:
return typed.Format(time.RFC3339Nano)
case time.Duration:
return typed.String()
case gocql.Duration:
return fmt.Sprintf("%dmo%dd%dns", typed.Months, typed.Days, typed.Nanoseconds)
case gocql.UUID:
return typed.String()
case net.IP:
return typed.String()
case *big.Int:
if typed == nil {
return ""
}
return typed.String()
case big.Int:
return typed.String()
case fmt.Stringer:
return typed.String()
}
valueOf := reflect.ValueOf(value)
for valueOf.Kind() == reflect.Pointer {
if valueOf.IsNil() {
return ""
}
valueOf = valueOf.Elem()
}
switch valueOf.Kind() {
case reflect.Map:
entries := make([]string, 0, valueOf.Len())
iterator := valueOf.MapRange()
for iterator.Next() {
entries = append(entries, cqlString(iterator.Key().Interface())+"="+cqlString(iterator.Value().Interface()))
}
sort.Strings(entries)
return "{" + strings.Join(entries, ", ") + "}"
case reflect.Slice, reflect.Array:
values := make([]string, valueOf.Len())
for index := range values {
values[index] = cqlString(valueOf.Index(index).Interface())
}
return "[" + strings.Join(values, ", ") + "]"
default:
return fmt.Sprint(value)
}
}
func cqlTypeName(typeInfo gocql.TypeInfo) string {
if typeInfo == nil {
return "unknown"
}
switch typed := typeInfo.(type) {
case gocql.CollectionType:
switch typed.Type() {
case gocql.TypeMap:
return "map<" + cqlTypeName(typed.Key) + ", " + cqlTypeName(typed.Elem) + ">"
case gocql.TypeList:
return "list<" + cqlTypeName(typed.Elem) + ">"
case gocql.TypeSet:
return "set<" + cqlTypeName(typed.Elem) + ">"
}
case gocql.TupleTypeInfo:
parts := make([]string, len(typed.Elems))
for index, element := range typed.Elems {
parts[index] = cqlTypeName(element)
}
return "tuple<" + strings.Join(parts, ", ") + ">"
case gocql.UDTTypeInfo:
return quoteCQLIdentifier(typed.Name)
case gocql.VectorType:
return fmt.Sprintf("vector<%s, %d>", cqlTypeName(typed.SubType), typed.Dimensions)
}
names := map[gocql.Type]string{
gocql.TypeCustom: "custom", gocql.TypeAscii: "ascii", gocql.TypeBigInt: "bigint",
gocql.TypeBlob: "blob", gocql.TypeBoolean: "boolean", gocql.TypeCounter: "counter",
gocql.TypeDecimal: "decimal", gocql.TypeDouble: "double", gocql.TypeFloat: "float",
gocql.TypeInt: "int", gocql.TypeText: "text", gocql.TypeTimestamp: "timestamp",
gocql.TypeUUID: "uuid", gocql.TypeVarchar: "text", gocql.TypeVarint: "varint",
gocql.TypeTimeUUID: "timeuuid", gocql.TypeInet: "inet", gocql.TypeDate: "date",
gocql.TypeTime: "time", gocql.TypeSmallInt: "smallint", gocql.TypeTinyInt: "tinyint",
gocql.TypeDuration: "duration", gocql.TypeUDT: "udt", gocql.TypeTuple: "tuple",
gocql.TypeList: "list", gocql.TypeMap: "map", gocql.TypeSet: "set",
}
if name := names[typeInfo.Type()]; name != "" {
return name
}
return "unknown"
}

View File

@ -0,0 +1,45 @@
package main
import (
"net"
"reflect"
"testing"
"time"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
func TestNormalizeCQLValuePreservesLegacyStringContract(t *testing.T) {
uuid, err := gocql.ParseUUID("00112233-4455-6677-8899-aabbccddeeff")
if err != nil {
t.Fatal(err)
}
tests := []struct {
value any
want any
}{
{nil, nil},
{42, "42"},
{true, "true"},
{[]byte{0x00, 0xff}, "0x00ff"},
{uuid, "00112233-4455-6677-8899-aabbccddeeff"},
{net.ParseIP("127.0.0.1"), "127.0.0.1"},
{time.Date(2026, 8, 3, 12, 34, 56, 7, time.UTC), "2026-08-03T12:34:56.000000007Z"},
{gocql.Duration{Months: 1, Days: 2, Nanoseconds: 3}, "1mo2d3ns"},
{[]int{1, 2}, "[1, 2]"},
{[]any{1, nil, "three"}, "[1, null, three]"},
{map[string]int{"b": 2, "a": 1}, "{a=1, b=2}"},
}
for _, test := range tests {
if got := normalizeCQLValue(test.value); !reflect.DeepEqual(got, test.want) {
t.Fatalf("normalizeCQLValue(%#v) = %#v, want %#v", test.value, got, test.want)
}
}
}
func TestCQLTypeNameUsesCQLSyntax(t *testing.T) {
typeInfo := gocql.NewNativeType(4, gocql.TypeList, "varchar")
if got := cqlTypeName(typeInfo); got != "list<text>" {
t.Fatalf("unexpected collection type name: %s", got)
}
}

View File

@ -1,7 +1,7 @@
driver strategy scope reason
access intentional-fallback java-jdbc-metadata Uses Access JDBC metadata; no portable stable server-side paging API, common constraints filter locally.
bigquery native-pushdown java-sql Uses dataset INFORMATION_SCHEMA.TABLES with type, normal-name filter, stable order, and literal LIMIT/OFFSET.
cassandra intentional-fallback java-cql Uses Cassandra system_schema; portable LIKE/fuzzy paging is not safe across versions, common constraints filter locally.
cassandra shared-fallback native-go Uses cassandra-gocql-driver schema metadata, then applies stable filtering and paging in the native agent.
dameng native-pushdown java-sql Uses ALL_OBJECTS with type, filter, stable order, and LIMIT/OFFSET with legacy fallback.
databend native-pushdown java-sql Uses system procedures for routine pushdown and shared table metadata constraints for table-like objects.
db2 native-pushdown java-sql Uses SYSCAT metadata with type, filter, stable order, and OFFSET/FETCH.

1 driver strategy scope reason
2 access intentional-fallback java-jdbc-metadata Uses Access JDBC metadata; no portable stable server-side paging API, common constraints filter locally.
3 bigquery native-pushdown java-sql Uses dataset INFORMATION_SCHEMA.TABLES with type, normal-name filter, stable order, and literal LIMIT/OFFSET.
4 cassandra intentional-fallback shared-fallback java-cql native-go Uses Cassandra system_schema; portable LIKE/fuzzy paging is not safe across versions, common constraints filter locally. Uses cassandra-gocql-driver schema metadata, then applies stable filtering and paging in the native agent.
5 dameng native-pushdown java-sql Uses ALL_OBJECTS with type, filter, stable order, and LIMIT/OFFSET with legacy fallback.
6 databend native-pushdown java-sql Uses system procedures for routine pushdown and shared table metadata constraints for table-like objects.
7 db2 native-pushdown java-sql Uses SYSCAT metadata with type, filter, stable order, and OFFSET/FETCH.

View File

@ -24,6 +24,8 @@ class DriverReleasePackagesTest(unittest.TestCase):
duckdb_source.write_bytes(b"\xcf\xfa\xed\xfetest-duckdb-agent")
rabbitmq_source = release_dir / "dbx-agent-rabbitmq-linux-x64"
rabbitmq_source.write_bytes(b"\x7fELFtest-rabbitmq-agent")
cassandra_source = release_dir / "dbx-agent-cassandra-linux-x64"
cassandra_source.write_bytes(b"\x7fELFtest-cassandra-agent")
java_source = release_dir / "dbx-agent-h2.jar"
java_source.write_bytes(b"test-jar")
versions = {
@ -34,6 +36,7 @@ class DriverReleasePackagesTest(unittest.TestCase):
"vastbase": "0.1.37",
"duckdb": "0.1.0",
"rabbitmq": "0.1.0",
"cassandra": "0.1.37",
}
renamed = version_agent_artifacts(release_dir, versions)
@ -42,7 +45,8 @@ class DriverReleasePackagesTest(unittest.TestCase):
versioned_vastbase = release_dir / "dbx-agent-vastbase-0.1.37-linux-x64"
versioned_duckdb = release_dir / "dbx-agent-duckdb-0.1.0-macos-aarch64"
versioned_rabbitmq = release_dir / "dbx-agent-rabbitmq-0.1.0-linux-x64"
self.assertEqual(renamed, [versioned_java, versioned_native, versioned_vastbase, versioned_duckdb, versioned_rabbitmq])
versioned_cassandra = release_dir / "dbx-agent-cassandra-0.1.37-linux-x64"
self.assertEqual(renamed, [versioned_java, versioned_cassandra, versioned_native, versioned_vastbase, versioned_duckdb, versioned_rabbitmq])
registry = {
"jres": {"21": {"version": "21", "platforms": {}}},
@ -54,6 +58,19 @@ class DriverReleasePackagesTest(unittest.TestCase):
"jre": "21",
"jar": {"url": f"https://example.com/{versioned_java.name}", "size": versioned_java.stat().st_size},
},
"cassandra": {
"version": "0.1.37",
"label": "Apache Cassandra",
"min_app_version": "0.6.0",
"jre": "21",
"jar": {"url": "https://example.com/legacy-placeholder.jar", "size": 0},
"native": {
"linux-x64": {
"url": f"https://example.com/{versioned_cassandra.name}",
"size": versioned_cassandra.stat().st_size,
}
},
},
"kingbase": {
"version": "0.1.34",
"label": "人大金仓 KingbaseES",
@ -116,6 +133,7 @@ class DriverReleasePackagesTest(unittest.TestCase):
outputs,
[
release_dir / "dbx-agent-h2-0.2.5.tar.zst",
release_dir / "dbx-agent-cassandra-0.1.37-linux-x64.tar.zst",
release_dir / "dbx-agent-kingbase-0.1.34-windows-x64.tar.zst",
release_dir / "dbx-agent-vastbase-0.1.37-linux-x64.tar.zst",
release_dir / "dbx-agent-duckdb-0.1.0-macos-aarch64.tar.zst",
@ -124,10 +142,11 @@ class DriverReleasePackagesTest(unittest.TestCase):
)
package_cases = [
(outputs[0], "h2", versioned_java, "jar", None),
(outputs[1], "kingbase", versioned_native, "native", "windows-x64"),
(outputs[2], "vastbase", versioned_vastbase, "native", "linux-x64"),
(outputs[3], "duckdb", versioned_duckdb, "native", "macos-aarch64"),
(outputs[4], "rabbitmq", versioned_rabbitmq, "native", "linux-x64"),
(outputs[1], "cassandra", versioned_cassandra, "native", "linux-x64"),
(outputs[2], "kingbase", versioned_native, "native", "windows-x64"),
(outputs[3], "vastbase", versioned_vastbase, "native", "linux-x64"),
(outputs[4], "duckdb", versioned_duckdb, "native", "macos-aarch64"),
(outputs[5], "rabbitmq", versioned_rabbitmq, "native", "linux-x64"),
]
for output, driver_name, source, artifact_type, platform in package_cases:
tar_bytes = subprocess.run(
@ -152,10 +171,11 @@ class DriverReleasePackagesTest(unittest.TestCase):
final_registry = json.loads((release_dir / "agent-registry.json").read_text(encoding="utf-8"))
release_artifacts = [
(final_registry["drivers"]["h2"]["jar"], outputs[0]),
(final_registry["drivers"]["kingbase"]["native"]["windows-x64"], outputs[1]),
(final_registry["drivers"]["vastbase"]["native"]["linux-x64"], outputs[2]),
(final_registry["drivers"]["duckdb"]["native"]["macos-aarch64"], outputs[3]),
(final_registry["drivers"]["rabbitmq"]["native"]["linux-x64"], outputs[4]),
(final_registry["drivers"]["cassandra"]["native"]["linux-x64"], outputs[1]),
(final_registry["drivers"]["kingbase"]["native"]["windows-x64"], outputs[2]),
(final_registry["drivers"]["vastbase"]["native"]["linux-x64"], outputs[3]),
(final_registry["drivers"]["duckdb"]["native"]["macos-aarch64"], outputs[4]),
(final_registry["drivers"]["rabbitmq"]["native"]["linux-x64"], outputs[5]),
]
for artifact, output in release_artifacts:
self.assertEqual(artifact["url"], f"https://example.com/{output.name}")
@ -164,7 +184,7 @@ class DriverReleasePackagesTest(unittest.TestCase):
self.assertEqual(len(artifact["sha256"]), 64)
removed = remove_raw_driver_artifacts(release_dir)
self.assertEqual(removed, [versioned_duckdb, versioned_java, versioned_native, versioned_rabbitmq, versioned_vastbase])
self.assertEqual(removed, [versioned_cassandra, versioned_duckdb, versioned_java, versioned_native, versioned_rabbitmq, versioned_vastbase])
self.assertTrue(all(output.is_file() for output in outputs))
def test_full_offline_bundle_includes_supported_windows_artifacts(self) -> None:

View File

@ -13,6 +13,7 @@ KOTLIN_SCAN_EXCLUDED_PARTS = {".git", ".gradle", "build"}
DEFAULT_AGENT_JRE_KEY = "21"
NON_JDBC_AGENT_MODULES = {"mongodb", "etcd", "zookeeper", "kafka", "rocketmq", "rabbitmq"}
NATIVE_ONLY_AGENT_MODULES = {
"cassandra": "drivers/cassandra-go",
"duckdb": "drivers/duckdb",
"oracle": "drivers/oracle-go",
"kingbase": "drivers/kingbase-go",

View File

@ -136,10 +136,10 @@ class ValidateAgentsTest(unittest.TestCase):
"include(*(infrastructureModules + driverModules))\n",
encoding="utf-8",
)
for driver in ("oracle-go", "kingbase-go", "vastbase-go", "xugu", "duckdb", "rabbitmq"):
for driver in ("cassandra-go", "oracle-go", "kingbase-go", "vastbase-go", "xugu", "duckdb", "rabbitmq"):
(root / "drivers" / driver).mkdir(parents=True)
(root / "versions.json").write_text(
json.dumps({"h2": "0.1.0", "oracle": "0.1.0", "kingbase": "0.1.0", "vastbase": "0.1.0", "xugu": "0.1.0", "rabbitmq": "0.1.0"}),
json.dumps({"h2": "0.1.0", "cassandra": "0.1.0", "oracle": "0.1.0", "kingbase": "0.1.0", "vastbase": "0.1.0", "xugu": "0.1.0", "rabbitmq": "0.1.0"}),
encoding="utf-8",
)

View File

@ -4,7 +4,7 @@ import json
from pathlib import Path
NATIVE_DRIVERS = ("oracle", "xugu", "kingbase", "vastbase", "duckdb", "rabbitmq")
NATIVE_DRIVERS = ("cassandra", "oracle", "xugu", "kingbase", "vastbase", "duckdb", "rabbitmq")
PLATFORMS = (
"macos-aarch64",
"macos-x64",