feat(cassandra): complete native agent migration

This commit is contained in:
t8y2 2026-08-04 01:21:17 +08:00
parent 8c584590ef
commit 44ef416d2f
No known key found for this signature in database
17 changed files with 1299 additions and 356 deletions

View File

@ -566,6 +566,47 @@ jobs:
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
- name: Cassandra native agent integration tests
shell: bash
working-directory: agents/drivers/cassandra-go
run: |
set -euo pipefail
for version in 3.11.19 5.0.6; do
name="dbx-cassandra-${version//./-}"
docker rm -fv "$name" >/dev/null 2>&1 || true
docker run -d --name "$name" \
-e CASSANDRA_CLUSTER_NAME="DBX Cassandra CI $version" \
-e CASSANDRA_DC=dc1 \
-e CASSANDRA_RACK=rack1 \
-e CASSANDRA_ENDPOINT_SNITCH=GossipingPropertyFileSnitch \
-e CASSANDRA_NUM_TOKENS=16 \
-e MAX_HEAP_SIZE=512M \
-e HEAP_NEWSIZE=100M \
-p 9042:9042 \
"cassandra:$version"
cleanup() {
docker rm -fv "$name" >/dev/null 2>&1 || true
}
trap cleanup EXIT
ready=false
for _ in $(seq 1 100); do
if docker exec "$name" cqlsh -e 'SELECT release_version FROM system.local' >/dev/null 2>&1; then
ready=true
break
fi
sleep 3
done
if [ "$ready" != "true" ]; then
docker logs "$name"
exit 1
fi
CASSANDRA_TEST_HOST=127.0.0.1 \
CASSANDRA_TEST_PORT=9042 \
go test -run '^TestCassandraIntegration$' -count=1 ./...
cleanup
trap - EXIT
done
- name: RabbitMQ native agent integration tests
shell: bash
working-directory: agents/drivers/rabbitmq

View File

@ -0,0 +1,61 @@
# Cassandra native Agent
The Cassandra Agent uses Apache `cassandra-gocql-driver` and implements the DBX
multi-session JSON-RPC protocol without a JVM.
## Compatibility
- Native protocol versions: v3-v5
- Declared server range: Apache Cassandra 2.1+
- Live validation: 2.2.19, 3.11.19, 4.1.10, and 5.0.6
- Authentication: username/password
- TLS: CA verification, optional client certificate/key, hostname verification
- Metadata: keyspaces, tables, columns, indexes, CQL table DDL, completion search
- Queries: legacy string result values, paging, cancellation, logged and unlogged batches
The Agent accepts both normal DBX connection fields and Cassandra JDBC-style
connection strings, including the wrapper's `host1--host2:9042` contact-point
syntax.
## JDBC URL parameter mapping
| JDBC parameter | Native behavior |
| --- | --- |
| `consistency` | GoCQL consistency |
| `fetchsize` | default page size |
| `retries` | retry/reconnection attempt count |
| `loadbalancing` | default, round-robin, DC-aware, or token-aware built-in policy |
| `localdatacenter` | DC-aware host selection |
| `retry` | default/simple, fallthrough, downgrading, or exponential built-in policy |
| `reconnection` | constant or exponential reconnection policy |
| `debug` | GoCQL debug logging to stderr |
| `enablessl` | TLS enablement |
| `sslenginefactory` | the standard `DefaultSslEngineFactory` maps to native TLS |
| `hostnameverification` | TLS hostname verification; enabled by default |
| `user`, `password` | password authentication |
| `requesttimeout`, `connecttimeout` | request and connection deadlines |
| `tcpnodelay`, `keepalive` | native TCP socket options |
| `compliancemode` | accepted; JDBC-only `java.sql` behavior is not applicable to JSON-RPC |
Java implementation hooks do not have a safe native equivalent. The Agent
returns a targeted connection error for `configfile`, `usekrb5=true`,
`secureconnectbundle`, custom `sslenginefactory` classes, and custom policy
classes. Translate Java HOCON settings to the supported URL parameters before
migrating a connection.
## Integration test
```bash
CASSANDRA_TEST_HOST=127.0.0.1 \
CASSANDRA_TEST_PORT=9042 \
CASSANDRA_TEST_USERNAME=cassandra \
CASSANDRA_TEST_PASSWORD=cassandra \
go test -run TestCassandraIntegration -v
```
Optional variables include `CASSANDRA_TEST_URL_PARAMS`, `CASSANDRA_TEST_SSL`,
`CASSANDRA_TEST_CA_CERT_PATH`, `CASSANDRA_TEST_CLIENT_CERT_PATH`, and
`CASSANDRA_TEST_CLIENT_KEY_PATH`.
See `bench/README.md` for the archived JDBC comparison workflow and measured
Cassandra 4.1.10 results.

View File

@ -0,0 +1,61 @@
# Cassandra Agent benchmark
This benchmark compares the same DBX JSON-RPC operations through the native
Apache `cassandra-gocql-driver` Agent and the archived Cassandra JDBC Agent.
It measures process startup, connection creation, RSS, latency, throughput,
artifact size, and shutdown behavior.
Each connection sample uses a fresh Agent process so JDBC runtime pooling cannot
turn later samples into warm reconnects. Query workloads use one persistent,
already-connected process per candidate.
## Prepare the fixture
The default workload expects `dbx_native_test.all_types` with at least 100 rows
and an integer primary key named `id`. Override the SQL variables below when
using another schema.
## Build the native Agent
From `agents/`:
```bash
go build -o /tmp/dbx-cassandra-bench/cassandra-go ./drivers/cassandra-go
```
Keep an archived JDBC Agent JAR as the baseline. The production Cassandra
module publishes only the native executable.
## Run
```bash
GO_AGENT=/tmp/dbx-cassandra-bench/cassandra-go \
JDBC_AGENT_JAR=/tmp/dbx-cassandra-bench/dbx-agent-cassandra.jar \
CASSANDRA_HOST=127.0.0.1 \
CASSANDRA_PORT=9042 \
CASSANDRA_KEYSPACE=dbx_native_test \
python3 drivers/cassandra-go/bench/agent_compare.py \
> /tmp/dbx-cassandra-bench/result.json
```
If Java is only available in a container, provide the full interactive command:
```bash
JDBC_AGENT_COMMAND='docker run --rm -i --name dbx-cassandra-jdbc-bench -v /tmp/dbx-cassandra-bench:/bench:ro eclipse-temurin:21-jre java -jar /bench/dbx-agent-cassandra.jar'
JDBC_RSS_COMMAND="docker inspect --format '{{.State.Pid}}' dbx-cassandra-jdbc-bench | xargs -I{} awk '/VmRSS/ {print \$2}' /proc/{}/status"
```
## Configuration
- `BENCH_CANDIDATES`: `go,jdbc` by default
- `BENCH_STARTUPS`: startup samples, default `10`
- `BENCH_CONNECTS`: connection samples, default `10`
- `BENCH_WARMUPS`: warmups before each workload, default `20`
- `CASSANDRA_USERNAME`, `CASSANDRA_PASSWORD`, `CASSANDRA_URL_PARAMS`
- `CASSANDRA_SSL`, `CASSANDRA_CA_CERT_PATH`, `CASSANDRA_CLIENT_CERT_PATH`, `CASSANDRA_CLIENT_KEY_PATH`
- `BENCH_SELECT_ONE_SQL`, `BENCH_DECODE_SQL`, `BENCH_PAGE_SQL`
- `BENCH_SELECT_ONE_COUNT`, `BENCH_DECODE_COUNT`, `BENCH_LIST_TABLES_COUNT`, `BENCH_PAGE_COUNT`
Run both candidates on the same host against the same Cassandra instance. Do
not compare a local native Agent with a remote JDBC Agent or change the query
shape between candidates.

View File

@ -0,0 +1,328 @@
#!/usr/bin/env python3
import json
import os
import shlex
import statistics
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path
@dataclass(frozen=True)
class Candidate:
name: str
command: list[str]
artifact: Path
rss_command: str = ""
class AgentProcess:
def __init__(self, candidate: Candidate):
self.candidate = candidate
self.process = subprocess.Popen(
candidate.command,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
bufsize=1,
)
self.request_id = 0
self.stderr_lines: list[str] = []
threading.Thread(target=self._drain_stderr, daemon=True).start()
self._wait_ready()
def _drain_stderr(self) -> None:
assert self.process.stderr is not None
for line in self.process.stderr:
self.stderr_lines.append(line.rstrip())
def _wait_ready(self) -> None:
assert self.process.stdout is not None
deadline = time.monotonic() + env_float("BENCH_READY_TIMEOUT", 30.0)
while time.monotonic() < deadline:
line = self.process.stdout.readline()
if line == "" and self.process.poll() is not None:
raise RuntimeError(self._failure("agent exited before ready"))
try:
if json.loads(line).get("ready") is True:
return
except (json.JSONDecodeError, AttributeError):
continue
raise TimeoutError(self._failure("timed out waiting for agent readiness"))
def call(self, method: str, params: dict | None = None) -> dict:
self.request_id += 1
request = {
"jsonrpc": "2.0",
"id": self.request_id,
"method": method,
"params": params or {},
}
assert self.process.stdin is not None
assert self.process.stdout is not None
self.process.stdin.write(json.dumps(request, separators=(",", ":")) + "\n")
self.process.stdin.flush()
while True:
line = self.process.stdout.readline()
if line == "" and self.process.poll() is not None:
raise RuntimeError(self._failure(f"agent exited during {method}"))
try:
response = json.loads(line)
except json.JSONDecodeError:
continue
if response.get("id") != self.request_id:
continue
if response.get("error") is not None:
raise RuntimeError(f"{self.candidate.name} {method}: {json.dumps(response['error'], ensure_ascii=False)}")
return response.get("result")
def rss_kib(self) -> int:
if self.candidate.rss_command:
output = subprocess.check_output(self.candidate.rss_command, shell=True, text=True).strip()
return int(output)
output = subprocess.check_output(
["ps", "-o", "rss=", "-p", str(self.process.pid)],
text=True,
).strip()
return int(output or "0")
def close(self) -> bool:
if self.process.poll() is not None:
return True
try:
self.call("shutdown")
except Exception:
pass
try:
self.process.wait(timeout=3)
return True
except subprocess.TimeoutExpired:
self.process.terminate()
try:
self.process.wait(timeout=2)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait(timeout=5)
return False
def _failure(self, message: str) -> str:
stderr = "\n".join(self.stderr_lines[-20:])
return f"{self.candidate.name}: {message}\n{stderr}".rstrip()
def main() -> None:
candidates = configured_candidates()
connection = connection_params()
startup_iterations = env_int("BENCH_STARTUPS", 10)
connect_iterations = env_int("BENCH_CONNECTS", 10)
warmups = env_int("BENCH_WARMUPS", 20)
workloads = configured_workloads(connection["database"])
results = []
for candidate in candidates:
startup_samples = benchmark_startup(candidate, startup_iterations)
connect_samples = benchmark_connect(candidate, connection, connect_iterations)
process = AgentProcess(candidate)
shutdown_clean = False
try:
process.call("connect", connection)
rss_kib = process.rss_kib()
workload_results = [benchmark_workload(process, workload, warmups) for workload in workloads]
process.call("disconnect")
finally:
shutdown_clean = process.close()
results.append(
{
"candidate": candidate.name,
"command": candidate.command,
"artifact_bytes": candidate.artifact.stat().st_size,
"startup_ms": statistics.median(startup_samples),
"startup_samples_ms": startup_samples,
"connect_ms": statistics.median(connect_samples),
"connect_samples_ms": connect_samples,
"rss_kib": rss_kib,
"shutdown_exited_within_3s": shutdown_clean,
"workloads": workload_results,
}
)
output = {
"host": os.uname().nodename,
"server": env_default("CASSANDRA_SERVER", f"{connection['host']}:{connection['port']}"),
"keyspace": connection["database"],
"startup_iterations": startup_iterations,
"connect_iterations": connect_iterations,
"warmups": warmups,
"results": results,
}
json.dump(output, sys.stdout, ensure_ascii=False, indent=2)
sys.stdout.write("\n")
def configured_candidates() -> list[Candidate]:
selected = {item.strip() for item in env_default("BENCH_CANDIDATES", "go,jdbc").split(",") if item.strip()}
candidates = []
if "go" in selected:
artifact = required_path("GO_AGENT")
candidates.append(Candidate("go-native", [str(artifact)], artifact, os.getenv("GO_RSS_COMMAND", "")))
if "jdbc" in selected:
artifact = required_path("JDBC_AGENT_JAR")
raw_command = os.getenv("JDBC_AGENT_COMMAND", "")
command = shlex.split(raw_command) if raw_command else [env_default("JAVA_BIN", "java"), "-jar", str(artifact)]
candidates.append(Candidate("jdbc-java", command, artifact, os.getenv("JDBC_RSS_COMMAND", "")))
if not candidates:
raise ValueError("BENCH_CANDIDATES selected no candidates")
return candidates
def connection_params() -> dict:
return {
"host": env_default("CASSANDRA_HOST", "127.0.0.1"),
"port": env_int("CASSANDRA_PORT", 9042),
"database": env_default("CASSANDRA_KEYSPACE", "dbx_native_test"),
"username": os.getenv("CASSANDRA_USERNAME", ""),
"password": os.getenv("CASSANDRA_PASSWORD", ""),
"url_params": os.getenv("CASSANDRA_URL_PARAMS", ""),
"connection_string": os.getenv("CASSANDRA_CONNECTION_STRING", ""),
"ssl": env_bool("CASSANDRA_SSL", False),
"ca_cert_path": os.getenv("CASSANDRA_CA_CERT_PATH", ""),
"client_cert_path": os.getenv("CASSANDRA_CLIENT_CERT_PATH", ""),
"client_key_path": os.getenv("CASSANDRA_CLIENT_KEY_PATH", ""),
}
def configured_workloads(keyspace: str) -> list[dict]:
table = env_default("CASSANDRA_BENCH_TABLE", "all_types")
qualified = f'"{keyspace}"."{table}"'
return [
{
"name": "select_one",
"method": "execute_query",
"params": {"sql": env_default("BENCH_SELECT_ONE_SQL", f"SELECT id, txt FROM {qualified} WHERE id = 1"), "schema": keyspace, "maxRows": 1},
"count": env_int("BENCH_SELECT_ONE_COUNT", 1000),
},
{
"name": "decode_all_types",
"method": "execute_query",
"params": {"sql": env_default("BENCH_DECODE_SQL", f"SELECT * FROM {qualified} WHERE id = 1"), "schema": keyspace, "maxRows": 1},
"count": env_int("BENCH_DECODE_COUNT", 500),
},
{
"name": "list_tables",
"method": "list_tables",
"params": {"schema": keyspace},
"count": env_int("BENCH_LIST_TABLES_COUNT", 500),
},
{
"name": "page_100",
"method": "execute_query_page",
"params": {
"sql": env_default("BENCH_PAGE_SQL", f"SELECT id, txt FROM {qualified}"),
"schema": keyspace,
"maxRows": 100,
"pageSize": 100,
},
"count": env_int("BENCH_PAGE_COUNT", 200),
},
]
def benchmark_startup(candidate: Candidate, iterations: int) -> list[float]:
samples = []
for _ in range(iterations):
start = time.perf_counter()
process = AgentProcess(candidate)
samples.append((time.perf_counter() - start) * 1000)
process.close()
return samples
def benchmark_connect(candidate: Candidate, connection: dict, iterations: int) -> list[float]:
samples = []
for _ in range(iterations):
process = AgentProcess(candidate)
try:
start = time.perf_counter()
process.call("connect", connection)
samples.append((time.perf_counter() - start) * 1000)
finally:
process.close()
return samples
def benchmark_workload(process: AgentProcess, workload: dict, warmups: int) -> dict:
for _ in range(warmups):
process.call(workload["method"], workload["params"])
samples = []
start = time.perf_counter()
for _ in range(workload["count"]):
operation_start = time.perf_counter()
process.call(workload["method"], workload["params"])
samples.append((time.perf_counter() - operation_start) * 1000)
elapsed = time.perf_counter() - start
ordered = sorted(samples)
return {
"name": workload["name"],
"count": workload["count"],
"elapsed_ms": elapsed * 1000,
"ops_per_sec": workload["count"] / elapsed,
"mean_ms": statistics.mean(samples),
"p50_ms": percentile(ordered, 0.50),
"p95_ms": percentile(ordered, 0.95),
"p99_ms": percentile(ordered, 0.99),
}
def percentile(values: list[float], fraction: float) -> float:
if not values:
return 0.0
index = min(len(values) - 1, max(0, round((len(values) - 1) * fraction)))
return values[index]
def required_path(name: str) -> Path:
value = os.getenv(name, "")
if not value:
raise ValueError(f"{name} is required")
path = Path(value).expanduser().resolve()
if not path.is_file():
raise FileNotFoundError(path)
return path
def env_default(name: str, fallback: str) -> str:
return os.getenv(name, "") or fallback
def env_int(name: str, fallback: int) -> int:
value = int(env_default(name, str(fallback)))
if value < 1:
raise ValueError(f"{name} must be positive")
return value
def env_float(name: str, fallback: float) -> float:
value = float(env_default(name, str(fallback)))
if value <= 0:
raise ValueError(f"{name} must be positive")
return value
def env_bool(name: str, fallback: bool) -> bool:
raw = os.getenv(name)
if raw is None or raw == "":
return fallback
normalized = raw.strip().lower()
if normalized in {"1", "true", "yes", "on"}:
return True
if normalized in {"0", "false", "no", "off"}:
return False
raise ValueError(f"{name} must be a boolean")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,36 @@
{
"date": "2026-08-03",
"host": "oss-rainyun-01",
"cassandra": "4.1.10",
"method": "five cold-process startup and connection samples; persistent process for query workloads",
"results": [
{
"candidate": "go-native",
"startup_ms": 8.994690957479179,
"connect_ms": 33.21323194541037,
"rss_kib": 10352,
"artifact_bytes": 6750370,
"shutdown_exited_within_3s": true,
"workloads": [
{"name": "select_one", "count": 1000, "elapsed_ms": 1940.8286979887635, "ops_per_sec": 515.2438239584345, "mean_ms": 1.9388069859705865, "p50_ms": 1.9103229278698564, "p95_ms": 2.221024944446981, "p99_ms": 2.3676720447838306},
{"name": "decode_all_types", "count": 500, "elapsed_ms": 1079.476205050014, "ops_per_sec": 463.18760678641746, "mean_ms": 2.156704908935353, "p50_ms": 2.141958102583885, "p95_ms": 2.4329390143975616, "p99_ms": 2.7680869679898024},
{"name": "list_tables", "count": 500, "elapsed_ms": 57.42691201157868, "ops_per_sec": 8706.71924513698, "mean_ms": 0.11419291398487985, "p50_ms": 0.10622991248965263, "p95_ms": 0.1487070694565773, "p99_ms": 0.19879091996699572},
{"name": "page_100", "count": 200, "elapsed_ms": 984.7569830017164, "ops_per_sec": 203.09579261917395, "mean_ms": 4.921203925041482, "p50_ms": 4.922428051941097, "p95_ms": 6.079918937757611, "p99_ms": 7.303814985789359}
]
},
{
"candidate": "jdbc-java",
"startup_ms": 657.9867920372635,
"connect_ms": 1778.3896300243214,
"rss_kib": 166232,
"artifact_bytes": 22984056,
"shutdown_exited_within_3s": false,
"workloads": [
{"name": "select_one", "count": 1000, "elapsed_ms": 3868.2536740088835, "ops_per_sec": 258.51458675502147, "mean_ms": 3.865854301955551, "p50_ms": 3.6270120181143284, "p95_ms": 5.733568919822574, "p99_ms": 7.344924029894173},
{"name": "decode_all_types", "count": 500, "elapsed_ms": 2014.9188039358705, "ops_per_sec": 248.14895718046694, "mean_ms": 4.0274618696421385, "p50_ms": 3.9998559514060616, "p95_ms": 4.94410190731287, "p99_ms": 5.3489640122279525},
{"name": "list_tables", "count": 500, "elapsed_ms": 763.6687039630488, "ops_per_sec": 654.7341764894339, "mean_ms": 1.5251681823283434, "p50_ms": 1.494601950980723, "p95_ms": 1.9315499812364578, "p99_ms": 2.339883940294385},
{"name": "page_100", "count": 200, "elapsed_ms": 8309.039836982265, "ops_per_sec": 24.070169829952025, "mean_ms": 41.53980694070924, "p50_ms": 42.0777719700709, "p95_ms": 51.87095201108605, "p99_ms": 57.71494994405657}
]
}
]
}

View File

@ -31,23 +31,37 @@ type cassandraConfig struct {
clientCertPath string
clientKeyPath string
hostVerification bool
tcpNoDelay bool
keepAlive bool
debug bool
retryPolicy string
retryCount int
reconnectionPolicy string
reconnectionBaseDelay time.Duration
reconnectionMaxDelay time.Duration
loadBalancingPolicy string
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,
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,
hostVerification: true,
tcpNoDelay: true,
retryCount: 3,
reconnectionBaseDelay: time.Second,
reconnectionMaxDelay: 60 * time.Second,
}
if cp.Port > 0 {
config.port = cp.Port
@ -101,14 +115,7 @@ func applyConnectionString(config *cassandraConfig, params url.Values, raw strin
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)
}
config.hosts = splitHosts(parsed.Host)
if port := parsed.Port(); port != "" {
parsedPort, parseErr := strconv.Atoi(port)
if parseErr != nil || parsedPort < 1 || parsedPort > 65535 {
@ -193,18 +200,60 @@ func applyCassandraURLParams(config *cassandraConfig, params url.Values) error {
config.pageSize = size
case "cqlversion":
config.cqlVersion = value
case "ssl":
case "ssl", "enablessl":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid ssl option: %w", err)
}
config.ssl = enabled
case "hostverification", "verifyhostname", "sslhostnameverification":
case "hostverification", "verifyhostname", "sslhostnameverification", "hostnameverification":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid host verification option: %w", err)
}
config.hostVerification = enabled
case "tcpnodelay":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid tcpnodelay option: %w", err)
}
config.tcpNoDelay = enabled
case "keepalive":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid keepalive option: %w", err)
}
config.keepAlive = enabled
case "user":
config.username = value
case "password":
config.password = value
case "debug":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid debug option: %w", err)
}
config.debug = enabled
case "retries":
count, err := strconv.Atoi(value)
if err != nil || count < 0 || count > 1000 {
return fmt.Errorf("retries must be between 0 and 1000")
}
config.retryCount = count
case "retry":
policy, err := normalizeRetryPolicy(value)
if err != nil {
return err
}
config.retryPolicy = policy
case "reconnection":
policy, baseDelay, maxDelay, err := parseReconnectionPolicy(value)
if err != nil {
return err
}
config.reconnectionPolicy = policy
config.reconnectionBaseDelay = baseDelay
config.reconnectionMaxDelay = maxDelay
case "disableinitialhostlookup":
disabled, err := strconv.ParseBool(value)
if err != nil {
@ -212,9 +261,35 @@ func applyCassandraURLParams(config *cassandraConfig, params url.Values) error {
}
config.disableInitialHostLookup = disabled
case "loadbalancing":
if value != "" && !strings.EqualFold(value, "DcInferringLoadBalancingPolicy") {
return fmt.Errorf("unsupported Cassandra loadbalancing policy: %s", value)
policy, err := normalizeLoadBalancingPolicy(value)
if err != nil {
return err
}
config.loadBalancingPolicy = policy
case "sslenginefactory":
if value != "" && !strings.EqualFold(simpleClassName(value), "DefaultSslEngineFactory") {
return fmt.Errorf("custom Cassandra sslenginefactory is not supported by the native agent: %s", value)
}
config.ssl = true
case "usekrb5":
enabled, err := strconv.ParseBool(value)
if err != nil {
return fmt.Errorf("invalid usekrb5 option: %w", err)
}
if enabled {
return fmt.Errorf("Cassandra Kerberos authentication is not supported by the native agent")
}
case "secureconnectbundle":
if value != "" {
return fmt.Errorf("Cassandra secure connect bundles are not supported by the native agent")
}
case "configfile":
if value != "" {
return fmt.Errorf("Cassandra Java driver configfile is not supported; translate it to native URL parameters")
}
case "compliancemode":
// JDBC compliance modes only alter java.sql behavior. The native DBX
// JSON-RPC contract already defines statement and transaction behavior.
default:
return fmt.Errorf("unsupported Cassandra URL parameter: %s", rawKey)
}
@ -231,6 +306,11 @@ func (config cassandraConfig) clusterConfig(keyspace string) (*gocql.ClusterConf
cluster.WriteTimeout = config.requestTimeout
cluster.NumConns = config.numConnections
cluster.PageSize = config.pageSize
cluster.Dialer = cassandraDialer{
timeout: config.connectTimeout,
keepAlive: config.keepAlive,
tcpNoDelay: config.tcpNoDelay,
}
cluster.DisableInitialHostLookup = config.disableInitialHostLookup
cluster.IgnorePeerAddr = config.disableInitialHostLookup
if config.protocolVersion != 0 {
@ -264,15 +344,20 @@ func (config cassandraConfig) clusterConfig(keyspace string) (*gocql.ClusterConf
EnableHostVerification: config.hostVerification,
}
}
if config.localDatacenter != "" {
cluster.PoolConfig.HostSelectionPolicy = gocql.TokenAwareHostPolicy(
gocql.DCAwareRoundRobinPolicy(config.localDatacenter),
)
if config.debug {
cluster.Logger = gocql.NewLogger(gocql.LogLevelDebug)
}
if err := applyRetryPolicies(cluster, config); err != nil {
return nil, err
}
if err := applyLoadBalancingPolicy(cluster, config); err != nil {
return nil, err
}
return cluster, nil
}
func splitHosts(raw string) []string {
raw = strings.ReplaceAll(raw, "--", ",")
parts := strings.FieldsFunc(raw, func(char rune) bool { return char == ',' || char == ';' })
hosts := make([]string, 0, len(parts))
for _, part := range parts {
@ -280,16 +365,14 @@ func splitHosts(raw string) []string {
if host == "" {
continue
}
if parsedHost, _, err := net.SplitHostPort(host); err == nil {
host = parsedHost
}
hosts = append(hosts, strings.Trim(host, "[]"))
hosts = append(hosts, host)
}
return hosts
}
func allLoopbackHosts(hosts []string) bool {
for _, host := range hosts {
host = hostNameOnly(host)
if strings.EqualFold(host, "localhost") {
continue
}
@ -301,6 +384,14 @@ func allLoopbackHosts(hosts []string) bool {
return len(hosts) > 0
}
func hostNameOnly(host string) string {
host = strings.TrimSpace(host)
if parsedHost, _, err := net.SplitHostPort(host); err == nil {
return parsedHost
}
return strings.Trim(host, "[]")
}
func parseDurationOption(value string) (time.Duration, error) {
if duration, err := time.ParseDuration(value); err == nil {
return duration, nil
@ -312,6 +403,158 @@ func parseDurationOption(value string) (time.Duration, error) {
return time.Duration(milliseconds) * time.Millisecond, nil
}
func normalizeRetryPolicy(value string) (string, error) {
name := strings.ToLower(simpleClassName(value))
switch name {
case "", "defaultretrypolicy", "simpleretrypolicy":
return "simple", nil
case "fallthroughretrypolicy":
return "fallthrough", nil
case "downgradingconsistencyretrypolicy":
return "downgrading", nil
case "exponentialbackoffretrypolicy":
return "exponential", nil
default:
return "", fmt.Errorf("unsupported Cassandra retry policy: %s", value)
}
}
func normalizeLoadBalancingPolicy(value string) (string, error) {
name := strings.ToLower(simpleClassName(value))
switch name {
case "", "dcinferringloadbalancingpolicy", "defaultloadbalancingpolicy":
return "default", nil
case "roundrobinpolicy":
return "round_robin", nil
case "dcawareroundrobinpolicy":
return "dc_aware", nil
case "tokenawarepolicy":
return "token_aware", nil
default:
return "", fmt.Errorf("unsupported Cassandra loadbalancing policy: %s", value)
}
}
func parseReconnectionPolicy(value string) (string, time.Duration, time.Duration, error) {
trimmed := strings.TrimSpace(value)
name := simpleClassName(trimmed)
parameters := ""
if open := strings.IndexByte(name, '('); open >= 0 {
parameters = strings.TrimSuffix(name[open+1:], ")")
name = name[:open]
}
policy := strings.ToLower(strings.TrimSpace(name))
baseDelay := time.Second
maxDelay := 60 * time.Second
if parameters != "" {
parts := strings.Split(parameters, ",")
for index, part := range parts {
part = strings.TrimSpace(strings.ReplaceAll(strings.ToLower(part), "(long)", ""))
seconds, err := strconv.Atoi(part)
if err != nil || seconds < 0 {
return "", 0, 0, fmt.Errorf("invalid Cassandra reconnection policy delay: %s", part)
}
if index == 0 {
baseDelay = time.Duration(seconds) * time.Second
} else if index == 1 {
maxDelay = time.Duration(seconds) * time.Second
} else {
return "", 0, 0, fmt.Errorf("too many Cassandra reconnection policy parameters")
}
}
}
switch policy {
case "", "constantreconnectionpolicy":
return "constant", baseDelay, baseDelay, nil
case "exponentialreconnectionpolicy":
return "exponential", baseDelay, maxDelay, nil
default:
return "", 0, 0, fmt.Errorf("unsupported Cassandra reconnection policy: %s", value)
}
}
func simpleClassName(value string) string {
value = strings.TrimSpace(value)
prefix := value
if open := strings.IndexByte(prefix, '('); open >= 0 {
prefix = prefix[:open]
}
if dot := strings.LastIndexByte(prefix, '.'); dot >= 0 {
return value[dot+1:]
}
return value
}
func applyRetryPolicies(cluster *gocql.ClusterConfig, config cassandraConfig) error {
switch config.retryPolicy {
case "":
case "simple":
cluster.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: config.retryCount}
case "fallthrough":
cluster.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: 0}
case "downgrading":
cluster.RetryPolicy = &gocql.DowngradingConsistencyRetryPolicy{}
case "exponential":
cluster.RetryPolicy = &gocql.ExponentialBackoffRetryPolicy{
NumRetries: config.retryCount,
Min: config.reconnectionBaseDelay,
Max: config.reconnectionMaxDelay,
}
default:
return fmt.Errorf("unsupported Cassandra retry policy: %s", config.retryPolicy)
}
if config.reconnectionPolicy != "" || config.retryCount != 3 {
switch config.reconnectionPolicy {
case "", "constant":
cluster.ReconnectionPolicy = &gocql.ConstantReconnectionPolicy{
MaxRetries: config.retryCount,
Interval: config.reconnectionBaseDelay,
}
case "exponential":
cluster.ReconnectionPolicy = &gocql.ExponentialReconnectionPolicy{
MaxRetries: config.retryCount,
InitialInterval: config.reconnectionBaseDelay,
MaxInterval: config.reconnectionMaxDelay,
}
default:
return fmt.Errorf("unsupported Cassandra reconnection policy: %s", config.reconnectionPolicy)
}
}
return nil
}
func applyLoadBalancingPolicy(cluster *gocql.ClusterConfig, config cassandraConfig) error {
policy := config.loadBalancingPolicy
if policy == "" {
policy = "default"
}
switch policy {
case "default":
if config.localDatacenter == "" {
return nil
}
cluster.PoolConfig.HostSelectionPolicy = gocql.TokenAwareHostPolicy(
gocql.DCAwareRoundRobinPolicy(config.localDatacenter),
)
case "round_robin":
cluster.PoolConfig.HostSelectionPolicy = gocql.RoundRobinHostPolicy()
case "dc_aware":
if config.localDatacenter == "" {
return fmt.Errorf("DCAwareRoundRobinPolicy requires localdatacenter")
}
cluster.PoolConfig.HostSelectionPolicy = gocql.DCAwareRoundRobinPolicy(config.localDatacenter)
case "token_aware":
fallback := gocql.RoundRobinHostPolicy()
if config.localDatacenter != "" {
fallback = gocql.DCAwareRoundRobinPolicy(config.localDatacenter)
}
cluster.PoolConfig.HostSelectionPolicy = gocql.TokenAwareHostPolicy(fallback)
default:
return fmt.Errorf("unsupported Cassandra loadbalancing policy: %s", policy)
}
return nil
}
func normalizeOptionName(value string) string {
return strings.NewReplacer("_", "", "-", "", ".", "").Replace(strings.ToLower(strings.TrimSpace(value)))
}

View File

@ -1,8 +1,11 @@
package main
import (
"reflect"
"testing"
"time"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
func TestParseCassandraConfigSupportsLegacyJDBCOptions(t *testing.T) {
@ -40,7 +43,7 @@ func TestParseCassandraConfigAcceptsConnectionString(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if len(config.hosts) != 1 || config.hosts[0] != "db.example.com" || config.port != 9142 {
if len(config.hosts) != 1 || config.hosts[0] != "db.example.com:9142" || config.port != 9142 {
t.Fatalf("unexpected endpoint: %#v", config)
}
if config.keyspace != "catalog" || config.username != "alice" || config.password != "secret" {
@ -51,6 +54,106 @@ func TestParseCassandraConfigAcceptsConnectionString(t *testing.T) {
}
}
func TestParseCassandraConfigCoversMappableJDBCWrapperOptions(t *testing.T) {
config, err := parseCassandraConfig(connectParams{
ConnectionString: "jdbc:cassandra://host1--host2:9142/catalog?" +
"user=query-user&password=query-secret&enablessl=true&hostnameverification=false&" +
"tcpnodelay=false&keepalive=true&debug=true&retries=7&retry=DefaultRetryPolicy&" +
"reconnection=ExponentialReconnectionPolicy((long)2,(long)30)&" +
"loadbalancing=TokenAwarePolicy&compliancemode=Liquibase",
})
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(config.hosts, []string{"host1", "host2:9142"}) || config.port != 9142 {
t.Fatalf("unexpected multi-host endpoint: hosts=%#v port=%d", config.hosts, config.port)
}
if config.username != "query-user" || config.password != "query-secret" {
t.Fatalf("unexpected query credentials: %#v", config)
}
if !config.ssl || config.hostVerification || config.tcpNoDelay || !config.keepAlive || !config.debug {
t.Fatalf("unexpected transport options: %#v", config)
}
if config.retryPolicy != "simple" || config.retryCount != 7 || config.reconnectionPolicy != "exponential" {
t.Fatalf("unexpected retry options: %#v", config)
}
if config.reconnectionBaseDelay != 2*time.Second || config.reconnectionMaxDelay != 30*time.Second {
t.Fatalf("unexpected reconnection delays: %#v", config)
}
if config.loadBalancingPolicy != "token_aware" {
t.Fatalf("unexpected load-balancing option: %#v", config)
}
cluster, err := config.clusterConfig(config.keyspace)
if err != nil {
t.Fatal(err)
}
dialer, ok := cluster.Dialer.(cassandraDialer)
if !ok || dialer.tcpNoDelay || !dialer.keepAlive {
t.Fatalf("unexpected socket dialer: %#v", cluster.Dialer)
}
retryPolicy, ok := cluster.RetryPolicy.(*gocql.SimpleRetryPolicy)
if !ok || retryPolicy.NumRetries != 7 {
t.Fatalf("unexpected query retry policy: %#v", cluster.RetryPolicy)
}
reconnectionPolicy, ok := cluster.ReconnectionPolicy.(*gocql.ExponentialReconnectionPolicy)
if !ok || reconnectionPolicy.MaxRetries != 7 || reconnectionPolicy.InitialInterval != 2*time.Second || reconnectionPolicy.MaxInterval != 30*time.Second {
t.Fatalf("unexpected reconnection policy: %#v", cluster.ReconnectionPolicy)
}
}
func TestParseCassandraConfigUsesSecureTransportDefaults(t *testing.T) {
config, err := parseCassandraConfig(connectParams{Host: "127.0.0.1:9042", SSL: true})
if err != nil {
t.Fatal(err)
}
if !config.hostVerification || !config.tcpNoDelay || config.keepAlive {
t.Fatalf("unexpected defaults: %#v", config)
}
if !config.disableInitialHostLookup {
t.Fatal("loopback host with explicit port must disable peer discovery")
}
}
func TestParseCassandraConfigAcceptsDefaultSSLEngineFactory(t *testing.T) {
config, err := parseCassandraConfig(connectParams{
Host: "localhost",
URLParams: "sslenginefactory=com.datastax.oss.driver.internal.core.ssl.DefaultSslEngineFactory&usekrb5=false",
})
if err != nil {
t.Fatal(err)
}
if !config.ssl {
t.Fatal("default SSL engine factory must enable TLS")
}
}
func TestParseCassandraConfigRejectsJavaOnlyJDBCOptions(t *testing.T) {
tests := []string{
"configfile=/tmp/application.conf",
"secureconnectbundle=/tmp/secure-connect.zip",
"usekrb5=true",
"sslenginefactory=example.CustomSslEngineFactory",
}
for _, urlParams := range tests {
if _, err := parseCassandraConfig(connectParams{Host: "localhost", URLParams: urlParams}); err == nil {
t.Fatalf("expected unsupported Java-only option error for %q", urlParams)
}
}
}
func TestParseReconnectionPolicySupportsFullyQualifiedClass(t *testing.T) {
policy, baseDelay, maxDelay, err := parseReconnectionPolicy(
"com.datastax.oss.driver.internal.core.connection.ExponentialReconnectionPolicy((long)1,(long)8)",
)
if err != nil {
t.Fatal(err)
}
if policy != "exponential" || baseDelay != time.Second || maxDelay != 8*time.Second {
t.Fatalf("unexpected policy: %s %s %s", policy, baseDelay, maxDelay)
}
}
func TestParseCassandraConfigRejectsUnsupportedLoadBalancingClass(t *testing.T) {
_, err := parseCassandraConfig(connectParams{
Host: "localhost",

View File

@ -0,0 +1,48 @@
package main
import (
"context"
"net"
"time"
)
const cassandraKeepAlivePeriod = 30 * time.Second
type cassandraDialer struct {
timeout time.Duration
keepAlive bool
tcpNoDelay bool
}
func (dialer cassandraDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
keepAlivePeriod := time.Duration(-1)
if dialer.keepAlive {
keepAlivePeriod = cassandraKeepAlivePeriod
}
connection, err := (&net.Dialer{
Timeout: dialer.timeout,
KeepAlive: keepAlivePeriod,
}).DialContext(ctx, network, address)
if err != nil {
return nil, err
}
tcpConnection, ok := connection.(*net.TCPConn)
if !ok {
return connection, nil
}
if err := tcpConnection.SetNoDelay(dialer.tcpNoDelay); err != nil {
connection.Close()
return nil, err
}
if err := tcpConnection.SetKeepAlive(dialer.keepAlive); err != nil {
connection.Close()
return nil, err
}
if dialer.keepAlive {
if err := tcpConnection.SetKeepAlivePeriod(cassandraKeepAlivePeriod); err != nil {
connection.Close()
return nil, err
}
}
return connection, nil
}

View File

@ -0,0 +1,206 @@
package main
import (
"encoding/json"
"fmt"
"os"
"strconv"
"strings"
"testing"
"time"
)
func TestCassandraIntegration(t *testing.T) {
host := strings.TrimSpace(os.Getenv("CASSANDRA_TEST_HOST"))
if host == "" {
t.Skip("Cassandra integration environment is not configured")
}
port := 9042
if rawPort := strings.TrimSpace(os.Getenv("CASSANDRA_TEST_PORT")); rawPort != "" {
parsedPort, err := strconv.Atoi(rawPort)
if err != nil {
t.Fatal(err)
}
port = parsedPort
}
ssl, err := strconv.ParseBool(envDefault("CASSANDRA_TEST_SSL", "false"))
if err != nil {
t.Fatal(err)
}
connection := connectParams{
Host: host,
Port: port,
Username: os.Getenv("CASSANDRA_TEST_USERNAME"),
Password: os.Getenv("CASSANDRA_TEST_PASSWORD"),
URLParams: os.Getenv("CASSANDRA_TEST_URL_PARAMS"),
SSL: ssl,
CACertPath: os.Getenv("CASSANDRA_TEST_CA_CERT_PATH"),
ClientCertPath: os.Getenv("CASSANDRA_TEST_CLIENT_CERT_PATH"),
ClientKeyPath: os.Getenv("CASSANDRA_TEST_CLIENT_KEY_PATH"),
}
runtime, err := newConnectionRuntime(connection)
if err != nil {
t.Fatal(err)
}
defer runtime.close()
server := newServer(runtime, connection)
if err := server.validateConnection(); err != nil {
t.Fatal(err)
}
suffix := strconv.FormatInt(time.Now().UnixNano(), 36)
keyspace := "dbx_native_it_" + suffix
table := "all_types"
pagedTable := "paged_rows"
mustCQL(t, server, "CREATE KEYSPACE "+quoteCQLIdentifier(keyspace)+" WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}", "")
t.Cleanup(func() {
_, _ = server.executeQuery(queryOptions{SQL: "DROP KEYSPACE IF EXISTS " + quoteCQLIdentifier(keyspace)})
})
mustCQL(t, server, "CREATE TABLE "+qualifiedCQLName(keyspace, table)+" ("+
"id int PRIMARY KEY, txt text, flag boolean, amount decimal, payload blob, created timestamp, address inet, "+
"tags set<text>, items list<int>, attrs map<text, int>, pair frozen<tuple<int, text>>)", keyspace)
mustCQL(t, server, "CREATE INDEX "+quoteCQLIdentifier(table+"_txt_idx")+" ON "+qualifiedCQLName(keyspace, table)+" (txt)", keyspace)
mustCQL(t, server, "INSERT INTO "+qualifiedCQLName(keyspace, table)+" "+
"(id, txt, flag, amount, payload, created, address, tags, items, attrs, pair) VALUES "+
"(1, 'hello', true, 12.34, 0x00ff, '2026-08-03T00:00:00Z', '127.0.0.1', {'a', 'b'}, [1, 2], {'a': 1}, (7, 'seven'))", keyspace)
mustCQL(t, server, "CREATE TABLE "+qualifiedCQLName(keyspace, pagedTable)+" (id int PRIMARY KEY, txt text)", keyspace)
for start := 0; start < 250; start += 50 {
statements := make([]string, 0, 50)
for id := start; id < start+50; id++ {
statements = append(statements, fmt.Sprintf("INSERT INTO %s (id, txt) VALUES (%d, 'row-%d')", qualifiedCQLName(keyspace, pagedTable), id, id))
}
mustStatements(t, server, keyspace, statements, false)
}
mustStatements(t, server, keyspace, []string{
"INSERT INTO " + qualifiedCQLName(keyspace, pagedTable) + " (id, txt) VALUES (1001, 'unlogged')",
}, false)
mustStatements(t, server, keyspace, []string{
"INSERT INTO " + qualifiedCQLName(keyspace, pagedTable) + " (id, txt) VALUES (1002, 'logged')",
}, true)
connectionInfo, err := server.connectionInfo()
if err != nil || strings.TrimSpace(fmt.Sprint(connectionInfo["database_version"])) == "" {
t.Fatalf("connection info failed: info=%v err=%v", connectionInfo, err)
}
databases, err := server.listDatabases()
if err != nil || !containsDatabase(databases, keyspace) {
t.Fatalf("keyspace metadata missing: databases=%v err=%v", databases, err)
}
tables, err := server.listTables(keyspace, metadataListConstraints{})
if err != nil || !containsTable(tables, table) || !containsTable(tables, pagedTable) {
t.Fatalf("table metadata missing: tables=%v err=%v", tables, err)
}
columns, err := server.getColumns(keyspace, table)
if err != nil || len(columns) != 11 || !containsPrimaryKeyColumn(columns, "id") {
t.Fatalf("column metadata mismatch: columns=%v err=%v", columns, err)
}
indexes, err := server.listIndexes(keyspace, table)
if err != nil || !containsIndex(indexes, table+"_txt_idx") {
t.Fatalf("index metadata missing: indexes=%v err=%v", indexes, err)
}
ddl, err := server.getTableDDL(keyspace, table)
if err != nil || !strings.Contains(ddl, "tuple<int, text>") || !strings.Contains(ddl, "PRIMARY KEY") {
t.Fatalf("table DDL mismatch: ddl=%q err=%v", ddl, err)
}
result, err := server.executeQuery(queryOptions{
SQL: "SELECT * FROM " + qualifiedCQLName(keyspace, table) + " WHERE id = 1",
Schema: keyspace,
})
if err != nil || len(result.Rows) != 1 || len(result.Rows[0]) != len(result.Columns) {
t.Fatalf("all-types query failed: result=%v err=%v", result, err)
}
for _, value := range result.Rows[0] {
if value != nil {
if _, ok := value.(string); !ok {
t.Fatalf("legacy result contract requires strings, got %T (%v)", value, value)
}
}
}
page, err := server.executeQueryPage(queryOptions{
SQL: "SELECT id, txt FROM " + qualifiedCQLName(keyspace, pagedTable),
Schema: keyspace,
MaxRows: 250,
}, 100)
if err != nil || len(page.Rows) != 100 || !page.HasMore || page.SessionID == nil {
t.Fatalf("first page mismatch: page=%v err=%v", page, err)
}
totalRows := len(page.Rows)
for page.HasMore {
page, err = server.fetchQueryPage(*page.SessionID, 100)
if err != nil {
t.Fatal(err)
}
totalRows += len(page.Rows)
}
if totalRows != 250 {
t.Fatalf("unexpected paged row count: %d", totalRows)
}
}
func envDefault(name, fallback string) string {
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
return value
}
return fallback
}
func qualifiedCQLName(keyspace, object string) string {
return quoteCQLIdentifier(keyspace) + "." + quoteCQLIdentifier(object)
}
func mustCQL(t *testing.T, server *server, sql, keyspace string) {
t.Helper()
if _, err := server.executeQuery(queryOptions{SQL: sql, Schema: keyspace}); err != nil {
t.Fatalf("execute %q: %v", sql, err)
}
}
func mustStatements(t *testing.T, server *server, keyspace string, statements []string, transactional bool) {
t.Helper()
rawStatements, _ := json.Marshal(statements)
rawSchema, _ := json.Marshal(keyspace)
if _, err := server.executeStatements(map[string]json.RawMessage{
"schema": rawSchema,
"statements": rawStatements,
}, transactional); err != nil {
t.Fatal(err)
}
}
func containsDatabase(databases []databaseInfo, name string) bool {
for _, database := range databases {
if database.Name == name {
return true
}
}
return false
}
func containsTable(tables []tableInfo, name string) bool {
for _, table := range tables {
if table.Name == name {
return true
}
}
return false
}
func containsIndex(indexes []indexInfo, name string) bool {
for _, index := range indexes {
if index.Name == name {
return true
}
}
return false
}
func containsPrimaryKeyColumn(columns []columnInfo, name string) bool {
for _, column := range columns {
if column.Name == name && column.IsPrimaryKey {
return true
}
}
return false
}

View File

@ -48,3 +48,20 @@ func TestTrimStatementSQL(t *testing.T) {
t.Fatalf("unexpected trimmed SQL: %q", got)
}
}
func TestIsSchemaChangingCQL(t *testing.T) {
for _, sql := range []string{
"CREATE TABLE app.events (id int PRIMARY KEY)",
" alter keyspace app with replication = {'class': 'SimpleStrategy'} ",
"DROP INDEX app.events_idx;",
} {
if !isSchemaChangingCQL(sql) {
t.Fatalf("expected schema-changing CQL: %q", sql)
}
}
for _, sql := range []string{"SELECT * FROM app.events", "INSERT INTO app.events (id) VALUES (1)", "TRUNCATE app.events"} {
if isSchemaChangingCQL(sql) {
t.Fatalf("unexpected schema-changing CQL: %q", sql)
}
}
}

View File

@ -35,6 +35,9 @@ func (s *server) executeQuery(options queryOptions) (queryResult, error) {
}
if len(columns) == 0 {
err := iter.Close()
if err == nil && isSchemaChangingCQL(options.SQL) {
s.runtime.invalidateMetadataSession()
}
result.ExecutionTimeMS = time.Since(start).Milliseconds()
return result, err
}
@ -327,3 +330,16 @@ func trimStatementSQL(sql string) string {
}
return trimmed
}
func isSchemaChangingCQL(sql string) bool {
fields := strings.Fields(trimStatementSQL(sql))
if len(fields) == 0 {
return false
}
switch strings.ToUpper(fields[0]) {
case "CREATE", "ALTER", "DROP":
return true
default:
return false
}
}

View File

@ -24,13 +24,15 @@ const (
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
mu sync.Mutex
config cassandraConfig
sessions map[string]*gocql.Session
retiredSessions []*gocql.Session
permits chan struct{}
metadataPermits chan struct{}
activeOperations int
references int
closed bool
}
func newConnectionRuntime(cp connectParams) (*connectionRuntime, error) {
@ -69,6 +71,23 @@ func (r *connectionRuntime) sessionFor(keyspace string) (*gocql.Session, error)
return session, nil
}
func (r *connectionRuntime) invalidateMetadataSession() {
var retiredSession *gocql.Session
r.mu.Lock()
if session := r.sessions[""]; session != nil {
delete(r.sessions, "")
if r.activeOperations == 0 {
retiredSession = session
} else {
r.retiredSessions = append(r.retiredSessions, session)
}
}
r.mu.Unlock()
if retiredSession != nil {
retiredSession.Close()
}
}
func (r *connectionRuntime) acquire(metadata bool) (func(), error) {
ctx, cancel := context.WithTimeout(context.Background(), operationPermitTimeout)
defer cancel()
@ -83,11 +102,25 @@ func (r *connectionRuntime) acquire(metadata bool) (func(), error) {
}
select {
case r.permits <- struct{}{}:
r.mu.Lock()
r.activeOperations++
r.mu.Unlock()
return func() {
var retiredSessions []*gocql.Session
r.mu.Lock()
r.activeOperations--
if r.activeOperations == 0 && len(r.retiredSessions) > 0 {
retiredSessions = r.retiredSessions
r.retiredSessions = nil
}
r.mu.Unlock()
<-r.permits
if metadataAcquired {
<-r.metadataPermits
}
for _, session := range retiredSessions {
session.Close()
}
}, nil
case <-ctx.Done():
if metadataAcquired {
@ -105,11 +138,16 @@ func (r *connectionRuntime) close() {
}
r.closed = true
sessions := r.sessions
retiredSessions := r.retiredSessions
r.sessions = map[string]*gocql.Session{}
r.retiredSessions = nil
r.mu.Unlock()
for _, session := range sessions {
session.Close()
}
for _, session := range retiredSessions {
session.Close()
}
}
func (r *runtimeServer) acquireRuntime(cp connectParams) (*connectionRuntime, string, error) {

View File

@ -0,0 +1,60 @@
package main
import (
"testing"
gocql "github.com/apache/cassandra-gocql-driver/v2"
)
func TestInvalidateMetadataSessionDefersCloseUntilOperationsFinish(t *testing.T) {
runtime := &connectionRuntime{
sessions: map[string]*gocql.Session{"": {}},
permits: make(chan struct{}, 2),
metadataPermits: make(chan struct{}, 1),
}
releaseFirst, err := runtime.acquire(false)
if err != nil {
t.Fatal(err)
}
releaseSecond, err := runtime.acquire(false)
if err != nil {
t.Fatal(err)
}
metadataSession := runtime.sessions[""]
runtime.invalidateMetadataSession()
if metadataSession.Closed() {
t.Fatal("metadata session closed while operations were active")
}
if len(runtime.retiredSessions) != 1 {
t.Fatalf("unexpected retired session count: %d", len(runtime.retiredSessions))
}
releaseFirst()
if metadataSession.Closed() {
t.Fatal("metadata session closed before the final operation completed")
}
releaseSecond()
if !metadataSession.Closed() {
t.Fatal("metadata session was not closed after the final operation completed")
}
if len(runtime.retiredSessions) != 0 {
t.Fatalf("retired sessions were not cleared: %d", len(runtime.retiredSessions))
}
}
func TestInvalidateMetadataSessionClosesImmediatelyWithoutOperations(t *testing.T) {
metadataSession := &gocql.Session{}
runtime := &connectionRuntime{sessions: map[string]*gocql.Session{"": metadataSession}}
runtime.invalidateMetadataSession()
if !metadataSession.Closed() {
t.Fatal("idle metadata session was not closed immediately")
}
if _, exists := runtime.sessions[""]; exists {
t.Fatal("invalidated metadata session remains cached")
}
}

View File

@ -1,10 +0,0 @@
dependencies {
implementation fileTree(dir: 'libs', include: ['*.jar'])
implementation 'com.ing.data:cassandra-jdbc-wrapper:4.12.0'
}
tasks.named('shadowJar') {
manifest {
attributes('Agent-Label': 'Apache Cassandra', 'Main-Class': 'com.dbx.agent.cassandra.CassandraAgent')
}
}

View File

@ -1,211 +0,0 @@
package com.dbx.agent.cassandra;
import com.dbx.agent.AbstractJdbcAgent;
import com.dbx.agent.ColumnInfo;
import com.dbx.agent.ConnectParams;
import com.dbx.agent.DatabaseInfo;
import com.dbx.agent.ForeignKeyInfo;
import com.dbx.agent.IndexInfo;
import com.dbx.agent.JdbcIdentifiers;
import com.dbx.agent.MultiSessionJsonRpcServer;
import com.dbx.agent.TableInfo;
import com.dbx.agent.TriggerInfo;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
public final class CassandraAgent extends AbstractJdbcAgent {
private static final String DC_INFERRING_LOAD_BALANCING = "loadbalancing=DcInferringLoadBalancingPolicy";
private static final Pattern TARGET_PATTERN = Pattern.compile("target[\"']?\\s*[:=]\\s*[\"']?([\\w]+)");
@Override
protected String driverClass() {
return "com.ing.data.cassandra.jdbc.CassandraDriver";
}
@Override
protected String buildJdbcUrl(ConnectParams params) {
return buildUrl(params);
}
@Override
public String setSchemaSQL(String schema) {
return "USE " + JdbcIdentifiers.INSTANCE.doubleQuote(schema);
}
@Override
public List<DatabaseInfo> listDatabases() {
return unchecked(() -> {
List<DatabaseInfo> result = new ArrayList<>();
String sql = "SELECT keyspace_name FROM system_schema.keyspaces";
try (java.sql.Statement stmt = requireConnected().createStatement();
ResultSet rs = stmt.executeQuery(sql)) {
while (rs.next()) {
result.add(new DatabaseInfo(rs.getString(1)));
}
}
result.sort(Comparator.comparing(DatabaseInfo::getName));
return result;
});
}
@Override
public List<String> listSchemas() {
List<String> result = new ArrayList<>();
for (DatabaseInfo database : listDatabases()) {
result.add(database.getName());
}
return result;
}
@Override
public List<TableInfo> listTables(String schema) {
return unchecked(() -> {
List<TableInfo> result = new ArrayList<>();
String sql = "SELECT table_name FROM system_schema.tables WHERE keyspace_name = ?";
try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) {
stmt.setString(1, schema);
try (ResultSet rs = stmt.executeQuery()) {
while (rs.next()) {
result.add(new TableInfo(rs.getString(1), "TABLE", null));
}
}
}
result.sort(Comparator.comparing(TableInfo::getName));
return result;
});
}
@Override
public List<ColumnInfo> getColumns(String schema, String table) {
return unchecked(() -> {
List<ColumnInfo> result = new ArrayList<>();
String sql = "SELECT column_name, type, kind FROM system_schema.columns WHERE keyspace_name = ? AND table_name = ?";
try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) {
stmt.setString(1, schema);
stmt.setString(2, table);
try (ResultSet rs = stmt.executeQuery()) {
while (rs.next()) {
String kind = coalesce(rs.getString("kind"));
boolean isPrimaryKey = "partition_key".equals(kind) || "clustering".equals(kind);
result.add(new ColumnInfo(
rs.getString("column_name"),
coalesce(rs.getString("type"), "unknown"),
!isPrimaryKey,
null,
isPrimaryKey,
kind.trim().isEmpty() ? null : kind,
null,
null,
null,
null
));
}
}
}
return result;
});
}
@Override
public List<IndexInfo> listIndexes(String schema, String table) {
return unchecked(() -> {
List<IndexInfo> result = new ArrayList<>();
String sql = "SELECT index_name, options FROM system_schema.indexes WHERE keyspace_name = ? AND table_name = ?";
try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) {
stmt.setString(1, schema);
stmt.setString(2, table);
try (ResultSet rs = stmt.executeQuery()) {
while (rs.next()) {
String indexName = coalesce(rs.getString("index_name"));
String options = coalesce(rs.getString("options"));
result.add(new IndexInfo(
indexName,
targetColumns(options),
false,
false,
null,
null,
null,
null
));
}
}
}
result.sort(Comparator.comparing(IndexInfo::getName));
return result;
});
}
@Override
public List<ForeignKeyInfo> listForeignKeys(String schema, String table) {
return Collections.emptyList();
}
@Override
public List<TriggerInfo> listTriggers(String schema, String table) {
return Collections.emptyList();
}
@Override
protected Object resultValue(ResultSet rs, int index, int sqlType) {
return unchecked(() -> {
Object value = rs.getObject(index);
return rs.wasNull() ? null : value == null ? null : value.toString();
});
}
static String buildUrl(ConnectParams params) {
String baseUrl = "jdbc:cassandra://" + params.getHost() + ":" + params.getPort();
String keyspace = coalesce(params.getDatabase()).trim();
// Cassandra rejects an empty keyspace path; omit it so DBX can connect first and list keyspaces.
String url = keyspace.isEmpty() ? baseUrl : baseUrl + "/" + keyspace;
String extraParams = coalesce(params.getUrl_params()).trim();
while (extraParams.startsWith("?") || extraParams.startsWith("&")) {
extraParams = extraParams.substring(1);
}
if (!hasDatacenterOrLoadBalancingParameter(extraParams)) {
extraParams = extraParams.isEmpty()
? DC_INFERRING_LOAD_BALANCING
: extraParams + (extraParams.endsWith("&") ? "" : "&") + DC_INFERRING_LOAD_BALANCING;
}
return extraParams.isEmpty() ? url : url + "?" + extraParams;
}
private static boolean hasDatacenterOrLoadBalancingParameter(String urlParams) {
for (String param : urlParams.split("&")) {
int valueSeparator = param.indexOf('=');
String paramName = valueSeparator < 0 ? param : param.substring(0, valueSeparator);
paramName = paramName.trim();
if (paramName.equals("localdatacenter") || paramName.equals("loadbalancing")) {
return true;
}
}
return false;
}
private static List<String> targetColumns(String options) {
Matcher matcher = TARGET_PATTERN.matcher(options);
if (!matcher.find()) {
return Collections.emptyList();
}
return Collections.singletonList(matcher.group(1));
}
private static String coalesce(String value) {
return value == null ? "" : value;
}
private static String coalesce(String value, String fallback) {
return value == null ? fallback : value;
}
public static void main(String[] args) {
new MultiSessionJsonRpcServer(CassandraAgent::new).run();
}
}

View File

@ -1,94 +0,0 @@
package com.dbx.agent.cassandra;
import com.dbx.agent.ConnectParams;
import com.dbx.agent.DatabaseAgent;
import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.assertEquals;
class CassandraAgentTest extends JdbcFakeExecutionBehaviorTest {
@Override
protected DatabaseAgent createAgent() {
return new CassandraAgent();
}
@Override
protected String resultSetSql() {
return "LIST ROLES";
}
@Test
void buildsServerUrlWhenKeyspaceIsEmpty() {
ConnectParams params = new ConnectParams("127.0.0.1", 9042, "", "cassandra", "cassandra", "", "", false);
assertEquals(
"jdbc:cassandra://127.0.0.1:9042?loadbalancing=DcInferringLoadBalancingPolicy",
CassandraAgent.buildUrl(params)
);
}
@Test
void buildsKeyspaceUrlWhenKeyspaceIsSet() {
ConnectParams params = new ConnectParams("127.0.0.1", 9042, "app_keyspace", "cassandra", "cassandra", "", "", false);
assertEquals(
"jdbc:cassandra://127.0.0.1:9042/app_keyspace?loadbalancing=DcInferringLoadBalancingPolicy",
CassandraAgent.buildUrl(params)
);
}
@Test
void appendsUrlParamsForMultiDcLocalDatacenter() {
ConnectParams params = new ConnectParams(
"127.0.0.1", 9042, "app_keyspace", "cassandra", "cassandra", "localdatacenter=dc1", "", false
);
assertEquals("jdbc:cassandra://127.0.0.1:9042/app_keyspace?localdatacenter=dc1", CassandraAgent.buildUrl(params));
}
@Test
void stripsLeadingQuestionMarkFromUrlParams() {
ConnectParams params = new ConnectParams(
"127.0.0.1", 9042, "", "cassandra", "cassandra", "?localdatacenter=dc1", "", false
);
assertEquals("jdbc:cassandra://127.0.0.1:9042?localdatacenter=dc1", CassandraAgent.buildUrl(params));
}
@Test
void appendsInferringPolicyAfterOtherUrlParams() {
ConnectParams params = new ConnectParams(
"127.0.0.1", 9042, "", "cassandra", "cassandra", "requesttimeout=10000", "", false
);
assertEquals(
"jdbc:cassandra://127.0.0.1:9042?requesttimeout=10000&loadbalancing=DcInferringLoadBalancingPolicy",
CassandraAgent.buildUrl(params)
);
}
@Test
void appendsInferringPolicyAfterTrailingSeparator() {
ConnectParams params = new ConnectParams(
"127.0.0.1", 9042, "", "cassandra", "cassandra", "requesttimeout=10000&", "", false
);
assertEquals(
"jdbc:cassandra://127.0.0.1:9042?requesttimeout=10000&loadbalancing=DcInferringLoadBalancingPolicy",
CassandraAgent.buildUrl(params)
);
}
@Test
void preservesCustomLoadBalancingPolicy() {
ConnectParams params = new ConnectParams(
"127.0.0.1", 9042, "", "cassandra", "cassandra", "loadbalancing=CustomPolicy", "", false
);
assertEquals(
"jdbc:cassandra://127.0.0.1:9042?loadbalancing=CustomPolicy",
CassandraAgent.buildUrl(params)
);
}
}

View File

@ -5,7 +5,7 @@ def driverModules = [
'access', 'dameng', 'goldendb', 'databend', 'databricks', 'saphana',
'teradata', 'vertica', 'firebird', 'exasol', 'oceanbase-oracle', 'gbase8a', 'gbase8s',
'bigquery', 'kylin', 'sundb', 'h2', 'h2-legacy', 'snowflake', 'trino', 'hive', 'spark',
'db2', 'informix', 'neo4j', 'cassandra', 'mongodb', 'highgo', 'uxdb', 'tdengine', 'yashandb', 'oscar',
'db2', 'informix', 'neo4j', 'mongodb', 'highgo', 'uxdb', 'tdengine', 'yashandb', 'oscar',
'iris', 'iotdb', 'etcd', 'zookeeper', 'kafka', 'rocketmq', 'sqlserver-legacy'
]