diff --git a/.github/scripts/bump-agent-versions.mjs b/.github/scripts/bump-agent-versions.mjs index 94993a949..b3f92a087 100644 --- a/.github/scripts/bump-agent-versions.mjs +++ b/.github/scripts/bump-agent-versions.mjs @@ -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; diff --git a/.github/scripts/bump-agent-versions.test.mjs b/.github/scripts/bump-agent-versions.test.mjs index 1f61f0c3d..5d58964d5 100644 --- a/.github/scripts/bump-agent-versions.test.mjs +++ b/.github/scripts/bump-agent-versions.test.mjs @@ -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" }, diff --git a/.github/scripts/label-pull-request.mjs b/.github/scripts/label-pull-request.mjs index f903a7ead..36b1e7baa 100644 --- a/.github/scripts/label-pull-request.mjs +++ b/.github/scripts/label-pull-request.mjs @@ -47,6 +47,7 @@ const LABEL_PALETTE = [ ]; const DRIVER_DATABASE_ALIASES = { + "cassandra-go": "cassandra", gbase8a: "gbase", gbase8s: "gbase", "h2-legacy": "h2", diff --git a/.github/scripts/label-pull-request.test.mjs b/.github/scripts/label-pull-request.test.mjs index ed89fb522..1737ae8a0 100644 --- a/.github/scripts/label-pull-request.test.mjs +++ b/.github/scripts/label-pull-request.test.mjs @@ -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"], ); }); diff --git a/.github/workflows/agents-release.yml b/.github/workflows/agents-release.yml index 3e987555a..6d2188ff0 100644 --- a/.github/workflows/agents-release.yml +++ b/.github/workflows/agents-release.yml @@ -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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fd06d1b01..3e5c8e11a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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 diff --git a/agents/README.md b/agents/README.md index f4fe28175..46d66c569 100644 --- a/agents/README.md +++ b/agents/README.md @@ -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//build/libs/*-all.jar ~/.dbx/agents/drivers/ 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 diff --git a/agents/README.zh-CN.md b/agents/README.zh-CN.md index fff0b60b8..d0990a585 100644 --- a/agents/README.zh-CN.md +++ b/agents/README.zh-CN.md @@ -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 agent，DBX 会自动下载并管理 JRE 21 安装。 +多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`oracle`、`kingbase`、`xugu` 和 `rabbitmq`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 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//build/libs/*-all.jar ~/.dbx/agents/drivers/ 重启 DBX 或断开重连数据库，使新 agent 进程加载替换后的 JAR。 -`oracle`、`kingbase`、`xugu` 和 `rabbitmq` 等原生 agent 使用驱动目录下的 `agent` 可执行文件而非 `agent.jar`。 +`cassandra`、`oracle`、`kingbase`、`xugu` 和 `rabbitmq` 等原生 agent 使用驱动目录下的 `agent` 可执行文件而非 `agent.jar`。 ## 版本管理 diff --git a/agents/build.gradle b/agents/build.gradle index fe86b937d..f26d57a1f 100644 --- a/agents/build.gradle +++ b/agents/build.gradle @@ -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', diff --git a/agents/drivers/cassandra-go/config.go b/agents/drivers/cassandra-go/config.go new file mode 100644 index 000000000..39b829fa9 --- /dev/null +++ b/agents/drivers/cassandra-go/config.go @@ -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))) +} diff --git a/agents/drivers/cassandra-go/config_test.go b/agents/drivers/cassandra-go/config_test.go new file mode 100644 index 000000000..e3136041e --- /dev/null +++ b/agents/drivers/cassandra-go/config_test.go @@ -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) + } +} diff --git a/agents/drivers/cassandra-go/go.mod b/agents/drivers/cassandra-go/go.mod new file mode 100644 index 000000000..2da3d187a --- /dev/null +++ b/agents/drivers/cassandra-go/go.mod @@ -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 diff --git a/agents/drivers/cassandra-go/go.sum b/agents/drivers/cassandra-go/go.sum new file mode 100644 index 000000000..d5efef166 --- /dev/null +++ b/agents/drivers/cassandra-go/go.sum @@ -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= diff --git a/agents/drivers/cassandra-go/main.go b/agents/drivers/cassandra-go/main.go new file mode 100644 index 000000000..0cd3c06ca --- /dev/null +++ b/agents/drivers/cassandra-go/main.go @@ -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 + } +} diff --git a/agents/drivers/cassandra-go/metadata.go b/agents/drivers/cassandra-go/metadata.go new file mode 100644 index 000000000..b30387bd6 --- /dev/null +++ b/agents/drivers/cassandra-go/metadata.go @@ -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 +} diff --git a/agents/drivers/cassandra-go/metadata_test.go b/agents/drivers/cassandra-go/metadata_test.go new file mode 100644 index 000000000..263b7279a --- /dev/null +++ b/agents/drivers/cassandra-go/metadata_test.go @@ -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) + } +} diff --git a/agents/drivers/cassandra-go/protocol_error.go b/agents/drivers/cassandra-go/protocol_error.go new file mode 100644 index 000000000..ae7b4ec5b --- /dev/null +++ b/agents/drivers/cassandra-go/protocol_error.go @@ -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() +} diff --git a/agents/drivers/cassandra-go/protocol_test.go b/agents/drivers/cassandra-go/protocol_test.go new file mode 100644 index 000000000..c4d691b65 --- /dev/null +++ b/agents/drivers/cassandra-go/protocol_test.go @@ -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) + } +} diff --git a/agents/drivers/cassandra-go/query.go b/agents/drivers/cassandra-go/query.go new file mode 100644 index 000000000..521e759ca --- /dev/null +++ b/agents/drivers/cassandra-go/query.go @@ -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 +} diff --git a/agents/drivers/cassandra-go/runtime.go b/agents/drivers/cassandra-go/runtime.go new file mode 100644 index 000000000..1021a6b2d --- /dev/null +++ b/agents/drivers/cassandra-go/runtime.go @@ -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 +} diff --git a/agents/drivers/cassandra-go/values.go b/agents/drivers/cassandra-go/values.go new file mode 100644 index 000000000..14afb36b5 --- /dev/null +++ b/agents/drivers/cassandra-go/values.go @@ -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" +} diff --git a/agents/drivers/cassandra-go/values_test.go b/agents/drivers/cassandra-go/values_test.go new file mode 100644 index 000000000..226121f62 --- /dev/null +++ b/agents/drivers/cassandra-go/values_test.go @@ -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" { + t.Fatalf("unexpected collection type name: %s", got) + } +} diff --git a/agents/metadata-constraint-coverage.tsv b/agents/metadata-constraint-coverage.tsv index b2120630a..8b9946ee9 100644 --- a/agents/metadata-constraint-coverage.tsv +++ b/agents/metadata-constraint-coverage.tsv @@ -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. diff --git a/agents/scripts/driver_release_packages_test.py b/agents/scripts/driver_release_packages_test.py index b96fcef03..51d5185d9 100644 --- a/agents/scripts/driver_release_packages_test.py +++ b/agents/scripts/driver_release_packages_test.py @@ -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: diff --git a/agents/scripts/validate_agents.py b/agents/scripts/validate_agents.py index fa7526f97..4d713a43f 100644 --- a/agents/scripts/validate_agents.py +++ b/agents/scripts/validate_agents.py @@ -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", diff --git a/agents/scripts/validate_agents_test.py b/agents/scripts/validate_agents_test.py index e946d8db3..318287fc6 100644 --- a/agents/scripts/validate_agents_test.py +++ b/agents/scripts/validate_agents_test.py @@ -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", ) diff --git a/agents/scripts/version_agent_artifacts.py b/agents/scripts/version_agent_artifacts.py index e1e8b1a60..6831a66a3 100644 --- a/agents/scripts/version_agent_artifacts.py +++ b/agents/scripts/version_agent_artifacts.py @@ -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",