From 44ef416d2fd72796b05c85921f23d5ff4c664045 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 4 Aug 2026 01:21:17 +0800 Subject: [PATCH] feat(cassandra): complete native agent migration --- .github/workflows/ci.yml | 41 +++ agents/drivers/cassandra-go/README.md | 61 ++++ agents/drivers/cassandra-go/bench/README.md | 61 ++++ .../cassandra-go/bench/agent_compare.py | 328 ++++++++++++++++++ .../bench/results/cassandra-4.1.10.json | 36 ++ agents/drivers/cassandra-go/config.go | 307 ++++++++++++++-- agents/drivers/cassandra-go/config_test.go | 105 +++++- agents/drivers/cassandra-go/dialer.go | 48 +++ .../drivers/cassandra-go/integration_test.go | 206 +++++++++++ agents/drivers/cassandra-go/protocol_test.go | 17 + agents/drivers/cassandra-go/query.go | 16 + agents/drivers/cassandra-go/runtime.go | 52 ++- agents/drivers/cassandra-go/runtime_test.go | 60 ++++ agents/drivers/cassandra/build.gradle | 10 - .../dbx/agent/cassandra/CassandraAgent.java | 211 ----------- .../agent/cassandra/CassandraAgentTest.java | 94 ----- agents/settings.gradle | 2 +- 17 files changed, 1299 insertions(+), 356 deletions(-) create mode 100644 agents/drivers/cassandra-go/README.md create mode 100644 agents/drivers/cassandra-go/bench/README.md create mode 100755 agents/drivers/cassandra-go/bench/agent_compare.py create mode 100644 agents/drivers/cassandra-go/bench/results/cassandra-4.1.10.json create mode 100644 agents/drivers/cassandra-go/dialer.go create mode 100644 agents/drivers/cassandra-go/integration_test.go create mode 100644 agents/drivers/cassandra-go/runtime_test.go delete mode 100644 agents/drivers/cassandra/build.gradle delete mode 100644 agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java delete mode 100644 agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3e5c8e11a..c8c1d7036 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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 diff --git a/agents/drivers/cassandra-go/README.md b/agents/drivers/cassandra-go/README.md new file mode 100644 index 000000000..b87b89ea7 --- /dev/null +++ b/agents/drivers/cassandra-go/README.md @@ -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. diff --git a/agents/drivers/cassandra-go/bench/README.md b/agents/drivers/cassandra-go/bench/README.md new file mode 100644 index 000000000..5d2a75043 --- /dev/null +++ b/agents/drivers/cassandra-go/bench/README.md @@ -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. diff --git a/agents/drivers/cassandra-go/bench/agent_compare.py b/agents/drivers/cassandra-go/bench/agent_compare.py new file mode 100755 index 000000000..42fa909da --- /dev/null +++ b/agents/drivers/cassandra-go/bench/agent_compare.py @@ -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() diff --git a/agents/drivers/cassandra-go/bench/results/cassandra-4.1.10.json b/agents/drivers/cassandra-go/bench/results/cassandra-4.1.10.json new file mode 100644 index 000000000..bcd25a4bb --- /dev/null +++ b/agents/drivers/cassandra-go/bench/results/cassandra-4.1.10.json @@ -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} + ] + } + ] +} diff --git a/agents/drivers/cassandra-go/config.go b/agents/drivers/cassandra-go/config.go index 39b829fa9..462ae3299 100644 --- a/agents/drivers/cassandra-go/config.go +++ b/agents/drivers/cassandra-go/config.go @@ -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))) } diff --git a/agents/drivers/cassandra-go/config_test.go b/agents/drivers/cassandra-go/config_test.go index e3136041e..8dd8aded7 100644 --- a/agents/drivers/cassandra-go/config_test.go +++ b/agents/drivers/cassandra-go/config_test.go @@ -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", diff --git a/agents/drivers/cassandra-go/dialer.go b/agents/drivers/cassandra-go/dialer.go new file mode 100644 index 000000000..a2051caf8 --- /dev/null +++ b/agents/drivers/cassandra-go/dialer.go @@ -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 +} diff --git a/agents/drivers/cassandra-go/integration_test.go b/agents/drivers/cassandra-go/integration_test.go new file mode 100644 index 000000000..017c13097 --- /dev/null +++ b/agents/drivers/cassandra-go/integration_test.go @@ -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, items list, attrs map, pair frozen>)", 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") || !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 +} diff --git a/agents/drivers/cassandra-go/protocol_test.go b/agents/drivers/cassandra-go/protocol_test.go index c4d691b65..ee125b50b 100644 --- a/agents/drivers/cassandra-go/protocol_test.go +++ b/agents/drivers/cassandra-go/protocol_test.go @@ -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) + } + } +} diff --git a/agents/drivers/cassandra-go/query.go b/agents/drivers/cassandra-go/query.go index 521e759ca..ddc288898 100644 --- a/agents/drivers/cassandra-go/query.go +++ b/agents/drivers/cassandra-go/query.go @@ -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 + } +} diff --git a/agents/drivers/cassandra-go/runtime.go b/agents/drivers/cassandra-go/runtime.go index 1021a6b2d..a03c1a285 100644 --- a/agents/drivers/cassandra-go/runtime.go +++ b/agents/drivers/cassandra-go/runtime.go @@ -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) { diff --git a/agents/drivers/cassandra-go/runtime_test.go b/agents/drivers/cassandra-go/runtime_test.go new file mode 100644 index 000000000..3f7de69a1 --- /dev/null +++ b/agents/drivers/cassandra-go/runtime_test.go @@ -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") + } +} diff --git a/agents/drivers/cassandra/build.gradle b/agents/drivers/cassandra/build.gradle deleted file mode 100644 index 38a8d6ffd..000000000 --- a/agents/drivers/cassandra/build.gradle +++ /dev/null @@ -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') - } -} diff --git a/agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java b/agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java deleted file mode 100644 index f7cdc196f..000000000 --- a/agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java +++ /dev/null @@ -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 listDatabases() { - return unchecked(() -> { - List 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 listSchemas() { - List result = new ArrayList<>(); - for (DatabaseInfo database : listDatabases()) { - result.add(database.getName()); - } - return result; - } - - @Override - public List listTables(String schema) { - return unchecked(() -> { - List 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 getColumns(String schema, String table) { - return unchecked(() -> { - List 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 listIndexes(String schema, String table) { - return unchecked(() -> { - List 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 listForeignKeys(String schema, String table) { - return Collections.emptyList(); - } - - @Override - public List 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 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(); - } -} diff --git a/agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java b/agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java deleted file mode 100644 index 5873d8abc..000000000 --- a/agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java +++ /dev/null @@ -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) - ); - } -} diff --git a/agents/settings.gradle b/agents/settings.gradle index 78a7f3b64..6b55cb674 100644 --- a/agents/settings.gradle +++ b/agents/settings.gradle @@ -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' ]