feat(cassandra): add native Go agent
This commit is contained in:
parent
a6f1a9a871
commit
5c2dfde334
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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" },
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ const LABEL_PALETTE = [
|
|||
];
|
||||
|
||||
const DRIVER_DATABASE_ALIASES = {
|
||||
"cassandra-go": "cassandra",
|
||||
gbase8a: "gbase",
|
||||
gbase8s: "gbase",
|
||||
"h2-legacy": "h2",
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
);
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo
|
|||
| db2 | IBM DB2 | DB2 JDBC |
|
||||
| informix | IBM Informix | Informix JDBC |
|
||||
| neo4j | Neo4j | Neo4j JDBC |
|
||||
| cassandra | Apache Cassandra | Cassandra JDBC |
|
||||
| cassandra | Apache Cassandra 2.1+ | Apache cassandra-gocql-driver native agent |
|
||||
| bigquery | Google BigQuery | BigQuery JDBC |
|
||||
| kylin | Apache Kylin | Kylin JDBC |
|
||||
| sundb | SunDB | SunDB JDBC |
|
||||
|
|
@ -49,7 +49,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo
|
|||
|
||||
## Multi-JRE Support
|
||||
|
||||
Most Java agents target JRE 21. Native agents, such as `duckdb`, `oracle`, `kingbase`, `xugu`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents.
|
||||
Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `oracle`, `kingbase`, `xugu`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents.
|
||||
|
||||
## JDBC Connection Pooling
|
||||
|
||||
|
|
@ -75,7 +75,7 @@ Set `DBX_AGENT_JDBC_POOL_ENABLED=false` for a runtime-level compatibility fallba
|
|||
|
||||
For new agents, prefer a **native (Go or Rust) driver** over a Java/JDBC agent whenever a mature, license-compatible native driver is available. Native agents ship as a single self-contained executable with no JRE, which significantly reduces memory footprint and startup time — the JVM baseline that every Java agent pays even when idle is avoided entirely.
|
||||
|
||||
- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/duckdb`, `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/xugu`, and `drivers/rabbitmq` (amqp091-go) as reference implementations. No JRE download or management is needed.
|
||||
- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/cassandra-go` (Apache cassandra-gocql-driver), `drivers/duckdb`, `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/xugu`, and `drivers/rabbitmq` (amqp091-go) as reference implementations. No JRE download or management is needed.
|
||||
- **Java/JDBC** — the default fallback when only a JDBC driver exists for the database, or when the native driver is immature or unmaintained. Most agents still fall in this category.
|
||||
|
||||
Native agents implement the same JSON-RPC contract and `versions.json` registration as Java agents; they ship an `agent` executable instead of `agent.jar`. If both native and Java source implementations exist for the same database, publish only the native artifact unless the Java variant has a separately registered compatibility profile, such as `oracle-legacy` / `oracle-10g`.
|
||||
|
|
@ -87,13 +87,14 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed).
|
|||
```bash
|
||||
./gradlew shadowJar
|
||||
(cd drivers/oracle-go && go build -o agent .)
|
||||
(cd drivers/cassandra-go && go build -o agent .)
|
||||
(cd drivers/kingbase-go && go build -o agent .)
|
||||
(cd drivers/vastbase-go && go build -o agent .)
|
||||
(cd drivers/xugu && go build -o agent .)
|
||||
(cd drivers/rabbitmq && go build -o agent .)
|
||||
```
|
||||
|
||||
Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/xugu`, and `drivers/rabbitmq`.
|
||||
Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/xugu`, and `drivers/rabbitmq`.
|
||||
|
||||
### Local DBX Runtime Test
|
||||
|
||||
|
|
@ -107,7 +108,7 @@ cp agents/drivers/<db_type>/build/libs/*-all.jar ~/.dbx/agents/drivers/<db_type>
|
|||
|
||||
Restart DBX or disconnect and reconnect the database so the new agent process loads the replacement JAR.
|
||||
|
||||
Native agents such as `oracle`, `kingbase`, `xugu`, and `rabbitmq` use the `agent` executable in the driver directory instead of `agent.jar`.
|
||||
Native agents such as `cassandra`, `oracle`, `kingbase`, `xugu`, and `rabbitmq` use the `agent` executable in the driver directory instead of `agent.jar`.
|
||||
|
||||
## Versioning
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<db_type>/build/libs/*-all.jar ~/.dbx/agents/drivers/<db_type>
|
|||
|
||||
重启 DBX 或断开重连数据库,使新 agent 进程加载替换后的 JAR。
|
||||
|
||||
`oracle`、`kingbase`、`xugu` 和 `rabbitmq` 等原生 agent 使用驱动目录下的 `agent` 可执行文件而非 `agent.jar`。
|
||||
`cassandra`、`oracle`、`kingbase`、`xugu` 和 `rabbitmq` 等原生 agent 使用驱动目录下的 `agent` 可执行文件而非 `agent.jar`。
|
||||
|
||||
## 版本管理
|
||||
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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)))
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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=
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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"
|
||||
}
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"net"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
gocql "github.com/apache/cassandra-gocql-driver/v2"
|
||||
)
|
||||
|
||||
func TestNormalizeCQLValuePreservesLegacyStringContract(t *testing.T) {
|
||||
uuid, err := gocql.ParseUUID("00112233-4455-6677-8899-aabbccddeeff")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tests := []struct {
|
||||
value any
|
||||
want any
|
||||
}{
|
||||
{nil, nil},
|
||||
{42, "42"},
|
||||
{true, "true"},
|
||||
{[]byte{0x00, 0xff}, "0x00ff"},
|
||||
{uuid, "00112233-4455-6677-8899-aabbccddeeff"},
|
||||
{net.ParseIP("127.0.0.1"), "127.0.0.1"},
|
||||
{time.Date(2026, 8, 3, 12, 34, 56, 7, time.UTC), "2026-08-03T12:34:56.000000007Z"},
|
||||
{gocql.Duration{Months: 1, Days: 2, Nanoseconds: 3}, "1mo2d3ns"},
|
||||
{[]int{1, 2}, "[1, 2]"},
|
||||
{[]any{1, nil, "three"}, "[1, null, three]"},
|
||||
{map[string]int{"b": 2, "a": 1}, "{a=1, b=2}"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := normalizeCQLValue(test.value); !reflect.DeepEqual(got, test.want) {
|
||||
t.Fatalf("normalizeCQLValue(%#v) = %#v, want %#v", test.value, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCQLTypeNameUsesCQLSyntax(t *testing.T) {
|
||||
typeInfo := gocql.NewNativeType(4, gocql.TypeList, "varchar")
|
||||
if got := cqlTypeName(typeInfo); got != "list<text>" {
|
||||
t.Fatalf("unexpected collection type name: %s", got)
|
||||
}
|
||||
}
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Reference in New Issue