From 27f3ea728758c2946c131d6e9e51118ac7eec49f Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 10 Aug 2026 13:53:43 +0800 Subject: [PATCH] feat(iotdb): replace JDBC agent with native Go driver --- .github/scripts/bump-agent-versions.mjs | 3 +- .github/scripts/bump-agent-versions.test.mjs | 13 + .../scripts/reuse-agent-release-assets.mjs | 2 +- .../reuse-agent-release-assets.test.mjs | 18 + .github/workflows/agents-release.yml | 48 +- .github/workflows/ci.yml | 10 +- README.md | 4 +- README.zh-CN.md | 4 +- agents/README.md | 11 +- agents/README.zh-CN.md | 11 +- agents/build.gradle | 2 +- agents/drivers/iotdb/README.md | 55 ++ agents/drivers/iotdb/bench/README.md | 59 ++ agents/drivers/iotdb/bench/build.gradle | 31 + agents/drivers/iotdb/bench/go/go.mod | 7 + agents/drivers/iotdb/bench/go/go.sum | 29 + agents/drivers/iotdb/bench/go/main.go | 284 ++++++ .../iotdb/bench/JdbcDriverBenchmark.java | 287 ++++++ agents/drivers/iotdb/bench/results/README.md | 32 + .../results/iotdb-2.0.8-macos-arm64.json | 839 ++++++++++++++++++ agents/drivers/iotdb/bench/run.py | 253 ++++++ agents/drivers/iotdb/bench/settings.gradle | 1 + agents/drivers/iotdb/build.gradle | 10 - agents/drivers/iotdb/driver.go | 333 +++++++ agents/drivers/iotdb/go.mod | 7 + agents/drivers/iotdb/go.sum | 4 + agents/drivers/iotdb/integration_test.go | 331 +++++++ agents/drivers/iotdb/main.go | 515 +++++++++++ agents/drivers/iotdb/main_test.go | 199 +++++ agents/drivers/iotdb/metadata.go | 703 +++++++++++++++ agents/drivers/iotdb/protocol_error.go | 138 +++ agents/drivers/iotdb/query.go | 515 +++++++++++ .../java/com/dbx/agent/iotdb/IoTDBAgent.java | 46 - .../com/dbx/agent/iotdb/IoTDBAgentTest.java | 32 - agents/metadata-constraint-coverage.tsv | 1 + .../scripts/driver_release_packages_test.py | 16 + agents/scripts/validate_agents.py | 1 + agents/scripts/validate_agents_test.py | 3 +- agents/scripts/version_agent_artifacts.py | 2 +- agents/settings.gradle | 2 +- 40 files changed, 4750 insertions(+), 111 deletions(-) create mode 100644 agents/drivers/iotdb/README.md create mode 100644 agents/drivers/iotdb/bench/README.md create mode 100644 agents/drivers/iotdb/bench/build.gradle create mode 100644 agents/drivers/iotdb/bench/go/go.mod create mode 100644 agents/drivers/iotdb/bench/go/go.sum create mode 100644 agents/drivers/iotdb/bench/go/main.go create mode 100644 agents/drivers/iotdb/bench/java/com/dbx/agent/iotdb/bench/JdbcDriverBenchmark.java create mode 100644 agents/drivers/iotdb/bench/results/README.md create mode 100644 agents/drivers/iotdb/bench/results/iotdb-2.0.8-macos-arm64.json create mode 100644 agents/drivers/iotdb/bench/run.py create mode 100644 agents/drivers/iotdb/bench/settings.gradle delete mode 100644 agents/drivers/iotdb/build.gradle create mode 100644 agents/drivers/iotdb/driver.go create mode 100644 agents/drivers/iotdb/go.mod create mode 100644 agents/drivers/iotdb/go.sum create mode 100644 agents/drivers/iotdb/integration_test.go create mode 100644 agents/drivers/iotdb/main.go create mode 100644 agents/drivers/iotdb/main_test.go create mode 100644 agents/drivers/iotdb/metadata.go create mode 100644 agents/drivers/iotdb/protocol_error.go create mode 100644 agents/drivers/iotdb/query.go delete mode 100644 agents/drivers/iotdb/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java delete mode 100644 agents/drivers/iotdb/src/test/java/com/dbx/agent/iotdb/IoTDBAgentTest.java diff --git a/.github/scripts/bump-agent-versions.mjs b/.github/scripts/bump-agent-versions.mjs index b1e273cd6a..46708b35cf 100644 --- a/.github/scripts/bump-agent-versions.mjs +++ b/.github/scripts/bump-agent-versions.mjs @@ -53,12 +53,13 @@ const nativeDriverDirectories = { duckdb: "duckdb", oracle: "oracle-go", kingbase: "kingbase-go", + iotdb: "iotdb", neo4j: "neo4j-go", vastbase: "vastbase-go", rabbitmq: "rabbitmq", tdengine: "tdengine", }; -const nativeDriverModules = new Set(["cassandra", "duckdb", "oracle", "xugu", "kingbase", "neo4j", "vastbase", "rabbitmq", "tdengine"]); +const nativeDriverModules = new Set(["cassandra", "duckdb", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "tdengine"]); function resolveAgentModule(moduleName, { legacyStandaloneModules, moduleExists, readModuleFile }) { let checkDir = null; diff --git a/.github/scripts/bump-agent-versions.test.mjs b/.github/scripts/bump-agent-versions.test.mjs index dba4b308d2..533e1fcd5c 100644 --- a/.github/scripts/bump-agent-versions.test.mjs +++ b/.github/scripts/bump-agent-versions.test.mjs @@ -96,6 +96,19 @@ test("bumps Neo4j from its native Go source directory", () => { assert.deepEqual(result.nativeModules, ["neo4j"]); }); +test("bumps IoTDB from its native Go source directory", () => { + const result = evaluateAgentVersionBump({ + versions: { iotdb: "0.1.30" }, + changedFiles: ["agents/drivers/iotdb/main.go"], + moduleExists: (path) => path === "agents/drivers/iotdb", + readModuleFile: () => "", + }); + + assert.equal(result.versions.iotdb, "0.1.31"); + assert.deepEqual(result.javaModules, []); + assert.deepEqual(result.nativeModules, ["iotdb"]); +}); + test("rebuilds native modules when shared native packaging changes", () => { const existing = new Set([ "agents/drivers/access", diff --git a/.github/scripts/reuse-agent-release-assets.mjs b/.github/scripts/reuse-agent-release-assets.mjs index 5499bba222..5c863ff2ac 100644 --- a/.github/scripts/reuse-agent-release-assets.mjs +++ b/.github/scripts/reuse-agent-release-assets.mjs @@ -15,7 +15,7 @@ import { basename, join } from "node:path"; import { tmpdir } from "node:os"; const REGISTRY_ASSET = "agent-registry.json"; -const NATIVE_MODULES = new Set(["duckdb", "oracle", "xugu", "kingbase", "neo4j", "vastbase", "rabbitmq", "tdengine"]); +const NATIVE_MODULES = new Set(["duckdb", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "tdengine"]); const PLATFORMS = [ "macos-aarch64", "macos-x64", diff --git a/.github/scripts/reuse-agent-release-assets.test.mjs b/.github/scripts/reuse-agent-release-assets.test.mjs index 9a57ec0cab..16c6453074 100644 --- a/.github/scripts/reuse-agent-release-assets.test.mjs +++ b/.github/scripts/reuse-agent-release-assets.test.mjs @@ -96,6 +96,24 @@ test("requires all Neo4j native platforms when reusing a release", () => { ); }); +test("requires all IoTDB native platforms when reusing a release", () => { + const native = Object.fromEntries( + platforms.slice(0, -1).map((platform, index) => [platform, artifact(`dbx-agent-iotdb-0.1.30-${platform}.tar.zst`, String(index + 1))]), + ); + const registry = { drivers: { iotdb: { version: "0.1.30", native } }, jres: {} }; + + assert.throws( + () => collectReusableAssetPlan({ + registry, + release: releaseFor(Object.values(native)), + versions: { iotdb: "0.1.30" }, + modules: ["iotdb"], + reuseJre: false, + }), + /missing=windows-x64/, + ); +}); + test("ignores zero-size legacy JAR placeholders for native-only modules", () => { const native = Object.fromEntries( platforms.map((platform, index) => [platform, artifact(`dbx-agent-duckdb-0.1.2-${platform}.tar.zst`, String(index + 1))]), diff --git a/.github/workflows/agents-release.yml b/.github/workflows/agents-release.yml index 2d1685d62f..62bbe49c1c 100644 --- a/.github/workflows/agents-release.yml +++ b/.github/workflows/agents-release.yml @@ -449,6 +449,46 @@ jobs: name: neo4j-native path: "release-native/dbx-agent-neo4j-*" + build-iotdb-native: + needs: [bump-versions] + if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'iotdb') }} + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: "1.25.x" + - name: Test IoTDB native agent + working-directory: agents/drivers/iotdb + run: go test ./... + - name: Cross-compile IoTDB native agent + shell: bash + run: | + mkdir -p release-native + cd agents/drivers/iotdb + declare -A TARGETS=( + ["macos-aarch64"]="darwin/arm64" + ["macos-x64"]="darwin/amd64" + ["linux-aarch64"]="linux/arm64" + ["linux-x64"]="linux/amd64" + ["windows-aarch64"]="windows/arm64" + ["windows-x64"]="windows/amd64" + ) + for platform in "${!TARGETS[@]}"; do + IFS=/ read -r goos goarch <<< "${TARGETS[$platform]}" + output="../../../release-native/dbx-agent-iotdb-${platform}" + if [[ "$goos" == "windows" ]]; then + output="${output}.exe" + fi + echo "Building $platform ($goos/$goarch)" + CGO_ENABLED=0 GOOS="$goos" GOARCH="$goarch" go build -trimpath -ldflags="-s -w" -o "$output" . + done + ls -lh ../../../release-native + - uses: actions/upload-artifact@v4 + with: + name: iotdb-native + path: "release-native/dbx-agent-iotdb-*" + build-duckdb-native: needs: [bump-versions] if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'duckdb') }} @@ -788,7 +828,7 @@ jobs: retention-days: 1 release: - needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-cassandra-native, build-kingbase-native, build-vastbase-native, build-neo4j-native, build-duckdb-native, build-tdengine-native, build-jre, reuse-previous-assets] + needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-cassandra-native, build-kingbase-native, build-vastbase-native, build-neo4j-native, build-iotdb-native, build-duckdb-native, build-tdengine-native, build-jre, reuse-previous-assets] if: ${{ always() && !contains(needs.*.result, 'failure') && !contains(needs.*.result, 'cancelled') }} runs-on: ubuntu-latest steps: @@ -912,6 +952,7 @@ jobs: rabbitmq) echo "RabbitMQ" ;; cassandra) echo "Apache Cassandra" ;; neo4j) echo "Neo4j" ;; + iotdb) echo "Apache IoTDB" ;; tdengine) echo "TDengine" ;; *) echo "$name" ;; esac @@ -976,7 +1017,7 @@ jobs: [ -n "$DRIVERS" ] && DRIVERS="${DRIVERS},"$'\n' DRIVERS="${DRIVERS}$(generate_jar_entry "$name" "$label" "$f" "$jre_key" "$version" "$external_driver" "$native_json")" done - for name in oracle xugu kingbase vastbase neo4j duckdb rabbitmq cassandra tdengine; do + for name in oracle xugu kingbase vastbase neo4j iotdb duckdb rabbitmq cassandra tdengine; do version=$(get_module_version "$name") [ -f "release/dbx-agent-${name}-${version}.jar" ] && continue native_json=$(generate_native_platforms "$name" "$version") @@ -1048,6 +1089,7 @@ jobs: rabbitmq) echo "RabbitMQ" ;; cassandra) echo "Apache Cassandra" ;; neo4j) echo "Neo4j" ;; + iotdb) echo "Apache IoTDB" ;; tdengine) echo "TDengine" ;; *) echo "$name" ;; esac @@ -1077,6 +1119,8 @@ jobs: LOG_PATH="agents/drivers/cassandra-go/" elif [ "$name" = "neo4j" ]; then LOG_PATH="agents/drivers/neo4j-go/" + elif [ "$name" = "iotdb" ]; then + LOG_PATH="agents/drivers/iotdb/" elif [ -d "agents/drivers/$name" ]; then LOG_PATH="agents/drivers/$name/" else diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aa761698e6..3c85b3f5a6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -701,7 +701,7 @@ jobs: - name: Setup Go uses: actions/setup-go@v5 with: - go-version: "1.24.x" + go-version: "1.25.x" - name: Setup Rust uses: dtolnay/rust-toolchain@1.97.1 @@ -752,6 +752,10 @@ jobs: run: go test ./... working-directory: agents/drivers/neo4j-go + - name: IoTDB native agent tests + run: go test ./... + working-directory: agents/drivers/iotdb + - name: Oracle native agent build run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-oracle-linux-x64 . working-directory: agents/drivers/oracle-go @@ -776,6 +780,10 @@ jobs: run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-neo4j-linux-x64 . working-directory: agents/drivers/neo4j-go + - name: IoTDB native agent build + run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-iotdb-linux-x64 . + working-directory: agents/drivers/iotdb + - name: TDengine native agent build run: cargo build --manifest-path drivers/tdengine/Cargo.toml --locked --release --bin dbx-tdengine-driver diff --git a/README.md b/README.md index ddde9ff035..394161b18d 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,7 @@ ### 70+ Databases, One Tool -MySQL, PostgreSQL, SQLite, Cloudflare D1, Redis, MongoDB, DuckDB, ClickHouse, SQL Server, Oracle, Elasticsearch, Easysearch, Qdrant, Milvus, Weaviate, MariaDB, TiDB, OceanBase, openGauss, GaussDB, KWDB, KingBase, Vastbase, GoldenDB, Doris, SelectDB, StarRocks, Manticore Search, Redshift, DM, TDengine, XuguDB, CockroachDB, Access, HighGo, UXDB, Dolt, and more. Agent/JDBC-oriented profiles extend DBX to H2, Snowflake, Trino, PrestoSQL, Hive, DB2, Informix, Neo4j, Cassandra, BigQuery, Kylin, SunDB, JDBCX, and custom JDBC connections. New native and agent-driven drivers also cover Databricks, SAP HANA, Teradata, Vertica, Firebird, Exasol, YashanDB, GBase 8a/8s, Databend, RQLite, Turso, InfluxDB, QuestDB, IoTDB, etcd, ZooKeeper, Nacos, IRIS, and more. Message queue admin is also available for Pulsar, Kafka, and RocketMQ. All in a single ~20 MB app. No bundled Chromium. +MySQL, PostgreSQL, SQLite, Cloudflare D1, Redis, MongoDB, DuckDB, ClickHouse, SQL Server, Oracle, Elasticsearch, Easysearch, Qdrant, Milvus, Weaviate, MariaDB, TiDB, OceanBase, openGauss, GaussDB, KWDB, KingBase, Vastbase, GoldenDB, Doris, SelectDB, StarRocks, Manticore Search, Redshift, DM, TDengine, XuguDB, CockroachDB, Access, HighGo, UXDB, Dolt, and more. Agent-based profiles extend DBX to H2, Snowflake, Trino, PrestoSQL, Hive, DB2, Informix, Neo4j, Cassandra, BigQuery, Kylin, SunDB, JDBCX, and custom JDBC connections. New native and agent-driven drivers also cover Databricks, SAP HANA, Teradata, Vertica, Firebird, Exasol, YashanDB, GBase 8a/8s, Databend, RQLite, Turso, InfluxDB, QuestDB, IoTDB, etcd, ZooKeeper, Nacos, IRIS, and more. Message queue admin is also available for Pulsar, Kafka, and RocketMQ. All in a single ~20 MB app. No bundled Chromium. ### Query Editor @@ -483,7 +483,7 @@ DBX is 20 MB with no runtime dependencies for its native database features (no s
What databases are supported? -MySQL, PostgreSQL, SQLite, Cloudflare D1, Redis, MongoDB, DuckDB, ClickHouse, SQL Server, Oracle, Elasticsearch, Easysearch, Qdrant, Milvus, Weaviate, MariaDB, TiDB, OceanBase, openGauss, GaussDB, KWDB, KingBase, Vastbase, GoldenDB, Doris, SelectDB, StarRocks, Manticore Search, Redshift, DM, TDengine, XuguDB, CockroachDB, Access, HighGo, UXDB, and more. Agent/JDBC-oriented profiles extend support to H2, Snowflake, Trino, PrestoSQL, Hive, DB2, Informix, Neo4j, Cassandra, BigQuery, Kylin, SunDB, JDBCX, Databricks, SAP HANA, Teradata, Vertica, Firebird, Exasol, YashanDB, GBase 8a/8s, Databend, RQLite, Turso, InfluxDB, QuestDB, IoTDB, etcd, ZooKeeper, Nacos, IRIS, and custom JDBC connections. Message queue admin (Pulsar, Kafka, RocketMQ) is also supported. +MySQL, PostgreSQL, SQLite, Cloudflare D1, Redis, MongoDB, DuckDB, ClickHouse, SQL Server, Oracle, Elasticsearch, Easysearch, Qdrant, Milvus, Weaviate, MariaDB, TiDB, OceanBase, openGauss, GaussDB, KWDB, KingBase, Vastbase, GoldenDB, Doris, SelectDB, StarRocks, Manticore Search, Redshift, DM, TDengine, XuguDB, CockroachDB, Access, HighGo, UXDB, and more. Agent-based profiles extend support to H2, Snowflake, Trino, PrestoSQL, Hive, DB2, Informix, Neo4j, Cassandra, BigQuery, Kylin, SunDB, JDBCX, Databricks, SAP HANA, Teradata, Vertica, Firebird, Exasol, YashanDB, GBase 8a/8s, Databend, RQLite, Turso, InfluxDB, QuestDB, IoTDB, etcd, ZooKeeper, Nacos, IRIS, and custom JDBC connections. Message queue admin (Pulsar, Kafka, RocketMQ) is also supported.
diff --git a/README.zh-CN.md b/README.zh-CN.md index 3900c6a1b6..3cd61555d1 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -90,7 +90,7 @@ ### 70+ 种数据库，一个工具搞定 -MySQL、PostgreSQL、SQLite、Cloudflare D1、Redis、MongoDB、DuckDB、ClickHouse、SQL Server、Oracle、Elasticsearch、Easysearch、MariaDB、TiDB、OceanBase、openGauss、GaussDB、KWDB、KingBase、Vastbase、GoldenDB、Doris、SelectDB、StarRocks、Manticore Search、Redshift、DM、TDengine、虚谷 XuguDB、CockroachDB、Access、HighGo、UXDB、Dolt 等数据库都能直接连接。Agent/JDBC 方向的配置还可扩展到 H2、Snowflake、Trino、Hive、DB2、Informix、Neo4j、Cassandra、BigQuery、Kylin、SunDB、JDBCX 和自定义 JDBC。新增的原生与 Agent 驱动还覆盖了 Databricks、SAP HANA、Teradata、Vertica、Firebird、Exasol、崖山 YashanDB、GBase、Databend、RQLite、Turso、InfluxDB、QuestDB、IoTDB、etcd、IRIS 等。全部装进约 20 MB 的应用里，不内嵌 Chromium。 +MySQL、PostgreSQL、SQLite、Cloudflare D1、Redis、MongoDB、DuckDB、ClickHouse、SQL Server、Oracle、Elasticsearch、Easysearch、MariaDB、TiDB、OceanBase、openGauss、GaussDB、KWDB、KingBase、Vastbase、GoldenDB、Doris、SelectDB、StarRocks、Manticore Search、Redshift、DM、TDengine、虚谷 XuguDB、CockroachDB、Access、HighGo、UXDB、Dolt 等数据库都能直接连接。Agent 配置还可扩展到 H2、Snowflake、Trino、Hive、DB2、Informix、Neo4j、Cassandra、BigQuery、Kylin、SunDB、JDBCX 和自定义 JDBC。新增的原生与 Agent 驱动还覆盖了 Databricks、SAP HANA、Teradata、Vertica、Firebird、Exasol、崖山 YashanDB、GBase、Databend、RQLite、Turso、InfluxDB、QuestDB、IoTDB、etcd、IRIS 等。全部装进约 20 MB 的应用里，不内嵌 Chromium。 ### 查询编辑器 @@ -486,7 +486,7 @@ DBX 仅 20 MB，无需运行时依赖（无需 Java、无需 Python）。AI 和
支持哪些数据库？ -MySQL、PostgreSQL、SQLite、Cloudflare D1、Redis、MongoDB、DuckDB、ClickHouse、SQL Server、Oracle、Elasticsearch、Easysearch、Qdrant、Milvus、Weaviate、MariaDB、TiDB、OceanBase、openGauss、GaussDB、KWDB、KingBase、Vastbase、GoldenDB、Doris、SelectDB、StarRocks、Manticore Search、Redshift、DM、TDengine、虚谷 XuguDB、CockroachDB、Access、HighGo、UXDB 等。JDBC 方向配置可扩展到 H2、Snowflake、Trino、PrestoSQL、Hive、DB2、Informix、Neo4j、Cassandra、BigQuery、Kylin、SunDB、JDBCX、Databricks、SAP HANA、Teradata、Vertica、Firebird、Exasol、崖山 YashanDB、GBase 8a/8s、Databend、RQLite、Turso、InfluxDB、QuestDB、IoTDB、etcd、ZooKeeper、Nacos、IRIS 及自定义 JDBC 连接，并支持消息队列管理（Pulsar、Kafka、RocketMQ）。 +MySQL、PostgreSQL、SQLite、Cloudflare D1、Redis、MongoDB、DuckDB、ClickHouse、SQL Server、Oracle、Elasticsearch、Easysearch、Qdrant、Milvus、Weaviate、MariaDB、TiDB、OceanBase、openGauss、GaussDB、KWDB、KingBase、Vastbase、GoldenDB、Doris、SelectDB、StarRocks、Manticore Search、Redshift、DM、TDengine、虚谷 XuguDB、CockroachDB、Access、HighGo、UXDB 等。Agent 配置可扩展到 H2、Snowflake、Trino、PrestoSQL、Hive、DB2、Informix、Neo4j、Cassandra、BigQuery、Kylin、SunDB、JDBCX、Databricks、SAP HANA、Teradata、Vertica、Firebird、Exasol、崖山 YashanDB、GBase 8a/8s、Databend、RQLite、Turso、InfluxDB、QuestDB、IoTDB、etcd、ZooKeeper、Nacos、IRIS 及自定义 JDBC 连接，并支持消息队列管理（Pulsar、Kafka、RocketMQ）。
diff --git a/agents/README.md b/agents/README.md index 2056c8a2a9..03d858345c 100644 --- a/agents/README.md +++ b/agents/README.md @@ -41,7 +41,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo | tdengine | TDengine 2.4+ | taos-connector-rust native WebSocket agent | | yashandb | 崖山 YashanDB | YashanDB JDBC | | xugu | 虚谷 XuguDB | XuguDB Go native agent | -| iotdb | Apache IoTDB | IoTDB JDBC | +| iotdb | Apache IoTDB | Apache IoTDB Go Client native agent | | etcd | etcd | jetcd | | zookeeper | Apache ZooKeeper | Apache Curator | | rabbitmq | RabbitMQ | amqp091-go native agent | @@ -49,7 +49,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo ## Multi-JRE Support -Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents. +Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents. ## JDBC Connection Pooling @@ -75,7 +75,7 @@ Set `DBX_AGENT_JDBC_POOL_ENABLED=false` for a runtime-level compatibility fallba For new agents, prefer a **native (Go or Rust) driver** over a Java/JDBC agent whenever a mature, license-compatible native driver is available. Native agents ship as a single self-contained executable with no JRE, which significantly reduces memory footprint and startup time — the JVM baseline that every Java agent pays even when idle is avoided entirely. -- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/cassandra-go` (Apache cassandra-gocql-driver), `drivers/duckdb`, `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/tdengine` (taos-connector-rust), `drivers/xugu`, and `drivers/rabbitmq` (amqp091-go) as reference implementations. No JRE download or management is needed. +- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/cassandra-go` (Apache cassandra-gocql-driver), `drivers/duckdb`, `drivers/iotdb` (Apache IoTDB Go Client), `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/tdengine` (taos-connector-rust), `drivers/xugu`, and `drivers/rabbitmq` (amqp091-go) as reference implementations. No JRE download or management is needed. - **Java/JDBC** — the default fallback when only a JDBC driver exists for the database, or when the native driver is immature or unmaintained. Most agents still fall in this category. Native agents implement the same JSON-RPC contract and `versions.json` registration as Java agents; they ship an `agent` executable instead of `agent.jar`. If both native and Java source implementations exist for the same database, publish only the native artifact unless the Java variant has a separately registered compatibility profile, such as `oracle-legacy` / `oracle-10g`. @@ -88,6 +88,7 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed). ./gradlew shadowJar (cd drivers/oracle-go && go build -o agent .) (cd drivers/cassandra-go && go build -o agent .) +(cd drivers/iotdb && go build -o agent .) (cd drivers/kingbase-go && go build -o agent .) (cd drivers/vastbase-go && go build -o agent .) (cargo build --manifest-path drivers/tdengine/Cargo.toml --release --locked) @@ -95,7 +96,7 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed). (cd drivers/rabbitmq && go build -o agent .) ``` -Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/tdengine`, `drivers/xugu`, and `drivers/rabbitmq`. +Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/duckdb`, `drivers/iotdb`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/tdengine`, `drivers/xugu`, and `drivers/rabbitmq`. ### Local DBX Runtime Test @@ -109,7 +110,7 @@ cp agents/drivers//build/libs/*-all.jar ~/.dbx/agents/drivers/ Restart DBX or disconnect and reconnect the database so the new agent process loads the replacement JAR. -Native agents such as `cassandra`, `oracle`, `kingbase`, `tdengine`, `xugu`, and `rabbitmq` use an `agent` executable instead of `agent.jar`. TDengine builds `target/release/dbx-tdengine-driver` from `drivers/tdengine/Cargo.toml`. +Native agents such as `cassandra`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, and `rabbitmq` use an `agent` executable instead of `agent.jar`. TDengine builds `target/release/dbx-tdengine-driver` from `drivers/tdengine/Cargo.toml`. ## Versioning diff --git a/agents/README.zh-CN.md b/agents/README.zh-CN.md index 9a9acc6d66..c9f4a4e5a2 100644 --- a/agents/README.zh-CN.md +++ b/agents/README.zh-CN.md @@ -41,7 +41,7 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数 | tdengine | TDengine 2.4+ | taos-connector-rust 原生 WebSocket agent | | yashandb | 崖山 YashanDB | YashanDB JDBC | | xugu | 虚谷 XuguDB | XuguDB Go 原生 agent | -| iotdb | Apache IoTDB | IoTDB JDBC | +| iotdb | Apache IoTDB | Apache IoTDB Go Client 原生 Agent | | etcd | etcd | jetcd | | zookeeper | Apache ZooKeeper | Apache Curator | | rabbitmq | RabbitMQ | amqp091-go 原生 agent | @@ -49,7 +49,7 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数 ## 多 JRE 支持 -多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`duckdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 JRE 21 安装。 +多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`duckdb`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 JRE 21 安装。 ## JDBC 连接池 @@ -75,7 +75,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J 对于新 agent，只要存在成熟、许可证兼容的原生驱动，优先选择**原生（Go 或 Rust）驱动**而非 Java/JDBC agent。原生 agent 以单一自包含可执行文件发布，无需 JRE，可显著降低内存占用和启动时间 —— 完全避开 Java agent 即便空闲也要付出的 JVM 基线开销。 -- **原生（Go/Rust）** —— 存在可用原生驱动时首选。参考 `drivers/cassandra-go`（Apache cassandra-gocql-driver）、`drivers/duckdb`、`drivers/oracle-go`（go-ora）、`drivers/kingbase-go`（gokb）、`drivers/vastbase-go`（openGauss connector）、`drivers/tdengine`（taos-connector-rust）、`drivers/xugu` 和 `drivers/rabbitmq`（amqp091-go）。无需 JRE 下载与管理。 +- **原生（Go/Rust）** —— 存在可用原生驱动时首选。参考 `drivers/cassandra-go`（Apache cassandra-gocql-driver）、`drivers/duckdb`、`drivers/iotdb`（Apache IoTDB Go Client）、`drivers/oracle-go`（go-ora）、`drivers/kingbase-go`（gokb）、`drivers/vastbase-go`（openGauss connector）、`drivers/tdengine`（taos-connector-rust）、`drivers/xugu` 和 `drivers/rabbitmq`（amqp091-go）。无需 JRE 下载与管理。 - **Java/JDBC** —— 当某数据库只有 JDBC 驱动，或原生驱动不成熟、缺乏维护时的默认兜底方案。多数 agent 仍属此类。 原生 agent 实现与 Java agent 相同的 JSON-RPC 契约和 `versions.json` 登记；它发布的是 `agent` 可执行文件而非 `agent.jar`。若同一数据库同时保留原生和 Java 源码实现，默认只发布原生产物；只有 Java 变体以独立兼容配置登记时才同时发布，例如 `oracle-legacy` / `oracle-10g`。 @@ -88,6 +88,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J ./gradlew shadowJar (cd drivers/oracle-go && go build -o agent .) (cd drivers/cassandra-go && go build -o agent .) +(cd drivers/iotdb && go build -o agent .) (cd drivers/kingbase-go && go build -o agent .) (cd drivers/vastbase-go && go build -o agent .) (cargo build --manifest-path drivers/tdengine/Cargo.toml --release --locked) @@ -95,7 +96,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J (cd drivers/rabbitmq && go build -o agent .) ``` -产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/cassandra-go`、`drivers/duckdb`、`drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/tdengine`、`drivers/xugu` 和 `drivers/rabbitmq` 构建。 +产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/cassandra-go`、`drivers/duckdb`、`drivers/iotdb`、`drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/tdengine`、`drivers/xugu` 和 `drivers/rabbitmq` 构建。 ### 本地 DBX 运行时测试 @@ -109,7 +110,7 @@ cp agents/drivers//build/libs/*-all.jar ~/.dbx/agents/drivers/ 重启 DBX 或断开重连数据库，使新 agent 进程加载替换后的 JAR。 -`cassandra`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq` 等原生 agent 使用可执行文件而非 `agent.jar`。TDengine 从 `drivers/tdengine/Cargo.toml` 构建 `target/release/dbx-tdengine-driver`。 +`cassandra`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq` 等原生 agent 使用可执行文件而非 `agent.jar`。TDengine 从 `drivers/tdengine/Cargo.toml` 构建 `target/release/dbx-tdengine-driver`。 ## 版本管理 diff --git a/agents/build.gradle b/agents/build.gradle index 04413a0d8e..ea4fe12c61 100644 --- a/agents/build.gradle +++ b/agents/build.gradle @@ -7,7 +7,7 @@ def legacyStandaloneProjects = ['mongodb', 'kafka', 'rocketmq'] as Set def pooledJdbcProjects = [ 'access', 'bigquery', 'dameng', 'databend', 'databricks', 'db2', 'exasol', 'firebird', 'gbase8a', 'gbase8s', 'goldendb', 'h2', 'h2-legacy', 'highgo', 'hive', - 'informix', 'iotdb', 'iris', 'kylin', 'oceanbase-oracle', 'oscar', 'saphana', + 'informix', 'iris', 'kylin', 'oceanbase-oracle', 'oscar', 'saphana', 'snowflake', 'spark', 'sqlserver-legacy', 'sundb', 'teradata', 'trino', 'uxdb', 'vertica', 'yashandb' ] as Set diff --git a/agents/drivers/iotdb/README.md b/agents/drivers/iotdb/README.md new file mode 100644 index 0000000000..a420a6163e --- /dev/null +++ b/agents/drivers/iotdb/README.md @@ -0,0 +1,55 @@ +# Apache IoTDB native Agent + +This module implements the DBX Agent protocol with the official Apache IoTDB +Go client. Tree SQL is the default; set `sql_dialect=table` for the Table model. + +## Build and test + +Go 1.25 or newer is required by the pinned upstream client revision. + +```bash +go test ./... +go test -race ./... +CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o agent . +``` + +Run the live Tree/Table, paging, multi-session, and cancellation tests against +an IoTDB server with: + +```bash +DBX_IOTDB_LIVE=1 go test -race ./... -count=1 +``` + +The live tests default to `127.0.0.1:6667` with `root/root`. Override them with +`DBX_IOTDB_HOST`, `DBX_IOTDB_PORT`, `DBX_IOTDB_USER`, and +`DBX_IOTDB_PASSWORD`. Tests create and remove isolated Tree and Table databases. + +## Connection options + +DBX connection fields and `jdbc:iotdb://...` connection strings are accepted. +The following URL parameters are supported: + +- `sql_dialect=tree|table` +- `database` or `db` +- `fetch_size` +- `time_zone` +- `connect_retry_max` +- `connect_timeout_ms` +- `enable_compression=true` +- `node_urls=host1:6667,host2:6667` for cluster sessions +- `ssl=true` and `insecure_skip_verify=true` + +The standard DBX certificate fields provide CA and client certificate paths +for TLS or mTLS connections. + +## Compatibility + +- Tree sessions expose databases, devices, timeseries columns, queries, DDL, + paging, and sequential batch/transaction execution. +- Table sessions expose databases, tables, TIME/TAG/FIELD column categories, + comments, queries, DDL, paging, and database switching. +- The Agent keeps one physical IoTDB session per logical DBX session and + invalidates that session after cancellation, timeout, or connection failure. + +The historical JDBC-versus-Go driver benchmark and raw result summary are kept +under [`bench`](bench/README.md). diff --git a/agents/drivers/iotdb/bench/README.md b/agents/drivers/iotdb/bench/README.md new file mode 100644 index 0000000000..049b8747f5 --- /dev/null +++ b/agents/drivers/iotdb/bench/README.md @@ -0,0 +1,59 @@ +# IoTDB JDBC vs Go driver benchmark + +This standalone benchmark preserves the JDBC-versus-Go comparison used before +the production IoTDB Agent migrated to Go. It compares the released Apache +IoTDB `2.0.8` JDBC and Go clients against the same Tree-model server and +fixture, measuring cold process plus connection startup and warm query latency +while fully decoding every returned cell. + +The default fixture contains 10,000 rows in `root.dbx_bench.d1`. Workloads are: + +- `SHOW DATABASES` +- one timestamp point query +- a 100-row range query +- a full scan of all fixture rows + +## Start IoTDB + +For example, run the matching standalone Docker image: + +```bash +docker run --rm --name dbx-iotdb-bench -p 6667:6667 apache/iotdb:2.0.8-standalone +``` + +Wait until port `6667` is ready before starting the benchmark. + +## Run + +From the repository root: + +```bash +python3 agents/drivers/iotdb/bench/run.py \ + > /tmp/dbx-iotdb-driver-benchmark.json +``` + +The runner builds both candidates from this benchmark directory, recreates the +fixture, alternates cold-start and connected-RSS samples, then runs three warm +workload rounds with alternating candidate order. It does not build or modify +the production IoTDB Agent. + +## Configuration + +- `IOTDB_HOST`, `IOTDB_PORT`, `IOTDB_USERNAME`, `IOTDB_PASSWORD` +- `BENCH_ROWS`, default `10000` +- `BENCH_FETCH_SIZE`, default `1024` +- `BENCH_STARTUPS`, default `5` +- `BENCH_RSS_SAMPLES`, default `3` +- `BENCH_ROUNDS`, default `3` +- `BENCH_ORDER`, default `jdbc,go`; use `go,jdbc` to check order effects +- `BENCH_WARMUPS`, default `3` +- `BENCH_METADATA_ITERATIONS`, default `20` +- `BENCH_POINT_ITERATIONS`, default `100` +- `BENCH_RANGE_ITERATIONS`, default `30` +- `BENCH_SCAN_ITERATIONS`, default `5` +- `BENCH_PREPARE=false` to preserve an existing fixture +- `BENCH_SKIP_BUILD=true` to reuse existing benchmark artifacts + +Run both candidates against the same server and do not change the fixture or +fetch size between candidates. This is a client-side comparison, not an IoTDB +server throughput benchmark. diff --git a/agents/drivers/iotdb/bench/build.gradle b/agents/drivers/iotdb/bench/build.gradle new file mode 100644 index 0000000000..818899faf4 --- /dev/null +++ b/agents/drivers/iotdb/bench/build.gradle @@ -0,0 +1,31 @@ +plugins { + id 'java' + id 'com.gradleup.shadow' version '9.0.0-beta12' +} + +repositories { + mavenCentral() +} + +dependencies { + implementation 'org.apache.iotdb:iotdb-jdbc:2.0.8' + implementation 'org.slf4j:slf4j-nop:2.0.16' +} + +sourceSets { + main { + java.srcDir 'java' + } +} + +tasks.named('shadowJar') { + archiveBaseName = 'dbx-iotdb-jdbc-benchmark' + archiveClassifier = '' + manifest { + attributes('Main-Class': 'com.dbx.agent.iotdb.bench.JdbcDriverBenchmark') + } +} + +tasks.register('benchmarkJar') { + dependsOn tasks.named('shadowJar') +} diff --git a/agents/drivers/iotdb/bench/go/go.mod b/agents/drivers/iotdb/bench/go/go.mod new file mode 100644 index 0000000000..565aba2fc4 --- /dev/null +++ b/agents/drivers/iotdb/bench/go/go.mod @@ -0,0 +1,7 @@ +module github.com/t8y2/dbx/agents/drivers/iotdb/bench + +go 1.23 + +require github.com/apache/iotdb-client-go/v2 v2.0.8 + +require github.com/apache/thrift v0.15.0 // indirect diff --git a/agents/drivers/iotdb/bench/go/go.sum b/agents/drivers/iotdb/bench/go/go.sum new file mode 100644 index 0000000000..b448acb43f --- /dev/null +++ b/agents/drivers/iotdb/bench/go/go.sum @@ -0,0 +1,29 @@ +github.com/apache/iotdb-client-go/v2 v2.0.8 h1:4Z2bmnoeH9gPNz0cDSfqbvwfyY0Ocbo18vtPfjKo34k= +github.com/apache/iotdb-client-go/v2 v2.0.8/go.mod h1:keIG+1Xfw8n2NukUbrEcYnnyN9VxP8OJLYvnHsg1k2s= +github.com/apache/thrift v0.15.0 h1:aGvdaR0v1t9XLgjtBYwxcBvBOTMqClzwE26CHOgjW1Y= +github.com/apache/thrift v0.15.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/agents/drivers/iotdb/bench/go/main.go b/agents/drivers/iotdb/bench/go/main.go new file mode 100644 index 0000000000..1bec1bdbb2 --- /dev/null +++ b/agents/drivers/iotdb/bench/go/main.go @@ -0,0 +1,284 @@ +package main + +import ( + "encoding/json" + "fmt" + "os" + "sort" + "strconv" + "time" + + "github.com/apache/iotdb-client-go/v2/client" +) + +const ( + clientVersion = "2.0.8" + database = "root.dbx_bench" + device = database + ".d1" +) + +type config struct { + Mode string + Host string + Port string + Username string + Password string + Rows int + FetchSize int32 + ConnectTimeoutMS int + QueryTimeoutMS int64 + Warmups int + MetadataIterations int + PointIterations int + RangeIterations int + ScanIterations int +} + +type workload struct { + Name string + SQL string + Iterations int +} + +type workloadResult struct { + Name string `json:"name"` + Iterations int `json:"iterations"` + Rows int64 `json:"rows"` + DecodedCells int64 `json:"decoded_cells"` + MeanMS float64 `json:"mean_ms"` + P50MS float64 `json:"p50_ms"` + P95MS float64 `json:"p95_ms"` + MinMS float64 `json:"min_ms"` + MaxMS float64 `json:"max_ms"` +} + +type benchmarkResult struct { + Driver string `json:"driver"` + ClientVersion string `json:"client_version"` + ConnectMS float64 `json:"connect_ms"` + FetchSize int32 `json:"fetch_size"` + Workloads []workloadResult `json:"workloads"` +} + +type queryObservation struct { + Rows int64 + DecodedCells int64 +} + +func main() { + conf := loadConfig() + session := client.NewSession(&client.Config{ + Host: conf.Host, + Port: conf.Port, + UserName: conf.Username, + Password: conf.Password, + FetchSize: conf.FetchSize, + TimeZone: client.DefaultTimeZone, + }) + + connectStarted := time.Now() + if err := session.Open(false, conf.ConnectTimeoutMS); err != nil { + fatal(err) + } + connectMS := elapsedMS(connectStarted) + defer session.Close() + + if conf.Mode == "probe" { + observation, err := executeQuery(&session, conf, "SHOW DATABASES") + if err != nil { + fatal(err) + } + writeJSON(map[string]any{ + "driver": "go", + "client_version": clientVersion, + "connect_ms": roundMillis(connectMS), + "rows": observation.Rows, + }) + return + } + if conf.Mode == "hold" { + observation, err := executeQuery(&session, conf, "SHOW DATABASES") + if err != nil { + fatal(err) + } + writeJSON(map[string]any{ + "driver": "go", + "client_version": clientVersion, + "ready": true, + "rows": observation.Rows, + }) + time.Sleep(time.Duration(envInt("BENCH_HOLD_MS", 30_000)) * time.Millisecond) + return + } + if conf.Mode != "benchmark" { + fatal(fmt.Errorf("unsupported IOTDB_BENCH_MODE: %s", conf.Mode)) + } + + workloads := []workload{ + {Name: "show_databases", SQL: "SHOW DATABASES", Iterations: conf.MetadataIterations}, + { + Name: "point_query", + SQL: fmt.Sprintf("SELECT s1,s2,s3 FROM %s WHERE time = %d", device, max(1, conf.Rows/2)), + Iterations: conf.PointIterations, + }, + {Name: "range_100", SQL: "SELECT s1,s2,s3 FROM " + device + " LIMIT 100", Iterations: conf.RangeIterations}, + {Name: "scan_all", SQL: "SELECT s1,s2,s3 FROM " + device, Iterations: conf.ScanIterations}, + } + + results := make([]workloadResult, 0, len(workloads)) + for _, item := range workloads { + warmups := conf.Warmups + if item.Name == "scan_all" && warmups > 1 { + warmups = 1 + } + for index := 0; index < warmups; index++ { + if _, err := executeQuery(&session, conf, item.SQL); err != nil { + fatal(err) + } + } + result, err := runWorkload(&session, conf, item) + if err != nil { + fatal(err) + } + results = append(results, result) + } + + writeJSON(benchmarkResult{ + Driver: "go", + ClientVersion: clientVersion, + ConnectMS: roundMillis(connectMS), + FetchSize: conf.FetchSize, + Workloads: results, + }) +} + +func runWorkload(session *client.Session, conf config, item workload) (workloadResult, error) { + samples := make([]float64, 0, item.Iterations) + var expected queryObservation + for index := 0; index < item.Iterations; index++ { + started := time.Now() + observation, err := executeQuery(session, conf, item.SQL) + if err != nil { + return workloadResult{}, err + } + samples = append(samples, elapsedMS(started)) + if index == 0 { + expected = observation + } else if expected != observation { + return workloadResult{}, fmt.Errorf("unstable result shape for %s", item.Name) + } + } + + ordered := append([]float64(nil), samples...) + sort.Float64s(ordered) + var total float64 + for _, sample := range samples { + total += sample + } + return workloadResult{ + Name: item.Name, + Iterations: item.Iterations, + Rows: expected.Rows, + DecodedCells: expected.DecodedCells, + MeanMS: roundMillis(total / float64(len(samples))), + P50MS: roundMillis(percentile(ordered, 0.50)), + P95MS: roundMillis(percentile(ordered, 0.95)), + MinMS: roundMillis(ordered[0]), + MaxMS: roundMillis(ordered[len(ordered)-1]), + }, nil +} + +func executeQuery(session *client.Session, conf config, sql string) (queryObservation, error) { + timeout := conf.QueryTimeoutMS + dataset, err := session.ExecuteQueryStatement(sql, &timeout) + if err != nil { + return queryObservation{}, err + } + defer dataset.Close() + + columns := len(dataset.GetColumnNames()) + var observation queryObservation + for { + hasNext, nextErr := dataset.Next() + if nextErr != nil { + return queryObservation{}, nextErr + } + if !hasNext { + break + } + observation.Rows++ + for column := 1; column <= columns; column++ { + if _, valueErr := dataset.GetObjectByIndex(int32(column)); valueErr != nil { + return queryObservation{}, valueErr + } + observation.DecodedCells++ + } + } + return observation, nil +} + +func loadConfig() config { + return config{ + Mode: env("IOTDB_BENCH_MODE", "benchmark"), + Host: env("IOTDB_HOST", "127.0.0.1"), + Port: env("IOTDB_PORT", "6667"), + Username: env("IOTDB_USERNAME", "root"), + Password: env("IOTDB_PASSWORD", "root"), + Rows: envInt("BENCH_ROWS", 10_000), + FetchSize: int32(envInt("BENCH_FETCH_SIZE", 1_024)), + ConnectTimeoutMS: envInt("BENCH_CONNECT_TIMEOUT_MS", 5_000), + QueryTimeoutMS: int64(envInt("BENCH_QUERY_TIMEOUT_MS", 30_000)), + Warmups: envInt("BENCH_WARMUPS", 3), + MetadataIterations: envInt("BENCH_METADATA_ITERATIONS", 20), + PointIterations: envInt("BENCH_POINT_ITERATIONS", 100), + RangeIterations: envInt("BENCH_RANGE_ITERATIONS", 30), + ScanIterations: envInt("BENCH_SCAN_ITERATIONS", 5), + } +} + +func percentile(values []float64, fraction float64) float64 { + index := int(float64(len(values))*fraction+0.999999) - 1 + if index < 0 { + index = 0 + } + if index >= len(values) { + index = len(values) - 1 + } + return values[index] +} + +func elapsedMS(started time.Time) float64 { + return float64(time.Since(started).Microseconds()) / 1_000 +} + +func roundMillis(value float64) float64 { + return float64(int64(value*1_000+0.5)) / 1_000 +} + +func env(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + return fallback +} + +func envInt(name string, fallback int) int { + value, err := strconv.Atoi(env(name, strconv.Itoa(fallback))) + if err != nil || value <= 0 { + fatal(fmt.Errorf("%s must be a positive integer", name)) + } + return value +} + +func writeJSON(value any) { + encoder := json.NewEncoder(os.Stdout) + encoder.SetEscapeHTML(false) + if err := encoder.Encode(value); err != nil { + fatal(err) + } +} + +func fatal(err error) { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) +} diff --git a/agents/drivers/iotdb/bench/java/com/dbx/agent/iotdb/bench/JdbcDriverBenchmark.java b/agents/drivers/iotdb/bench/java/com/dbx/agent/iotdb/bench/JdbcDriverBenchmark.java new file mode 100644 index 0000000000..7c2d615fdb --- /dev/null +++ b/agents/drivers/iotdb/bench/java/com/dbx/agent/iotdb/bench/JdbcDriverBenchmark.java @@ -0,0 +1,287 @@ +package com.dbx.agent.iotdb.bench; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Locale; + +public final class JdbcDriverBenchmark { + private static final String CLIENT_VERSION = "2.0.8"; + private static final String DATABASE = "root.dbx_bench"; + private static final String DEVICE = DATABASE + ".d1"; + + private JdbcDriverBenchmark() { + } + + public static void main(String[] args) throws Exception { + Class.forName("org.apache.iotdb.jdbc.IoTDBDriver"); + Config config = Config.fromEnvironment(); + switch (config.mode()) { + case "prepare" -> prepare(config); + case "probe" -> probe(config); + case "hold" -> hold(config); + case "benchmark" -> benchmark(config); + default -> throw new IllegalArgumentException("Unsupported IOTDB_BENCH_MODE: " + config.mode()); + } + } + + private static void prepare(Config config) throws SQLException { + try (Connection connection = openConnection(config); + Statement statement = connection.createStatement()) { + try { + statement.execute("DROP DATABASE " + DATABASE); + } catch (SQLException ignored) { + } + statement.execute("CREATE DATABASE " + DATABASE); + statement.execute("CREATE TIMESERIES " + DEVICE + ".s1 WITH DATATYPE=INT64, ENCODING=RLE"); + statement.execute("CREATE TIMESERIES " + DEVICE + ".s2 WITH DATATYPE=DOUBLE, ENCODING=GORILLA"); + statement.execute("CREATE TIMESERIES " + DEVICE + ".s3 WITH DATATYPE=TEXT, ENCODING=PLAIN"); + + String insertSql = "INSERT INTO " + DEVICE + "(time,s1,s2,s3) VALUES(?,?,?,?)"; + try (PreparedStatement prepared = connection.prepareStatement(insertSql)) { + for (int row = 1; row <= config.rows(); row++) { + prepared.setLong(1, row); + prepared.setLong(2, row * 10L); + prepared.setDouble(3, row / 10.0); + prepared.setString(4, "row-" + row); + prepared.addBatch(); + if (row % config.insertBatchSize() == 0) { + prepared.executeBatch(); + } + } + if (config.rows() % config.insertBatchSize() != 0) { + prepared.executeBatch(); + } + } + } + System.out.printf(Locale.ROOT, "{\"prepared_rows\":%d,\"database\":\"%s\"}%n", config.rows(), DATABASE); + } + + private static void probe(Config config) throws SQLException { + long started = System.nanoTime(); + try (Connection connection = openConnection(config)) { + double connectMs = elapsedMillis(started); + QueryObservation observation = executeQuery(connection, config.fetchSize(), "SHOW DATABASES"); + System.out.printf( + Locale.ROOT, + "{\"driver\":\"jdbc\",\"client_version\":\"%s\",\"connect_ms\":%.3f,\"rows\":%d}%n", + CLIENT_VERSION, + connectMs, + observation.rows() + ); + } + } + + private static void hold(Config config) throws Exception { + try (Connection connection = openConnection(config)) { + QueryObservation observation = executeQuery(connection, config.fetchSize(), "SHOW DATABASES"); + System.out.printf( + Locale.ROOT, + "{\"driver\":\"jdbc\",\"client_version\":\"%s\",\"ready\":true,\"rows\":%d}%n", + CLIENT_VERSION, + observation.rows() + ); + System.out.flush(); + Thread.sleep(envInt("BENCH_HOLD_MS", 30_000)); + } + } + + private static void benchmark(Config config) throws SQLException { + long connectStarted = System.nanoTime(); + try (Connection connection = openConnection(config)) { + double connectMs = elapsedMillis(connectStarted); + List workloads = List.of( + new Workload("show_databases", "SHOW DATABASES", config.metadataIterations()), + new Workload( + "point_query", + "SELECT s1,s2,s3 FROM " + DEVICE + " WHERE time = " + Math.max(1, config.rows() / 2), + config.pointIterations() + ), + new Workload("range_100", "SELECT s1,s2,s3 FROM " + DEVICE + " LIMIT 100", config.rangeIterations()), + new Workload("scan_all", "SELECT s1,s2,s3 FROM " + DEVICE, config.scanIterations()) + ); + + List results = new ArrayList<>(); + for (Workload workload : workloads) { + int warmups = workload.name().equals("scan_all") ? Math.min(1, config.warmups()) : config.warmups(); + for (int index = 0; index < warmups; index++) { + executeQuery(connection, config.fetchSize(), workload.sql()); + } + results.add(runWorkload(connection, config.fetchSize(), workload)); + } + + System.out.println(resultJson(config, connectMs, results)); + } + } + + private static WorkloadResult runWorkload(Connection connection, int fetchSize, Workload workload) throws SQLException { + List samples = new ArrayList<>(workload.iterations()); + long expectedRows = -1; + long expectedCells = -1; + for (int index = 0; index < workload.iterations(); index++) { + long started = System.nanoTime(); + QueryObservation observation = executeQuery(connection, fetchSize, workload.sql()); + samples.add(elapsedMillis(started)); + if (expectedRows < 0) { + expectedRows = observation.rows(); + expectedCells = observation.decodedCells(); + } else if (expectedRows != observation.rows() || expectedCells != observation.decodedCells()) { + throw new IllegalStateException("Unstable result shape for " + workload.name()); + } + } + return WorkloadResult.from(workload.name(), samples, expectedRows, expectedCells); + } + + private static QueryObservation executeQuery(Connection connection, int fetchSize, String sql) throws SQLException { + long rows = 0; + long decodedCells = 0; + try (Statement statement = connection.createStatement()) { + statement.setFetchSize(fetchSize); + try (ResultSet resultSet = statement.executeQuery(sql)) { + ResultSetMetaData metadata = resultSet.getMetaData(); + int columns = metadata.getColumnCount(); + while (resultSet.next()) { + rows++; + for (int column = 1; column <= columns; column++) { + resultSet.getObject(column); + decodedCells++; + } + } + } + } + return new QueryObservation(rows, decodedCells); + } + + private static Connection openConnection(Config config) throws SQLException { + String url = "jdbc:iotdb://" + config.host() + ":" + config.port() + "/"; + return DriverManager.getConnection(url, config.username(), config.password()); + } + + private static String resultJson(Config config, double connectMs, List results) { + StringBuilder json = new StringBuilder(512); + json.append('{') + .append("\"driver\":\"jdbc\",") + .append("\"client_version\":\"").append(CLIENT_VERSION).append("\",") + .append("\"connect_ms\":").append(decimal(connectMs)).append(',') + .append("\"fetch_size\":").append(config.fetchSize()).append(',') + .append("\"workloads\":["); + for (int index = 0; index < results.size(); index++) { + if (index > 0) { + json.append(','); + } + json.append(results.get(index).toJson()); + } + return json.append("]}").toString(); + } + + private static double elapsedMillis(long started) { + return (System.nanoTime() - started) / 1_000_000.0; + } + + private static String decimal(double value) { + return String.format(Locale.ROOT, "%.3f", value); + } + + private static String env(String name, String fallback) { + String value = System.getenv(name); + return value == null || value.isBlank() ? fallback : value; + } + + private static int envInt(String name, int fallback) { + return Integer.parseInt(env(name, Integer.toString(fallback))); + } + + private record Config( + String mode, + String host, + int port, + String username, + String password, + int rows, + int fetchSize, + int insertBatchSize, + int warmups, + int metadataIterations, + int pointIterations, + int rangeIterations, + int scanIterations + ) { + private static Config fromEnvironment() { + return new Config( + env("IOTDB_BENCH_MODE", "benchmark"), + env("IOTDB_HOST", "127.0.0.1"), + envInt("IOTDB_PORT", 6667), + env("IOTDB_USERNAME", "root"), + env("IOTDB_PASSWORD", "root"), + envInt("BENCH_ROWS", 10_000), + envInt("BENCH_FETCH_SIZE", 1_024), + envInt("BENCH_INSERT_BATCH_SIZE", 500), + envInt("BENCH_WARMUPS", 3), + envInt("BENCH_METADATA_ITERATIONS", 20), + envInt("BENCH_POINT_ITERATIONS", 100), + envInt("BENCH_RANGE_ITERATIONS", 30), + envInt("BENCH_SCAN_ITERATIONS", 5) + ); + } + } + + private record Workload(String name, String sql, int iterations) { + } + + private record QueryObservation(long rows, long decodedCells) { + } + + private record WorkloadResult( + String name, + int iterations, + long rows, + long decodedCells, + double meanMs, + double p50Ms, + double p95Ms, + double minMs, + double maxMs + ) { + private static WorkloadResult from(String name, List samples, long rows, long decodedCells) { + List ordered = new ArrayList<>(samples); + Collections.sort(ordered); + double total = samples.stream().mapToDouble(Double::doubleValue).sum(); + return new WorkloadResult( + name, + samples.size(), + rows, + decodedCells, + total / samples.size(), + percentile(ordered, 0.50), + percentile(ordered, 0.95), + ordered.get(0), + ordered.get(ordered.size() - 1) + ); + } + + private static double percentile(List values, double fraction) { + int index = Math.max(0, (int) Math.ceil(values.size() * fraction) - 1); + return values.get(Math.min(values.size() - 1, index)); + } + + private String toJson() { + return "{\"name\":\"" + name + "\"" + + ",\"iterations\":" + iterations + + ",\"rows\":" + rows + + ",\"decoded_cells\":" + decodedCells + + ",\"mean_ms\":" + decimal(meanMs) + + ",\"p50_ms\":" + decimal(p50Ms) + + ",\"p95_ms\":" + decimal(p95Ms) + + ",\"min_ms\":" + decimal(minMs) + + ",\"max_ms\":" + decimal(maxMs) + + '}'; + } + } +} diff --git a/agents/drivers/iotdb/bench/results/README.md b/agents/drivers/iotdb/bench/results/README.md new file mode 100644 index 0000000000..b18f143abe --- /dev/null +++ b/agents/drivers/iotdb/bench/results/README.md @@ -0,0 +1,32 @@ +# IoTDB 2.0.8 local driver comparison + +Date: 2026-08-10 + +Environment: + +- Apple M2 Pro, macOS arm64 +- Java 21.0.7 +- Go 1.26.5 +- Apache IoTDB 2.0.8 standalone +- JDBC and Go clients both pinned to 2.0.8 +- 10,000 Tree-model rows, fetch size 1,024 +- 15 alternating cold starts, 5 connected-RSS samples, 5 alternating workload rounds + +| Metric | JDBC | Go | Go reduction | +| --- | ---: | ---: | ---: | +| Cold process + connection p50 | 172.612 ms | 9.315 ms | 94.6% | +| Connected RSS p50 | 62,368 KB | 7,008 KB | 88.8% | +| Artifact size | 25,864,492 B | 6,475,202 B | 75.0% | +| In-process connection mean | 66.061 ms | 1.932 ms | 97.1% | +| `SHOW DATABASES` mean | 0.775 ms | 0.428 ms | 44.8% | +| Point query mean | 0.813 ms | 0.595 ms | 26.8% | +| 100-row query mean | 1.437 ms | 1.010 ms | 29.7% | +| 10,000-row full decode mean | 5.455 ms | 4.941 ms | 9.4% | + +The result supports a Go Agent primarily for startup, connection, memory, and +distribution-size improvements. Warm query differences are much smaller and +shrink as server work and result decoding increase. This benchmark compares +drivers directly; it does not include the DBX JSON-RPC Agent layer or prove +feature and version compatibility parity. + +See `iotdb-2.0.8-macos-arm64.json` for all raw rounds. diff --git a/agents/drivers/iotdb/bench/results/iotdb-2.0.8-macos-arm64.json b/agents/drivers/iotdb/bench/results/iotdb-2.0.8-macos-arm64.json new file mode 100644 index 0000000000..46e11d215b --- /dev/null +++ b/agents/drivers/iotdb/bench/results/iotdb-2.0.8-macos-arm64.json @@ -0,0 +1,839 @@ +{ + "server": { + "host": "127.0.0.1", + "port": 6667 + }, + "config": { + "rows": 10000, + "fetch_size": 1024, + "warmups": 3 + }, + "artifact_bytes": { + "jdbc": 25864492, + "go": 6475202 + }, + "startup_ms": { + "jdbc": { + "count": 15, + "mean": 181.087, + "p50": 172.612, + "p95": 268.332, + "min": 160.077, + "max": 268.332 + }, + "go": { + "count": 15, + "mean": 35.459, + "p50": 9.315, + "p95": 406.199, + "min": 6.986, + "max": 406.199 + } + }, + "connected_rss_kb": { + "jdbc": { + "count": 5, + "mean": 62323.2, + "p50": 62368.0, + "p95": 62528.0, + "min": 62112.0, + "max": 62528.0 + }, + "go": { + "count": 5, + "mean": 7062.4, + "p50": 7008.0, + "p95": 7248.0, + "min": 6992.0, + "max": 7248.0 + } + }, + "results": { + "jdbc": { + "driver": "jdbc", + "client_version": "2.0.8", + "connect_ms": { + "count": 5, + "mean": 66.061, + "p50": 66.0, + "p95": 69.002, + "min": 63.535, + "max": 69.002 + }, + "workloads": { + "show_databases": { + "rows": 2, + "decoded_cells": 10, + "round_mean_ms": { + "count": 5, + "mean": 0.775, + "p50": 0.771, + "p95": 0.816, + "min": 0.733, + "max": 0.816 + }, + "round_p50_ms": { + "count": 5, + "mean": 0.78, + "p50": 0.773, + "p95": 0.833, + "min": 0.726, + "max": 0.833 + }, + "round_p95_ms": { + "count": 5, + "mean": 0.89, + "p50": 0.885, + "p95": 0.979, + "min": 0.796, + "max": 0.979 + } + }, + "point_query": { + "rows": 1, + "decoded_cells": 4, + "round_mean_ms": { + "count": 5, + "mean": 0.813, + "p50": 0.846, + "p95": 0.885, + "min": 0.722, + "max": 0.885 + }, + "round_p50_ms": { + "count": 5, + "mean": 0.77, + "p50": 0.784, + "p95": 0.847, + "min": 0.688, + "max": 0.847 + }, + "round_p95_ms": { + "count": 5, + "mean": 1.113, + "p50": 1.175, + "p95": 1.254, + "min": 0.94, + "max": 1.254 + } + }, + "range_100": { + "rows": 100, + "decoded_cells": 400, + "round_mean_ms": { + "count": 5, + "mean": 1.437, + "p50": 1.495, + "p95": 1.72, + "min": 1.205, + "max": 1.72 + }, + "round_p50_ms": { + "count": 5, + "mean": 1.35, + "p50": 1.365, + "p95": 1.644, + "min": 1.143, + "max": 1.644 + }, + "round_p95_ms": { + "count": 5, + "mean": 2.193, + "p50": 2.119, + "p95": 2.825, + "min": 1.639, + "max": 2.825 + } + }, + "scan_all": { + "rows": 10000, + "decoded_cells": 40000, + "round_mean_ms": { + "count": 5, + "mean": 5.455, + "p50": 5.476, + "p95": 5.715, + "min": 5.102, + "max": 5.715 + }, + "round_p50_ms": { + "count": 5, + "mean": 5.146, + "p50": 5.254, + "p95": 5.351, + "min": 4.701, + "max": 5.351 + }, + "round_p95_ms": { + "count": 5, + "mean": 6.962, + "p50": 7.04, + "p95": 7.809, + "min": 5.891, + "max": 7.809 + } + } + }, + "rounds": [ + { + "driver": "jdbc", + "client_version": "2.0.8", + "connect_ms": 69.002, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.755, + "p50_ms": 0.753, + "p95_ms": 0.885, + "min_ms": 0.6, + "max_ms": 0.988 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.885, + "p50_ms": 0.847, + "p95_ms": 1.254, + "min_ms": 0.54, + "max_ms": 1.625 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.72, + "p50_ms": 1.644, + "p95_ms": 2.641, + "min_ms": 1.173, + "max_ms": 2.729 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.295, + "p50_ms": 5.254, + "p95_ms": 5.891, + "min_ms": 4.579, + "max_ms": 5.891 + } + ] + }, + { + "driver": "jdbc", + "client_version": "2.0.8", + "connect_ms": 67.177, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.816, + "p50_ms": 0.833, + "p95_ms": 0.979, + "min_ms": 0.62, + "max_ms": 1.151 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.865, + "p50_ms": 0.806, + "p95_ms": 1.223, + "min_ms": 0.573, + "max_ms": 2.011 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.495, + "p50_ms": 1.43, + "p95_ms": 2.825, + "min_ms": 0.972, + "max_ms": 3.059 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.689, + "p50_ms": 5.132, + "p95_ms": 6.746, + "min_ms": 5.001, + "max_ms": 6.746 + } + ] + }, + { + "driver": "jdbc", + "client_version": "2.0.8", + "connect_ms": 66.0, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.771, + "p50_ms": 0.773, + "p95_ms": 0.862, + "min_ms": 0.6, + "max_ms": 0.982 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.846, + "p50_ms": 0.784, + "p95_ms": 1.175, + "min_ms": 0.589, + "max_ms": 2.059 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.265, + "p50_ms": 1.166, + "p95_ms": 1.739, + "min_ms": 0.974, + "max_ms": 2.395 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.715, + "p50_ms": 5.351, + "p95_ms": 7.809, + "min_ms": 4.569, + "max_ms": 7.809 + } + ] + }, + { + "driver": "jdbc", + "client_version": "2.0.8", + "connect_ms": 63.535, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.798, + "p50_ms": 0.813, + "p95_ms": 0.928, + "min_ms": 0.618, + "max_ms": 0.963 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.722, + "p50_ms": 0.688, + "p95_ms": 0.94, + "min_ms": 0.567, + "max_ms": 1.719 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.205, + "p50_ms": 1.143, + "p95_ms": 1.639, + "min_ms": 0.94, + "max_ms": 2.42 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.102, + "p50_ms": 4.701, + "p95_ms": 7.04, + "min_ms": 4.094, + "max_ms": 7.04 + } + ] + }, + { + "driver": "jdbc", + "client_version": "2.0.8", + "connect_ms": 64.589, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.733, + "p50_ms": 0.726, + "p95_ms": 0.796, + "min_ms": 0.641, + "max_ms": 0.914 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.749, + "p50_ms": 0.723, + "p95_ms": 0.973, + "min_ms": 0.549, + "max_ms": 1.207 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.499, + "p50_ms": 1.365, + "p95_ms": 2.119, + "min_ms": 1.009, + "max_ms": 2.608 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.476, + "p50_ms": 5.292, + "p95_ms": 7.324, + "min_ms": 4.345, + "max_ms": 7.324 + } + ] + } + ] + }, + "go": { + "driver": "go", + "client_version": "2.0.8", + "connect_ms": { + "count": 5, + "mean": 1.932, + "p50": 1.9, + "p95": 2.455, + "min": 1.575, + "max": 2.455 + }, + "workloads": { + "show_databases": { + "rows": 2, + "decoded_cells": 10, + "round_mean_ms": { + "count": 5, + "mean": 0.428, + "p50": 0.417, + "p95": 0.478, + "min": 0.388, + "max": 0.478 + }, + "round_p50_ms": { + "count": 5, + "mean": 0.408, + "p50": 0.398, + "p95": 0.437, + "min": 0.386, + "max": 0.437 + }, + "round_p95_ms": { + "count": 5, + "mean": 0.543, + "p50": 0.552, + "p95": 0.619, + "min": 0.43, + "max": 0.619 + } + }, + "point_query": { + "rows": 1, + "decoded_cells": 4, + "round_mean_ms": { + "count": 5, + "mean": 0.595, + "p50": 0.565, + "p95": 0.723, + "min": 0.51, + "max": 0.723 + }, + "round_p50_ms": { + "count": 5, + "mean": 0.567, + "p50": 0.538, + "p95": 0.702, + "min": 0.488, + "max": 0.702 + }, + "round_p95_ms": { + "count": 5, + "mean": 0.785, + "p50": 0.788, + "p95": 0.929, + "min": 0.639, + "max": 0.929 + } + }, + "range_100": { + "rows": 100, + "decoded_cells": 400, + "round_mean_ms": { + "count": 5, + "mean": 1.01, + "p50": 1.009, + "p95": 1.061, + "min": 0.94, + "max": 1.061 + }, + "round_p50_ms": { + "count": 5, + "mean": 0.977, + "p50": 0.995, + "p95": 1.023, + "min": 0.913, + "max": 1.023 + }, + "round_p95_ms": { + "count": 5, + "mean": 1.319, + "p50": 1.304, + "p95": 1.749, + "min": 1.094, + "max": 1.749 + } + }, + "scan_all": { + "rows": 10000, + "decoded_cells": 40000, + "round_mean_ms": { + "count": 5, + "mean": 4.941, + "p50": 4.804, + "p95": 5.474, + "min": 4.624, + "max": 5.474 + }, + "round_p50_ms": { + "count": 5, + "mean": 4.858, + "p50": 4.838, + "p95": 5.15, + "min": 4.567, + "max": 5.15 + }, + "round_p95_ms": { + "count": 5, + "mean": 5.466, + "p50": 5.076, + "p95": 7.106, + "min": 4.858, + "max": 7.106 + } + } + }, + "rounds": [ + { + "driver": "go", + "client_version": "2.0.8", + "connect_ms": 2.455, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.478, + "p50_ms": 0.434, + "p95_ms": 0.619, + "min_ms": 0.382, + "max_ms": 0.638 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.723, + "p50_ms": 0.702, + "p95_ms": 0.929, + "min_ms": 0.498, + "max_ms": 1.183 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.061, + "p50_ms": 1.021, + "p95_ms": 1.323, + "min_ms": 0.901, + "max_ms": 1.38 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.474, + "p50_ms": 5.129, + "p95_ms": 7.106, + "min_ms": 4.77, + "max_ms": 7.106 + } + ] + }, + { + "driver": "go", + "client_version": "2.0.8", + "connect_ms": 1.575, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.449, + "p50_ms": 0.437, + "p95_ms": 0.546, + "min_ms": 0.364, + "max_ms": 0.561 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.651, + "p50_ms": 0.616, + "p95_ms": 0.884, + "min_ms": 0.481, + "max_ms": 1.035 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.046, + "p50_ms": 1.023, + "p95_ms": 1.304, + "min_ms": 0.882, + "max_ms": 1.401 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 5.157, + "p50_ms": 5.15, + "p95_ms": 5.39, + "min_ms": 4.891, + "max_ms": 5.39 + } + ] + }, + { + "driver": "go", + "client_version": "2.0.8", + "connect_ms": 1.9, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.417, + "p50_ms": 0.398, + "p95_ms": 0.569, + "min_ms": 0.343, + "max_ms": 0.585 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.565, + "p50_ms": 0.538, + "p95_ms": 0.788, + "min_ms": 0.404, + "max_ms": 0.86 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 0.94, + "p50_ms": 0.932, + "p95_ms": 1.094, + "min_ms": 0.822, + "max_ms": 1.166 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 4.647, + "p50_ms": 4.606, + "p95_ms": 4.858, + "min_ms": 4.48, + "max_ms": 4.858 + } + ] + }, + { + "driver": "go", + "client_version": "2.0.8", + "connect_ms": 1.707, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.388, + "p50_ms": 0.387, + "p95_ms": 0.43, + "min_ms": 0.328, + "max_ms": 0.433 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.525, + "p50_ms": 0.493, + "p95_ms": 0.685, + "min_ms": 0.423, + "max_ms": 1.744 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 1.009, + "p50_ms": 0.913, + "p95_ms": 1.749, + "min_ms": 0.829, + "max_ms": 2.153 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 4.624, + "p50_ms": 4.567, + "p95_ms": 4.901, + "min_ms": 4.413, + "max_ms": 4.901 + } + ] + }, + { + "driver": "go", + "client_version": "2.0.8", + "connect_ms": 2.024, + "fetch_size": 1024, + "workloads": [ + { + "name": "show_databases", + "iterations": 20, + "rows": 2, + "decoded_cells": 10, + "mean_ms": 0.409, + "p50_ms": 0.386, + "p95_ms": 0.552, + "min_ms": 0.321, + "max_ms": 0.554 + }, + { + "name": "point_query", + "iterations": 100, + "rows": 1, + "decoded_cells": 4, + "mean_ms": 0.51, + "p50_ms": 0.488, + "p95_ms": 0.639, + "min_ms": 0.355, + "max_ms": 1.366 + }, + { + "name": "range_100", + "iterations": 30, + "rows": 100, + "decoded_cells": 400, + "mean_ms": 0.994, + "p50_ms": 0.995, + "p95_ms": 1.125, + "min_ms": 0.891, + "max_ms": 1.197 + }, + { + "name": "scan_all", + "iterations": 5, + "rows": 10000, + "decoded_cells": 40000, + "mean_ms": 4.804, + "p50_ms": 4.838, + "p95_ms": 5.076, + "min_ms": 4.501, + "max_ms": 5.076 + } + ] + } + ] + } + }, + "environment": { + "date": "2026-08-10", + "machine": "Apple M2 Pro", + "os": "macOS Darwin 24.5.0 arm64", + "java": "21.0.7", + "go": "1.26.5", + "iotdb_server": "2.0.8 standalone", + "jdbc_client": "2.0.8", + "go_client": "2.0.8" + } +} diff --git a/agents/drivers/iotdb/bench/run.py b/agents/drivers/iotdb/bench/run.py new file mode 100644 index 0000000000..001f87a069 --- /dev/null +++ b/agents/drivers/iotdb/bench/run.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3 +import json +import os +import socket +import statistics +import subprocess +import sys +import threading +import time +from pathlib import Path + + +BENCH_DIR = Path(__file__).resolve().parent +IOTDB_PROJECT = BENCH_DIR.parent +AGENTS_DIR = IOTDB_PROJECT.parents[1] +JDBC_JAR = BENCH_DIR / "build" / "libs" / "dbx-iotdb-jdbc-benchmark.jar" +GO_BINARY = BENCH_DIR / "build" / "iotdb-go-benchmark" + + +def main() -> None: + environment = os.environ.copy() + wait_for_server(environment) + if not env_bool("BENCH_SKIP_BUILD", False): + build_candidates() + ensure_artifacts() + if env_bool("BENCH_PREPARE", True): + run_json(["java", "-jar", str(JDBC_JAR)], environment | {"IOTDB_BENCH_MODE": "prepare"}) + + commands = { + "jdbc": ["java", "-jar", str(JDBC_JAR)], + "go": [str(GO_BINARY)], + } + names = benchmark_order(commands) + startup_samples = {name: [] for name in commands} + startup_iterations = env_int("BENCH_STARTUPS", 5) + for iteration in range(startup_iterations): + order = names if iteration % 2 == 0 else list(reversed(names)) + for name in order: + started = time.perf_counter() + run_json(commands[name], environment | {"IOTDB_BENCH_MODE": "probe"}) + startup_samples[name].append((time.perf_counter() - started) * 1_000) + + rss_samples = {name: [] for name in commands} + for iteration in range(env_int("BENCH_RSS_SAMPLES", 3)): + order = names if iteration % 2 == 0 else list(reversed(names)) + for name in order: + rss_samples[name].append(measure_rss(commands[name], environment)) + + rounds = {name: [] for name in commands} + for iteration in range(env_int("BENCH_ROUNDS", 3)): + order = names if iteration % 2 == 0 else list(reversed(names)) + for name in order: + rounds[name].append(run_json(commands[name], environment | {"IOTDB_BENCH_MODE": "benchmark"})) + + output = { + "server": { + "host": environment.get("IOTDB_HOST", "127.0.0.1"), + "port": int(environment.get("IOTDB_PORT", "6667")), + }, + "config": { + "rows": env_int("BENCH_ROWS", 10_000), + "fetch_size": env_int("BENCH_FETCH_SIZE", 1_024), + "warmups": env_int("BENCH_WARMUPS", 3), + }, + "artifact_bytes": { + "jdbc": JDBC_JAR.stat().st_size, + "go": GO_BINARY.stat().st_size, + }, + "startup_ms": {name: summarize(samples) for name, samples in startup_samples.items()}, + "connected_rss_kb": {name: summarize(samples) for name, samples in rss_samples.items()}, + "results": {name: summarize_rounds(values) for name, values in rounds.items()}, + } + print(json.dumps(output, ensure_ascii=False, indent=2)) + + +def build_candidates() -> None: + subprocess.run( + [str(AGENTS_DIR / "gradlew"), "-p", str(BENCH_DIR), "benchmarkJar", "--console=plain"], + check=True, + stdout=sys.stderr, + stderr=sys.stderr, + ) + GO_BINARY.parent.mkdir(parents=True, exist_ok=True) + subprocess.run( + ["go", "build", "-o", str(GO_BINARY), "."], + cwd=BENCH_DIR / "go", + check=True, + stdout=sys.stderr, + stderr=sys.stderr, + ) + + +def wait_for_server(environment: dict[str, str]) -> None: + host = environment.get("IOTDB_HOST", "127.0.0.1") + port = int(environment.get("IOTDB_PORT", "6667")) + deadline = time.monotonic() + env_float("BENCH_SERVER_TIMEOUT", 60.0) + while time.monotonic() < deadline: + try: + with socket.create_connection((host, port), timeout=1): + return + except OSError: + time.sleep(0.5) + raise TimeoutError(f"IoTDB is not reachable at {host}:{port}") + + +def run_json(command: list[str], environment: dict[str, str]) -> dict: + completed = subprocess.run( + command, + env=environment, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + timeout=env_float("BENCH_COMMAND_TIMEOUT", 180.0), + ) + if completed.returncode != 0: + raise RuntimeError( + f"command failed ({' '.join(command)}):\nstdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + for line in reversed(completed.stdout.splitlines()): + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(value, dict): + return value + raise RuntimeError(f"command returned no JSON ({' '.join(command)}):\n{completed.stdout}") + + +def measure_rss(command: list[str], environment: dict[str, str]) -> float: + process = subprocess.Popen( + command, + env=environment | {"IOTDB_BENCH_MODE": "hold"}, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stderr_lines: list[str] = [] + + def drain_stderr() -> None: + assert process.stderr is not None + stderr_lines.extend(line.rstrip() for line in process.stderr) + + threading.Thread(target=drain_stderr, daemon=True).start() + try: + assert process.stdout is not None + deadline = time.monotonic() + env_float("BENCH_READY_TIMEOUT", 30.0) + while time.monotonic() < deadline: + line = process.stdout.readline() + if line == "" and process.poll() is not None: + raise RuntimeError(f"RSS probe exited early: {'; '.join(stderr_lines)}") + try: + payload = json.loads(line) + except json.JSONDecodeError: + continue + if payload.get("ready") is True: + break + else: + raise TimeoutError(f"RSS probe timed out: {'; '.join(stderr_lines)}") + time.sleep(0.2) + completed = subprocess.run( + ["ps", "-o", "rss=", "-p", str(process.pid)], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + ) + return float(completed.stdout.strip()) + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +def summarize(samples: list[float]) -> dict: + ordered = sorted(samples) + return { + "count": len(samples), + "mean": round(statistics.mean(samples), 3), + "p50": round(percentile(ordered, 0.50), 3), + "p95": round(percentile(ordered, 0.95), 3), + "min": round(ordered[0], 3), + "max": round(ordered[-1], 3), + } + + +def summarize_rounds(rounds: list[dict]) -> dict: + workloads: dict[str, list[dict]] = {} + for result in rounds: + for workload in result["workloads"]: + workloads.setdefault(workload["name"], []).append(workload) + return { + "driver": rounds[0]["driver"], + "client_version": rounds[0]["client_version"], + "connect_ms": summarize([result["connect_ms"] for result in rounds]), + "workloads": { + name: { + "rows": values[0]["rows"], + "decoded_cells": values[0]["decoded_cells"], + "round_mean_ms": summarize([value["mean_ms"] for value in values]), + "round_p50_ms": summarize([value["p50_ms"] for value in values]), + "round_p95_ms": summarize([value["p95_ms"] for value in values]), + } + for name, values in workloads.items() + }, + "rounds": rounds, + } + + +def percentile(values: list[float], fraction: float) -> float: + index = max(0, min(len(values) - 1, int(len(values) * fraction + 0.999999) - 1)) + return values[index] + + +def ensure_artifacts() -> None: + for path in (JDBC_JAR, GO_BINARY): + if not path.is_file(): + raise FileNotFoundError(path) + + +def benchmark_order(commands: dict[str, list[str]]) -> list[str]: + names = [name.strip() for name in os.getenv("BENCH_ORDER", "jdbc,go").split(",") if name.strip()] + if len(names) != len(commands) or set(names) != set(commands): + raise ValueError(f"BENCH_ORDER must contain exactly: {','.join(commands)}") + return names + + +def env_int(name: str, fallback: int) -> int: + value = int(os.getenv(name, str(fallback))) + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def env_float(name: str, fallback: float) -> float: + value = float(os.getenv(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 + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +if __name__ == "__main__": + main() diff --git a/agents/drivers/iotdb/bench/settings.gradle b/agents/drivers/iotdb/bench/settings.gradle new file mode 100644 index 0000000000..a5ba48564f --- /dev/null +++ b/agents/drivers/iotdb/bench/settings.gradle @@ -0,0 +1 @@ +rootProject.name = 'dbx-iotdb-driver-benchmark' diff --git a/agents/drivers/iotdb/build.gradle b/agents/drivers/iotdb/build.gradle deleted file mode 100644 index adaf4bf828..0000000000 --- a/agents/drivers/iotdb/build.gradle +++ /dev/null @@ -1,10 +0,0 @@ -dependencies { - implementation 'org.apache.iotdb:iotdb-jdbc:2.0.8' - implementation 'org.slf4j:slf4j-nop:2.0.16' -} - -tasks.named('jar') { - manifest { - attributes('Agent-Label': 'Apache IoTDB', 'Main-Class': 'com.dbx.agent.iotdb.IoTDBAgent') - } -} diff --git a/agents/drivers/iotdb/driver.go b/agents/drivers/iotdb/driver.go new file mode 100644 index 0000000000..504ab0e54f --- /dev/null +++ b/agents/drivers/iotdb/driver.go @@ -0,0 +1,333 @@ +package main + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "net" + "net/url" + "strconv" + "strings" + "time" + + "github.com/apache/iotdb-client-go/v2/client" +) + +const iotdbGoClientVersion = "2.0.9-0.20260807074554-e59fc7f55df1" + +type iotdbSession interface { + ExecuteQueryStatement(context.Context, string, *int64) (*client.SessionDataSet, error) + ExecuteNonQueryStatement(context.Context, string) error + Close() error +} + +type nativeIoTDBSession struct { + session *client.Session +} + +func (s *nativeIoTDBSession) ExecuteQueryStatement(ctx context.Context, sql string, _ *int64) (*client.SessionDataSet, error) { + return executeIoTDBStatement(ctx, s.session, sql) +} + +func (s *nativeIoTDBSession) ExecuteNonQueryStatement(ctx context.Context, sql string) error { + dataset, err := executeIoTDBStatement(ctx, s.session, sql) + if err != nil { + return err + } + if dataset != nil { + return dataset.Close() + } + return nil +} + +func (s *nativeIoTDBSession) Close() error { + return s.session.Close() +} + +func executeIoTDBStatement(ctx context.Context, session *client.Session, sql string) (dataset *client.SessionDataSet, err error) { + defer func() { + if recovered := recover(); recovered != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + dataset = nil + err = ctxErr + return + } + dataset = nil + err = fmt.Errorf("IoTDB Go client panicked while executing statement: %v", recovered) + } + }() + return session.ExecuteStatementWithContext(ctx, sql) +} + +type connectionConfig struct { + Host string + Port int + NodeURLs []string + Username string + Password string + Database string + Dialect string + FetchSize int32 + TimeZone string + ConnectRetryMax int + ConnectTimeoutMS int + EnableCompression bool + TLSConfig *client.TLSConfig + TLSInsecureSkipVerify bool +} + +type sessionClient struct { + session iotdbSession + dialect string +} + +func parseConnectionConfig(params connectParams) (connectionConfig, error) { + config := connectionConfig{ + Host: strings.TrimSpace(params.Host), + Port: params.Port, + Username: params.Username, + Password: params.Password, + Database: strings.TrimSpace(params.Database), + Dialect: client.TreeSqlDialect, + FetchSize: defaultFetchSize, + TimeZone: client.DefaultTimeZone, + ConnectRetryMax: client.DefaultConnectRetryMax, + ConnectTimeoutMS: int(defaultConnectTimeout / time.Millisecond), + } + if config.Host == "" { + config.Host = "127.0.0.1" + } + if config.Port <= 0 { + config.Port = defaultIoTDBPort + } + if config.Username == "" { + config.Username = "root" + } + if config.Password == "" { + config.Password = "root" + } + + query := url.Values{} + if raw := strings.TrimSpace(params.ConnectionString); raw != "" { + normalized := strings.TrimPrefix(raw, "jdbc:") + parsed, err := url.Parse(normalized) + if err != nil { + return connectionConfig{}, fmt.Errorf("parse IoTDB connection string: %w", err) + } + if parsed.Scheme != "" && !strings.EqualFold(parsed.Scheme, "iotdb") { + return connectionConfig{}, fmt.Errorf("unsupported IoTDB connection scheme: %s", parsed.Scheme) + } + if parsed.Hostname() != "" { + config.Host = parsed.Hostname() + } + if parsed.Port() != "" { + port, err := strconv.Atoi(parsed.Port()) + if err != nil || port <= 0 { + return connectionConfig{}, fmt.Errorf("invalid IoTDB port: %s", parsed.Port()) + } + config.Port = port + } + if parsed.User != nil { + if username := parsed.User.Username(); username != "" { + config.Username = username + } + if password, ok := parsed.User.Password(); ok { + config.Password = password + } + } + if database := strings.Trim(strings.TrimSpace(parsed.Path), "/"); database != "" { + config.Database = database + } + query = parsed.Query() + } + if raw := strings.TrimSpace(params.URLParams); raw != "" { + values, err := url.ParseQuery(strings.TrimPrefix(raw, "?")) + if err != nil { + return connectionConfig{}, fmt.Errorf("parse IoTDB URL parameters: %w", err) + } + for key, entries := range values { + query[key] = entries + } + } + + if dialect := strings.ToLower(strings.TrimSpace(firstQueryValue(query, "sql_dialect", "dialect"))); dialect != "" { + switch dialect { + case client.TreeSqlDialect, client.TableSqlDialect: + config.Dialect = dialect + default: + return connectionConfig{}, fmt.Errorf("unsupported IoTDB SQL dialect: %s", dialect) + } + } + if database := strings.TrimSpace(firstQueryValue(query, "database", "db")); database != "" { + config.Database = database + } + if value := firstQueryValue(query, "fetch_size", "fetchSize"); value != "" { + parsed, err := positiveInt(value, "fetch_size") + if err != nil { + return connectionConfig{}, err + } + config.FetchSize = int32(parsed) + } + if value := firstQueryValue(query, "time_zone", "timezone", "zone_id"); value != "" { + config.TimeZone = value + } + if value := firstQueryValue(query, "connect_retry_max", "connectRetryMax"); value != "" { + parsed, err := positiveInt(value, "connect_retry_max") + if err != nil { + return connectionConfig{}, err + } + config.ConnectRetryMax = parsed + } + if value := firstQueryValue(query, "connect_timeout_ms", "connection_timeout_ms"); value != "" { + parsed, err := positiveInt(value, "connect_timeout_ms") + if err != nil { + return connectionConfig{}, err + } + config.ConnectTimeoutMS = parsed + } + config.EnableCompression = queryBool(query, "enable_compression", "rpc_compression") + config.NodeURLs = parseNodeURLs(firstQueryValue(query, "node_urls", "nodes")) + if len(config.NodeURLs) == 0 { + config.NodeURLs = []string{net.JoinHostPort(config.Host, strconv.Itoa(config.Port))} + } + + tlsEnabled := params.SSL || queryBool(query, "ssl", "useSSL", "use_ssl", "tls") + if tlsEnabled { + config.TLSInsecureSkipVerify = queryBool(query, "insecure_skip_verify", "tls_insecure_skip_verify") + config.TLSConfig = &client.TLSConfig{ + Config: &tls.Config{ + ServerName: config.Host, + MinVersion: tls.VersionTLS12, + InsecureSkipVerify: config.TLSInsecureSkipVerify, + }, + CAFile: strings.TrimSpace(params.CACertPath), + CertFile: strings.TrimSpace(params.ClientCertPath), + KeyFile: strings.TrimSpace(params.ClientKeyPath), + } + if (config.TLSConfig.CertFile == "") != (config.TLSConfig.KeyFile == "") { + return connectionConfig{}, errors.New("both client_cert_path and client_key_path are required for IoTDB mTLS") + } + } + return config, nil +} + +func newSessionClient(config connectionConfig) (*sessionClient, error) { + var session client.Session + var err error + database := config.Database + if config.Dialect == client.TableSqlDialect { + database = "" + } + if len(config.NodeURLs) > 1 { + session, err = client.NewClusterSession(&client.ClusterConfig{ + NodeUrls: config.NodeURLs, + UserName: config.Username, + Password: config.Password, + FetchSize: config.FetchSize, + TimeZone: config.TimeZone, + ConnectRetryMax: config.ConnectRetryMax, + Database: database, + TLSConfig: config.TLSConfig, + }) + if err == nil { + err = session.OpenCluster(config.EnableCompression) + } + } else { + session = client.NewSession(&client.Config{ + Host: config.Host, + Port: strconv.Itoa(config.Port), + UserName: config.Username, + Password: config.Password, + FetchSize: config.FetchSize, + TimeZone: config.TimeZone, + ConnectRetryMax: config.ConnectRetryMax, + Database: database, + TLSConfig: config.TLSConfig, + }) + err = session.Open(config.EnableCompression, config.ConnectTimeoutMS) + } + if err != nil { + return nil, err + } + wrapped := &nativeIoTDBSession{session: &session} + connected := &sessionClient{session: wrapped, dialect: config.Dialect} + if config.Dialect == client.TableSqlDialect { + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(config.ConnectTimeoutMS)*time.Millisecond) + defer cancel() + if err := wrapped.ExecuteNonQueryStatement(ctx, "SET SQL_DIALECT=TABLE"); err != nil { + _ = wrapped.Close() + return nil, err + } + if strings.TrimSpace(config.Database) != "" { + if err := wrapped.ExecuteNonQueryStatement(ctx, "USE "+quoteTableIdentifier(config.Database)); err != nil { + _ = wrapped.Close() + return nil, err + } + } + } + return connected, nil +} + +func (s *sessionClient) Close() error { + if s == nil || s.session == nil { + return nil + } + return s.session.Close() +} + +func (s *server) ensureClient() (*sessionClient, error) { + s.clientMu.Lock() + defer s.clientMu.Unlock() + if s.client != nil { + return s.client, nil + } + connected, err := newSessionClient(s.config) + if err != nil { + return nil, err + } + s.client = connected + return connected, nil +} + +func (s *server) invalidateClient(target *sessionClient) { + s.clientMu.Lock() + if s.client == target { + s.client = nil + } + s.clientMu.Unlock() + _ = target.Close() +} + +func firstQueryValue(values url.Values, keys ...string) string { + for _, key := range keys { + if value := strings.TrimSpace(values.Get(key)); value != "" { + return value + } + } + return "" +} + +func queryBool(values url.Values, keys ...string) bool { + value := strings.ToLower(firstQueryValue(values, keys...)) + return value == "1" || value == "true" || value == "yes" || value == "on" +} + +func positiveInt(value, name string) (int, error) { + parsed, err := strconv.Atoi(value) + if err != nil || parsed <= 0 { + return 0, fmt.Errorf("%s must be a positive integer", name) + } + return parsed, nil +} + +func parseNodeURLs(value string) []string { + parts := strings.FieldsFunc(value, func(char rune) bool { return char == ',' || char == ';' }) + result := make([]string, 0, len(parts)) + for _, part := range parts { + if normalized := strings.TrimSpace(part); normalized != "" { + result = append(result, normalized) + } + } + return result +} diff --git a/agents/drivers/iotdb/go.mod b/agents/drivers/iotdb/go.mod new file mode 100644 index 0000000000..2c9ea79d5b --- /dev/null +++ b/agents/drivers/iotdb/go.mod @@ -0,0 +1,7 @@ +module github.com/t8y2/dbx/agents/drivers/iotdb + +go 1.25 + +require github.com/apache/iotdb-client-go/v2 v2.0.9-0.20260807074554-e59fc7f55df1 + +require github.com/apache/thrift v0.24.0 // indirect diff --git a/agents/drivers/iotdb/go.sum b/agents/drivers/iotdb/go.sum new file mode 100644 index 0000000000..9257c861bd --- /dev/null +++ b/agents/drivers/iotdb/go.sum @@ -0,0 +1,4 @@ +github.com/apache/iotdb-client-go/v2 v2.0.9-0.20260807074554-e59fc7f55df1 h1:UAfvet8VA2ZUCsIBusSqNDDecFGot8uDshegeDdLBRo= +github.com/apache/iotdb-client-go/v2 v2.0.9-0.20260807074554-e59fc7f55df1/go.mod h1:Ef61z8w2tRnuAcQKESFfiH03kDvfQniVPUgOXCxWnNw= +github.com/apache/thrift v0.24.0 h1:zy31L1a49QTNB2bG1BBfMXol3yJrTH975G3pPubQVLQ= +github.com/apache/thrift v0.24.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= diff --git a/agents/drivers/iotdb/integration_test.go b/agents/drivers/iotdb/integration_test.go new file mode 100644 index 0000000000..8e92b12059 --- /dev/null +++ b/agents/drivers/iotdb/integration_test.go @@ -0,0 +1,331 @@ +package main + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "strconv" + "strings" + "sync" + "testing" + "time" +) + +func TestLiveIoTDBAgentTreeAndTable(t *testing.T) { + if os.Getenv("DBX_IOTDB_LIVE") != "1" { + t.Skip("set DBX_IOTDB_LIVE=1 to run against a real IoTDB server") + } + suffix := strconv.Itoa(os.Getpid()) + treeDatabase := "root.dbx_go_agent_" + suffix + treeDevice := treeDatabase + ".d1" + treeServer := liveIoTDBServer(t, connectParams{Database: treeDatabase}) + defer treeServer.disconnect() + defer func() { _ = treeServer.executeNonQuery("DELETE DATABASE "+treeDatabase, "", 0) }() + + mustExecuteNonQuery(t, treeServer, "CREATE DATABASE "+treeDatabase, "") + mustExecuteNonQuery(t, treeServer, "CREATE TIMESERIES "+treeDevice+".s1 WITH DATATYPE=INT64, ENCODING=RLE", "") + for index := 1; index <= 4; index++ { + mustExecuteNonQuery(t, treeServer, fmt.Sprintf("INSERT INTO %s(time,s1) VALUES(%d,%d)", treeDevice, index, index*10), "") + } + + tables, err := treeServer.listTables(treeDatabase, metadataListConstraints{}) + if err != nil || len(tables) != 1 || tables[0].Name != "d1" { + t.Fatalf("tree listTables() = %#v, %v", tables, err) + } + columns, err := treeServer.getColumns(treeDatabase, "d1") + if err != nil || len(columns) != 1 || columns[0].Name != "s1" || columns[0].DataType != "INT64" { + t.Fatalf("tree getColumns() = %#v, %v", columns, err) + } + page, err := treeServer.executeQueryPage(queryOptions{SQL: "SELECT * FROM " + treeDevice, MaxRows: 3}, 2) + if err != nil || len(page.Rows) != 2 || !page.HasMore || page.SessionID == nil { + t.Fatalf("tree first page = %#v, %v", page, err) + } + lastPage, err := treeServer.fetchQueryPage(*page.SessionID, 2) + if err != nil || len(lastPage.Rows) != 1 || lastPage.HasMore || !lastPage.Truncated { + t.Fatalf("tree final page = %#v, %v", lastPage, err) + } + ddl, err := treeServer.getTableDDL(treeDatabase, "d1") + if err != nil || !strings.Contains(ddl, "CREATE TIMESERIES "+treeDevice+".s1") { + t.Fatalf("tree DDL = %q, %v", ddl, err) + } + + tableDatabase := "dbx_go_agent_" + suffix + tableParams := liveIoTDBParams() + tableParams.URLParams = "sql_dialect=table&time_zone=Asia%2FShanghai" + tableServer := liveIoTDBServer(t, tableParams) + defer tableServer.disconnect() + defer func() { _ = tableServer.executeNonQuery("DROP DATABASE "+quoteTableIdentifier(tableDatabase), "", 0) }() + + mustExecuteNonQuery(t, tableServer, "CREATE DATABASE "+quoteTableIdentifier(tableDatabase), "") + mustExecuteNonQuery(t, tableServer, + "CREATE TABLE "+quoteTableIdentifier(tableDatabase)+"."+quoteTableIdentifier("d1")+ + " (time TIMESTAMP TIME, device STRING TAG, s1 INT64 FIELD COMMENT 'value') COMMENT 'DBX live test' WITH (TTL='INF')", + tableDatabase, + ) + mustExecuteNonQuery(t, tableServer, + "INSERT INTO "+quoteTableIdentifier("d1")+"(time,device,s1) VALUES(1,'a',10),(2,'a',20)", + tableDatabase, + ) + tableColumns, err := tableServer.getColumns(tableDatabase, "d1") + if err != nil || len(tableColumns) != 3 || !tableColumns[0].IsPrimaryKey || !tableColumns[1].IsPrimaryKey { + t.Fatalf("table getColumns() = %#v, %v", tableColumns, err) + } + comment, err := tableServer.getTableComment(tableDatabase, "d1") + if err != nil || comment == nil || *comment != "DBX live test" { + t.Fatalf("table comment = %#v, %v", comment, err) + } + result, err := tableServer.executeQuery(queryOptions{SQL: "SELECT * FROM d1 ORDER BY time", Database: tableDatabase, MaxRows: 10}) + if err != nil || len(result.Rows) != 2 || result.Rows[0][0] != "1970-01-01T08:00:00.001+08:00" { + t.Fatalf("table query = %#v, %v", result, err) + } + tableDDL, err := tableServer.getTableDDL(tableDatabase, "d1") + if err != nil || !strings.Contains(tableDDL, "COMMENT 'DBX live test'") || !strings.Contains(tableDDL, `"device" STRING TAG`) { + t.Fatalf("table DDL = %q, %v", tableDDL, err) + } +} + +func TestLiveIoTDBAgentProcessMultiSessionAndCancellation(t *testing.T) { + if os.Getenv("DBX_IOTDB_LIVE") != "1" { + t.Skip("set DBX_IOTDB_LIVE=1 to run against a real IoTDB server") + } + process := startIoTDBAgentProcess(t) + defer process.close() + treeDatabase := "root.dbx_go_process_" + strconv.Itoa(os.Getpid()) + treeDevice := treeDatabase + ".d1" + treeParams := liveIoTDBParamsMap() + treeParams["agentSessionId"] = "tree" + treeParams["database"] = treeDatabase + process.call(t, "open_session", treeParams) + tableParams := liveIoTDBParamsMap() + tableParams["agentSessionId"] = "table" + tableParams["url_params"] = "sql_dialect=table" + process.call(t, "open_session", tableParams) + process.call(t, "execute_query", map[string]any{"agentSessionId": "tree", "sql": "CREATE DATABASE " + treeDatabase}) + process.call(t, "execute_query", map[string]any{ + "agentSessionId": "tree", "sql": "CREATE TIMESERIES " + treeDevice + ".s1 WITH DATATYPE=INT64, ENCODING=RLE", + }) + process.call(t, "execute_query", map[string]any{ + "agentSessionId": "tree", "sql": "INSERT INTO " + treeDevice + "(time,s1) VALUES(1,1)", + }) + + treeInfo := process.call(t, "connection_info", map[string]any{"agentSessionId": "tree"}) + tableInfo := process.call(t, "connection_info", map[string]any{"agentSessionId": "table"}) + if treeInfo.(map[string]any)["sqlDialect"] != "TREE" || tableInfo.(map[string]any)["sqlDialect"] != "TABLE" { + t.Fatalf("unexpected process dialects: tree=%#v table=%#v", treeInfo, tableInfo) + } + + queryID := process.send(t, "execute_query", map[string]any{ + "agentSessionId": "tree", + "sql": "SELECT COUNT(s1) FROM " + treeDevice + " GROUP BY ([0, 1000000000), 1ms)", + "maxRows": 100000000, + }) + time.Sleep(250 * time.Millisecond) + process.call(t, "cancel_session", map[string]any{"agentSessionId": "tree"}) + response := process.waitFor(t, queryID, 10*time.Second) + if response.Error == nil || response.Error.Data == nil || response.Error.Data.Category != "canceled" { + t.Fatalf("unexpected cancellation response: %#v", response) + } + process.call(t, "validate_session", map[string]any{"agentSessionId": "tree"}) + process.call(t, "execute_query", map[string]any{"agentSessionId": "tree", "sql": "DELETE DATABASE " + treeDatabase}) + process.call(t, "shutdown", nil) +} + +func TestIoTDBAgentHelperProcess(t *testing.T) { + if os.Getenv("DBX_IOTDB_AGENT_HELPER") != "1" { + return + } + main() + os.Exit(0) +} + +func liveIoTDBServer(t *testing.T, params connectParams) *server { + t.Helper() + base := liveIoTDBParams() + if params.Database != "" { + base.Database = params.Database + } + if params.URLParams != "" { + base.URLParams = params.URLParams + } + server, err := newServer(base) + if err != nil { + t.Fatal(err) + } + if err := server.validateConnection(); err != nil { + server.disconnect() + t.Fatal(err) + } + return server +} + +func liveIoTDBParams() connectParams { + return connectParams{ + Host: envOr("DBX_IOTDB_HOST", "127.0.0.1"), + Port: envIntOr("DBX_IOTDB_PORT", defaultIoTDBPort), + Username: envOr("DBX_IOTDB_USER", "root"), + Password: envOr("DBX_IOTDB_PASSWORD", "root"), + } +} + +func liveIoTDBParamsMap() map[string]any { + params := liveIoTDBParams() + return map[string]any{ + "host": params.Host, "port": params.Port, "username": params.Username, "password": params.Password, + } +} + +func mustExecuteNonQuery(t *testing.T, server *server, sql, database string) { + t.Helper() + if err := server.executeNonQuery(sql, database, 0); err != nil { + t.Fatalf("execute %q: %v", sql, err) + } +} + +func envOr(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + return fallback +} + +func envIntOr(name string, fallback int) int { + if value, err := strconv.Atoi(os.Getenv(name)); err == nil && value > 0 { + return value + } + return fallback +} + +type agentProcessResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result any `json:"result"` + Error *rpcError `json:"error"` +} + +type iotdbAgentProcess struct { + command *exec.Cmd + stdin *json.Encoder + responses chan agentProcessResponse + stderr *synchronizedBuffer + nextID int64 + pendingMu sync.Mutex + pending map[int64]agentProcessResponse +} + +func startIoTDBAgentProcess(t *testing.T) *iotdbAgentProcess { + t.Helper() + command := exec.Command(os.Args[0], "-test.run=^TestIoTDBAgentHelperProcess$") + command.Env = append(os.Environ(), "DBX_IOTDB_AGENT_HELPER=1") + stdout, err := command.StdoutPipe() + if err != nil { + t.Fatal(err) + } + stdin, err := command.StdinPipe() + if err != nil { + t.Fatal(err) + } + process := &iotdbAgentProcess{ + command: command, stdin: json.NewEncoder(stdin), responses: make(chan agentProcessResponse, 16), + stderr: &synchronizedBuffer{}, pending: map[int64]agentProcessResponse{}, + } + command.Stderr = process.stderr + if err := command.Start(); err != nil { + t.Fatal(err) + } + scanner := bufio.NewScanner(stdout) + if !scanner.Scan() || scanner.Text() != `{"ready":true}` { + process.close() + t.Fatalf("unexpected ready line %q: %s", scanner.Text(), process.stderr.String()) + } + go func() { + for scanner.Scan() { + var response agentProcessResponse + if json.Unmarshal(scanner.Bytes(), &response) == nil { + process.responses <- response + } + } + close(process.responses) + }() + return process +} + +func (p *iotdbAgentProcess) send(t *testing.T, method string, params map[string]any) int64 { + t.Helper() + p.nextID++ + request := map[string]any{"jsonrpc": "2.0", "id": p.nextID, "method": method, "params": params} + if err := p.stdin.Encode(request); err != nil { + t.Fatal(err) + } + return p.nextID +} + +func (p *iotdbAgentProcess) call(t *testing.T, method string, params map[string]any) any { + t.Helper() + id := p.send(t, method, params) + response := p.waitFor(t, id, 10*time.Second) + if response.Error != nil { + t.Fatalf("%s failed: %#v", method, response.Error) + } + return response.Result +} + +func (p *iotdbAgentProcess) waitFor(t *testing.T, id int64, timeout time.Duration) agentProcessResponse { + t.Helper() + p.pendingMu.Lock() + if response, ok := p.pending[id]; ok { + delete(p.pending, id) + p.pendingMu.Unlock() + return response + } + p.pendingMu.Unlock() + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + for { + select { + case response, ok := <-p.responses: + if !ok { + t.Fatalf("agent exited while waiting for %d: %s", id, p.stderr.String()) + } + var responseID int64 + if err := json.Unmarshal(response.ID, &responseID); err != nil { + t.Fatalf("invalid response id %s", response.ID) + } + if responseID == id { + return response + } + p.pendingMu.Lock() + p.pending[responseID] = response + p.pendingMu.Unlock() + case <-ctx.Done(): + t.Fatalf("timed out waiting for response %d: %s", id, p.stderr.String()) + } + } +} + +func (p *iotdbAgentProcess) close() { + if p.command.Process == nil || p.command.ProcessState != nil { + return + } + _ = p.command.Process.Kill() + _ = p.command.Wait() +} + +type synchronizedBuffer struct { + mu sync.Mutex + builder strings.Builder +} + +func (b *synchronizedBuffer) Write(value []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.builder.Write(value) +} + +func (b *synchronizedBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.builder.String() +} diff --git a/agents/drivers/iotdb/main.go b/agents/drivers/iotdb/main.go new file mode 100644 index 0000000000..0e2358c5cd --- /dev/null +++ b/agents/drivers/iotdb/main.go @@ -0,0 +1,515 @@ +package main + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "runtime" + "strconv" + "strings" + "sync" + "time" +) + +const ( + protocolVersion = 2 + defaultIoTDBPort = 6667 + defaultMaxRows = 10000 + defaultPageSize = 1000 + defaultFetchSize = 1024 + defaultConnectTimeout = 15 * time.Second + legacyAgentSessionID = "__legacy__" + maxAgentSessions = 256 +) + +type request struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params map[string]json.RawMessage `json:"params"` +} + +type response struct { + JSONRPC string `json:"jsonrpc,omitempty"` + ID json.RawMessage `json:"id,omitempty"` + Result any `json:"result,omitempty"` + Error *rpcError `json:"error,omitempty"` +} + +type connectParams struct { + Host string `json:"host"` + Port int `json:"port"` + Database string `json:"database"` + Username string `json:"username"` + Password string `json:"password"` + URLParams string `json:"url_params"` + ConnectionString string `json:"connection_string"` + SSL bool `json:"ssl"` + CACertPath string `json:"ca_cert_path"` + ClientCertPath string `json:"client_cert_path"` + ClientKeyPath string `json:"client_key_path"` + SessionRole string `json:"sessionRole"` +} + +type queryOptions struct { + SQL string `json:"sql"` + Database string `json:"database"` + Schema string `json:"schema"` + MaxRows int `json:"maxRows"` + FetchSize int `json:"fetchSize"` + TimeoutSecs int `json:"timeoutSecs"` +} + +type queryResult struct { + Columns []string `json:"columns"` + ColumnTypes []string `json:"column_types"` + Rows [][]any `json:"rows"` + AffectedRows int64 `json:"affected_rows"` + ExecutionTimeMS int64 `json:"execution_time_ms"` + Truncated bool `json:"truncated"` +} + +type queryPageResult struct { + Columns []string `json:"columns"` + ColumnTypes []string `json:"column_types"` + Rows [][]any `json:"rows"` + AffectedRows int64 `json:"affected_rows"` + ExecutionTimeMS int64 `json:"execution_time_ms"` + Truncated bool `json:"truncated"` + SessionID *string `json:"session_id"` + HasMore bool `json:"has_more"` +} + +type server struct { + params connectParams + config connectionConfig + + clientMu sync.Mutex + client *sessionClient + + querySessions map[string]*querySession + nextSessionID uint64 + + activeMu sync.Mutex + activeCancel context.CancelFunc +} + +type agentSession struct { + server *server + mu sync.Mutex +} + +type runtimeServer struct { + mu sync.RWMutex + sessions map[string]*agentSession +} + +func main() { + configureRuntimeParallelism() + runtimeServer := newRuntimeServer() + encoder := json.NewEncoder(os.Stdout) + var encoderMu sync.Mutex + var requests sync.WaitGroup + fmt.Fprintln(os.Stdout, `{"ready":true}`) + + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 0, 64*1024), 512*1024*1024) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var envelope request + if json.Unmarshal([]byte(line), &envelope) == nil && envelope.Method == "shutdown" { + requests.Wait() + result, _ := runtimeServer.handleLine(line) + encoderMu.Lock() + _ = encoder.Encode(result) + encoderMu.Unlock() + return + } + requests.Add(1) + go func(line string) { + defer requests.Done() + result, _ := runtimeServer.handleLine(line) + encoderMu.Lock() + defer encoderMu.Unlock() + if err := encoder.Encode(result); err != nil { + fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err) + } + }(line) + } + requests.Wait() +} + +func configureRuntimeParallelism() { + if raw := strings.TrimSpace(os.Getenv("DBX_AGENT_IOTDB_GOMAXPROCS")); raw != "" { + if configured, err := strconv.Atoi(raw); err == nil && configured > 0 { + runtime.GOMAXPROCS(configured) + return + } + } + if strings.TrimSpace(os.Getenv("GOMAXPROCS")) == "" { + runtime.GOMAXPROCS(min(runtime.NumCPU(), 4)) + } +} + +func newRuntimeServer() *runtimeServer { + return &runtimeServer{sessions: map[string]*agentSession{}} +} + +func (r *runtimeServer) handleLine(line string) (response, bool) { + var req request + if err := json.Unmarshal([]byte(line), &req); err != nil { + return errorResponse(nil, "", "", err), false + } + if len(req.ID) == 0 { + req.ID = json.RawMessage("1") + } + result, shutdown, err := r.dispatch(req.Method, req.Params) + if err != nil { + return errorResponse(req.ID, req.Method, stringParam(req.Params, "agentSessionId"), err), false + } + return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown +} + +func (r *runtimeServer) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) { + switch method { + case "handshake": + return handshakeResult(true), false, nil + case "open_session": + id := stringParam(params, "agentSessionId") + if id == "" { + return nil, false, errors.New("agentSessionId is required") + } + var connection connectParams + if err := decodeParams(params, &connection); err != nil { + return nil, false, err + } + return map[string]bool{"ok": true}, false, r.openSession(id, connection) + case "close_session": + return map[string]bool{"ok": true}, false, r.closeSession(stringParam(params, "agentSessionId")) + case "validate_session": + session, err := r.session(stringParam(params, "agentSessionId")) + if err != nil { + return nil, false, err + } + session.mu.Lock() + defer session.mu.Unlock() + return map[string]bool{"ok": true}, false, session.server.validateConnection() + case "cancel_session": + session, err := r.session(stringParam(params, "agentSessionId")) + if err != nil { + return nil, false, err + } + session.server.cancelActiveQuery() + return map[string]bool{"ok": true}, false, nil + case "test_connection": + var connection connectParams + if err := decodeParams(params, &connection); err != nil { + return nil, false, err + } + result, err := testConnection(connection) + return result, false, err + case "connect": + var connection connectParams + if err := decodeParams(params, &connection); err != nil { + return nil, false, err + } + _ = r.closeSession(legacyAgentSessionID) + return map[string]bool{"ok": true}, false, r.openSession(legacyAgentSessionID, connection) + case "disconnect": + return map[string]bool{"ok": true}, false, r.closeSession(legacyAgentSessionID) + case "shutdown": + return map[string]bool{"ok": true}, true, r.closeAllSessions() + default: + id := stringParam(params, "agentSessionId") + if id == "" { + id = legacyAgentSessionID + } + session, err := r.session(id) + if err != nil { + return nil, false, err + } + session.mu.Lock() + defer session.mu.Unlock() + return session.server.dispatch(method, params) + } +} + +func (r *runtimeServer) openSession(id string, params connectParams) error { + r.mu.Lock() + if _, exists := r.sessions[id]; exists { + r.mu.Unlock() + return fmt.Errorf("agent session already exists: %s", id) + } + if len(r.sessions) >= maxAgentSessions { + r.mu.Unlock() + return fmt.Errorf("agent session limit reached: %d", maxAgentSessions) + } + r.mu.Unlock() + + server, err := newServer(params) + if err != nil { + return err + } + if err := server.validateConnection(); err != nil { + server.disconnect() + return err + } + + r.mu.Lock() + defer r.mu.Unlock() + if _, exists := r.sessions[id]; exists { + server.disconnect() + return fmt.Errorf("agent session already exists: %s", id) + } + r.sessions[id] = &agentSession{server: server} + return nil +} + +func (r *runtimeServer) session(id string) (*agentSession, error) { + r.mu.RLock() + session := r.sessions[id] + r.mu.RUnlock() + if session == nil { + return nil, fmt.Errorf("agent session not found: %s", id) + } + return session, nil +} + +func (r *runtimeServer) closeSession(id string) error { + r.mu.Lock() + session := r.sessions[id] + delete(r.sessions, id) + r.mu.Unlock() + if session == nil { + return nil + } + session.server.cancelActiveQuery() + session.mu.Lock() + defer session.mu.Unlock() + return session.server.disconnect() +} + +func (r *runtimeServer) closeAllSessions() error { + r.mu.RLock() + ids := make([]string, 0, len(r.sessions)) + for id := range r.sessions { + ids = append(ids, id) + } + r.mu.RUnlock() + var firstErr error + for _, id := range ids { + if err := r.closeSession(id); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + +func newServer(params connectParams) (*server, error) { + config, err := parseConnectionConfig(params) + if err != nil { + return nil, err + } + return &server{ + params: params, + config: config, + querySessions: map[string]*querySession{}, + }, nil +} + +func (s *server) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) { + switch method { + case "handshake": + return handshakeResult(false), false, nil + case "validate_connection": + return map[string]bool{"ok": true}, false, s.validateConnection() + case "connection_info": + result, err := s.connectionInfo() + return result, false, err + case "list_databases": + result, err := s.listDatabases() + return result, false, err + case "list_schemas": + result, err := s.listSchemas() + return result, false, err + case "list_tables": + result, err := s.listTables(stringParam(params, "schema"), metadataListConstraintsFromParams(params)) + return result, false, err + case "get_table_comment": + result, err := s.getTableComment(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "list_objects": + result, err := s.listObjects(stringParam(params, "schema"), metadataListConstraintsFromParams(params)) + return result, false, err + case "list_data_types": + return iotdbDataTypes(), false, nil + case "completion_assistant_search_v1": + var input completionAssistantRequest + if err := decodeParams(params, &input); err != nil { + return nil, false, err + } + result, err := s.completionAssistantSearch(input) + return result, false, err + case "get_columns": + result, err := s.getColumns(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "list_indexes": + result, err := s.listIndexes(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "list_foreign_keys", "list_triggers", "list_constraints", "list_partitions", "list_subpartitions": + return []any{}, false, nil + case "get_object_source", "get_table_ddl": + result, err := s.getTableDDL(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "get_explain_info": + return map[string]any{"plan": "", "has_actual_stats": false}, false, nil + case "execute_query": + result, err := s.executeQuery(queryOptionsFromParams(params)) + return result, false, err + case "execute_query_page", "start_table_read": + result, err := s.executeQueryPage(queryOptionsFromParams(params), intParam(params, "pageSize")) + return result, false, err + case "fetch_query_page", "fetch_table_read_page": + result, err := s.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize")) + return result, false, err + case "close_query_session", "close_table_read_session": + return s.closeQuerySession(stringParam(params, "sessionId")), false, nil + case "execute_transaction": + result, err := s.executeStatements(params, true) + return result, false, err + case "execute_batch": + result, err := s.executeStatements(params, false) + return result, false, err + case "disconnect": + return map[string]bool{"ok": true}, false, s.disconnect() + case "shutdown": + return map[string]bool{"ok": true}, true, s.disconnect() + default: + return nil, false, fmt.Errorf("unknown method: %s", method) + } +} + +func handshakeResult(multiSession bool) map[string]any { + capabilities := []string{ + "connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "structured_error_v1", + } + if multiSession { + capabilities = append(capabilities, "multi_session") + } + return map[string]any{ + "protocolVersion": protocolVersion, + "agentProtocolVersion": protocolVersion, + "capabilities": capabilities, + } +} + +func testConnection(params connectParams) (map[string]any, error) { + server, err := newServer(params) + if err != nil { + return nil, err + } + defer server.disconnect() + if err := server.validateConnection(); err != nil { + return nil, err + } + info, err := server.connectionInfo() + if err != nil { + return nil, err + } + return map[string]any{"ok": true, "info": info}, nil +} + +func (s *server) disconnect() error { + s.cancelActiveQuery() + if err := s.closeAllQuerySessions(); err != nil { + return err + } + s.clientMu.Lock() + client := s.client + s.client = nil + s.clientMu.Unlock() + if client != nil { + return client.Close() + } + return nil +} + +func (s *server) setActiveOperation(cancel context.CancelFunc) { + s.activeMu.Lock() + s.activeCancel = cancel + s.activeMu.Unlock() +} + +func (s *server) clearActiveOperation(cancel context.CancelFunc) { + cancel() + s.activeMu.Lock() + s.activeCancel = nil + s.activeMu.Unlock() +} + +func (s *server) cancelActiveQuery() { + s.activeMu.Lock() + cancel := s.activeCancel + s.activeMu.Unlock() + if cancel != nil { + cancel() + } +} + +func queryOptionsFromParams(params map[string]json.RawMessage) queryOptions { + return queryOptions{ + SQL: stringParam(params, "sql"), + Database: stringParam(params, "database"), + Schema: stringParam(params, "schema"), + MaxRows: intParam(params, "maxRows"), + FetchSize: intParam(params, "fetchSize"), + TimeoutSecs: intParam(params, "timeoutSecs"), + } +} + +func decodeParams(params map[string]json.RawMessage, target any) error { + data, err := json.Marshal(params) + if err != nil { + return err + } + return json.Unmarshal(data, target) +} + +func stringParam(params map[string]json.RawMessage, key string) string { + if raw, ok := params[key]; ok { + var value string + if json.Unmarshal(raw, &value) == nil { + return value + } + } + return "" +} + +func intParam(params map[string]json.RawMessage, key string) int { + if raw, ok := params[key]; ok { + var value int + if json.Unmarshal(raw, &value) == nil { + return value + } + } + return 0 +} + +func stringSliceParam(params map[string]json.RawMessage, key string) []string { + if raw, ok := params[key]; ok { + var value []string + if json.Unmarshal(raw, &value) == nil && value != nil { + return value + } + } + return []string{} +} + +func errorResponse(id json.RawMessage, method, sessionID string, err error) response { + return response{JSONRPC: "2.0", ID: id, Error: classifyRPCError(method, sessionID, err)} +} diff --git a/agents/drivers/iotdb/main_test.go b/agents/drivers/iotdb/main_test.go new file mode 100644 index 0000000000..33e0ac527f --- /dev/null +++ b/agents/drivers/iotdb/main_test.go @@ -0,0 +1,199 @@ +package main + +import ( + "context" + "encoding/json" + "errors" + "reflect" + "testing" + "time" + + "github.com/apache/iotdb-client-go/v2/client" +) + +func TestHandshakeAdvertisesNativeCapabilities(t *testing.T) { + result, shutdown, err := newRuntimeServer().dispatch("handshake", nil) + if err != nil || shutdown { + t.Fatalf("unexpected handshake: shutdown=%t err=%v", shutdown, err) + } + capabilities := result.(map[string]any)["capabilities"].([]string) + want := []string{ + "connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", + "structured_error_v1", "multi_session", + } + if !reflect.DeepEqual(capabilities, want) { + t.Fatalf("unexpected capabilities: %#v", capabilities) + } +} + +func TestHandleLineClassifiesMissingSession(t *testing.T) { + response, _ := newRuntimeServer().handleLine( + `{"jsonrpc":"2.0","id":7,"method":"validate_session","params":{"agentSessionId":"missing"}}`, + ) + if response.Error == nil || response.Error.Data == nil { + t.Fatalf("expected structured error: %#v", response) + } + if response.Error.Data.Stage != "validate" || response.Error.Data.Category != "protocol" { + t.Fatalf("unexpected error classification: %#v", response.Error.Data) + } +} + +func TestParseConnectionConfig(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:iotdb://alice:secret@[::1]:7777/dbx_table?sql_dialect=table&fetch_size=2048&connect_retry_max=4", + URLParams: "enable_compression=true&node_urls=node1:6667,node2:6667&time_zone=UTC", + }) + if err != nil { + t.Fatal(err) + } + if config.Host != "::1" || config.Port != 7777 || config.Username != "alice" || config.Password != "secret" { + t.Fatalf("unexpected endpoint config: %#v", config) + } + if config.Database != "dbx_table" || config.Dialect != client.TableSqlDialect || config.FetchSize != 2048 { + t.Fatalf("unexpected session config: %#v", config) + } + if !config.EnableCompression || config.ConnectRetryMax != 4 || config.TimeZone != "UTC" { + t.Fatalf("unexpected optional config: %#v", config) + } + if !reflect.DeepEqual(config.NodeURLs, []string{"node1:6667", "node2:6667"}) { + t.Fatalf("unexpected cluster nodes: %#v", config.NodeURLs) + } +} + +func TestParseConnectionConfigTLS(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + Host: "iotdb.example.com", SSL: true, CACertPath: "/tmp/ca.pem", + ClientCertPath: "/tmp/client.pem", ClientKeyPath: "/tmp/client.key", + URLParams: "insecure_skip_verify=true", + }) + if err != nil { + t.Fatal(err) + } + if config.TLSConfig == nil || config.TLSConfig.CAFile != "/tmp/ca.pem" || !config.TLSInsecureSkipVerify { + t.Fatalf("unexpected TLS config: %#v", config.TLSConfig) + } + if _, err := parseConnectionConfig(connectParams{SSL: true, ClientCertPath: "/tmp/client.pem"}); err == nil { + t.Fatal("expected incomplete mTLS configuration to fail") + } + if _, err := parseConnectionConfig(connectParams{URLParams: "sql_dialect=relational"}); err == nil { + t.Fatal("expected unsupported SQL dialect to fail") + } +} + +func TestMetadataHelpers(t *testing.T) { + values := []string{"alpha", "beta", "gamma"} + if got := applyMetadataWindow(values, 1, 1); !reflect.DeepEqual(got, []string{"beta"}) { + t.Fatalf("unexpected metadata window: %#v", got) + } + if !metadataNameMatches("DeviceOne", "vice") || metadataNameMatches("DeviceOne", "table") { + t.Fatal("unexpected metadata name matching") + } + if got := treeDevicePath("root.db", "root.db.d1"); got != "root.db.d1" { + t.Fatalf("treeDevicePath() = %q", got) + } + if got := quoteTreePath("root.db.with-dash"); got != "root.db.`with-dash`" { + t.Fatalf("quoteTreePath() = %q", got) + } + if got := quoteTableIdentifier(`a"b`); got != `"a""b"` { + t.Fatalf("quoteTableIdentifier() = %q", got) + } +} + +func TestNormalizeIoTDBValues(t *testing.T) { + when := time.Date(2026, time.August, 10, 12, 34, 56, 123, time.FixedZone("CST", 8*60*60)) + if got := normalizeIoTDBValue(when, "TIMESTAMP"); got != "2026-08-10T12:34:56.000000123+08:00" { + t.Fatalf("unexpected timestamp: %#v", got) + } + if got := normalizeIoTDBValue(when, "DATE"); got != "2026-08-10" { + t.Fatalf("unexpected date: %#v", got) + } + if got := normalizeIoTDBValue([]byte{0xde, 0xad}, "BLOB"); got != "dead" { + t.Fatalf("unexpected blob: %#v", got) + } +} + +func TestClassifyIoTDBErrors(t *testing.T) { + canceled := classifyRPCError("execute_query", "session-1", context.Canceled) + if canceled.Data.Category != "canceled" || canceled.Data.SessionDisposition != "quarantine" { + t.Fatalf("unexpected cancellation classification: %#v", canceled) + } + timedOut := classifyRPCError("execute_query", "session-1", context.DeadlineExceeded) + if timedOut.Data.Category != "timeout" || timedOut.Data.SessionDisposition != "quarantine" { + t.Fatalf("unexpected timeout classification: %#v", timedOut) + } +} + +func TestDecodeQueryOptions(t *testing.T) { + params := map[string]json.RawMessage{ + "sql": json.RawMessage(`"SHOW DATABASES"`), + "maxRows": json.RawMessage(`100`), + "fetchSize": json.RawMessage(`25`), + "timeoutSecs": json.RawMessage(`5`), + } + options := queryOptionsFromParams(params) + if options.SQL != "SHOW DATABASES" || options.MaxRows != 100 || options.FetchSize != 25 || options.TimeoutSecs != 5 { + t.Fatalf("unexpected query options: %#v", options) + } +} + +func TestCancelActiveOperationInvalidatesClient(t *testing.T) { + blocking := &blockingIoTDBSession{closed: make(chan struct{})} + connected := &sessionClient{session: blocking, dialect: client.TreeSqlDialect} + server := &server{client: connected} + done := make(chan error, 1) + go func() { + _, err := runCancelable(server, 0, connected, func(ctx context.Context) (struct{}, error) { + <-ctx.Done() + return struct{}{}, ctx.Err() + }) + done <- err + }() + deadline := time.Now().Add(time.Second) + for { + server.activeMu.Lock() + active := server.activeCancel != nil + server.activeMu.Unlock() + if active { + break + } + if time.Now().After(deadline) { + t.Fatal("operation never became active") + } + time.Sleep(time.Millisecond) + } + server.cancelActiveQuery() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("runCancelable() error = %v", err) + } + case <-time.After(time.Second): + t.Fatal("canceled operation did not return") + } + server.clientMu.Lock() + defer server.clientMu.Unlock() + if server.client != nil { + t.Fatal("canceled client must be invalidated") + } +} + +type blockingIoTDBSession struct { + closed chan struct{} +} + +func (s *blockingIoTDBSession) ExecuteQueryStatement(context.Context, string, *int64) (*client.SessionDataSet, error) { + return nil, errors.New("not implemented") +} + +func (s *blockingIoTDBSession) ExecuteNonQueryStatement(context.Context, string) error { + return errors.New("not implemented") +} + +func (s *blockingIoTDBSession) Close() error { + select { + case <-s.closed: + default: + close(s.closed) + } + return nil +} diff --git a/agents/drivers/iotdb/metadata.go b/agents/drivers/iotdb/metadata.go new file mode 100644 index 0000000000..ac4c39f382 --- /dev/null +++ b/agents/drivers/iotdb/metadata.go @@ -0,0 +1,703 @@ +package main + +import ( + "encoding/json" + "fmt" + "regexp" + "sort" + "strconv" + "strings" + + "github.com/apache/iotdb-client-go/v2/client" +) + +const metadataQueryLimit = 100000 + +var ( + iotdbTypes = []string{ + "BOOLEAN", "INT32", "INT64", "FLOAT", "DOUBLE", "TEXT", "STRING", "BLOB", "TIMESTAMP", "DATE", + } + simpleTreeNodePattern = regexp.MustCompile(`^[A-Za-z0-9_]+$`) +) + +type databaseInfo struct { + Name string `json:"name"` +} + +type tableInfo struct { + Name string `json:"name"` + TableType string `json:"table_type"` + Comment *string `json:"comment"` + ParentSchema *string `json:"parent_schema,omitempty"` + ParentName *string `json:"parent_name,omitempty"` +} + +type objectInfo struct { + Name string `json:"name"` + ObjectType string `json:"object_type"` + Schema string `json:"schema"` + Comment *string `json:"comment"` + Valid *bool `json:"valid,omitempty"` +} + +type columnInfo struct { + Name string `json:"name"` + DataType string `json:"data_type"` + IsNullable bool `json:"is_nullable"` + ColumnDefault *string `json:"column_default"` + IsPrimaryKey bool `json:"is_primary_key"` + Extra *string `json:"extra"` + Comment *string `json:"comment"` + NumericPrecision *int `json:"numeric_precision"` + NumericScale *int `json:"numeric_scale"` + CharacterMaximumLength *int `json:"character_maximum_length"` +} + +type indexInfo struct { + Name string `json:"name"` + Columns []string `json:"columns"` + IsUnique bool `json:"is_unique"` + IsPrimary bool `json:"is_primary"` + Filter *string `json:"filter"` + IndexType *string `json:"index_type"` + IncludedColumns []string `json:"included_columns"` + Comment *string `json:"comment"` +} + +func (i indexInfo) MarshalJSON() ([]byte, error) { + type alias indexInfo + value := alias(i) + if value.Columns == nil { + value.Columns = []string{} + } + if value.IncludedColumns == nil { + value.IncludedColumns = []string{} + } + return json.Marshal(value) +} + +type metadataListConstraints struct { + Filter string + Limit int + Offset int + ObjectTypes []string +} + +type completionAssistantRequest struct { + ConnectionID string `json:"connection_id"` + Database string `json:"database"` + Schema string `json:"schema"` + ObjectKinds []string `json:"object_kinds"` + Mask string `json:"mask"` + CaseSensitive bool `json:"case_sensitive"` + GlobalSearch bool `json:"global_search"` + MaxResults int `json:"max_results"` + ParentSchema string `json:"parent_schema"` + ParentName string `json:"parent_name"` + MatchMode string `json:"match_mode"` +} + +type completionAssistantCandidate struct { + Name string `json:"name"` + Kind string `json:"kind"` + Database *string `json:"database"` + Schema *string `json:"schema"` + ParentSchema *string `json:"parent_schema"` + ParentName *string `json:"parent_name"` + Comment *string `json:"comment"` + DataType *string `json:"data_type"` +} + +type completionAssistantResponse struct { + Candidates []completionAssistantCandidate `json:"candidates"` + Incomplete bool `json:"incomplete"` + FallbackUsed bool `json:"fallback_used"` +} + +func iotdbDataTypes() []string { + return append([]string(nil), iotdbTypes...) +} + +func (s *server) connectionInfo() (map[string]any, error) { + version := "" + if values, err := s.queryValues("SHOW VERSION", "", 1, 5); err == nil && len(values.Rows) > 0 { + version = metadataRowString(values, values.Rows[0], "Version") + } else if err != nil { + return nil, err + } + username := s.params.Username + if values, err := s.queryValues("SHOW CURRENT_USER", "", 1, 5); err == nil && len(values.Rows) > 0 { + if current := metadataRowString(values, values.Rows[0], "CurrentUser"); current != "" { + username = current + } + } + dialect := strings.ToUpper(s.config.Dialect) + if values, err := s.queryValues("SHOW CURRENT_SQL_DIALECT", "", 1, 5); err == nil && len(values.Rows) > 0 { + if current := metadataRowString(values, values.Rows[0], "CurrentSqlDialect"); current != "" { + dialect = strings.ToUpper(current) + } + } + quote := "`" + identifierCase := "mixed" + if s.config.Dialect == client.TableSqlDialect { + quote = `"` + identifierCase = "lower" + } + return map[string]any{ + "database": s.config.Database, + "schema": s.config.Database, + "username": username, + "version": version, + "sqlDialect": dialect, + "identifierQuote": quote, + "compatibilityMode": "iotdb-" + strings.ToLower(dialect), + "databaseInfo": map[string]string{ + "productName": "Apache IoTDB", + "productVersion": version, + "unquotedIdentifierCase": identifierCase, + "quotedIdentifierCase": "mixed", + "driverName": "Apache IoTDB Go Client", + "driverVersion": "v" + iotdbGoClientVersion, + }, + }, nil +} + +func (s *server) listDatabases() ([]databaseInfo, error) { + values, err := s.queryValues("SHOW DATABASES", "", metadataQueryLimit, 0) + if err != nil { + return nil, err + } + result := make([]databaseInfo, 0, len(values.Rows)) + seen := map[string]bool{} + for _, row := range values.Rows { + name := metadataRowString(values, row, "Database") + if name == "" || seen[name] { + continue + } + seen[name] = true + result = append(result, databaseInfo{Name: name}) + } + sort.Slice(result, func(left, right int) bool { return result[left].Name < result[right].Name }) + return result, nil +} + +func (s *server) listSchemas() ([]string, error) { + databases, err := s.listDatabases() + if err != nil { + return nil, err + } + result := make([]string, len(databases)) + for index, database := range databases { + result[index] = database.Name + } + return result, nil +} + +func (s *server) listTables(schema string, constraints metadataListConstraints) ([]tableInfo, error) { + if !acceptsIoTDBTable(constraints.ObjectTypes) { + return []tableInfo{}, nil + } + if s.config.Dialect == client.TableSqlDialect { + return s.listTableDialectTables(schema, constraints) + } + return s.listTreeDialectDevices(schema, constraints) +} + +func (s *server) listTreeDialectDevices(schema string, constraints metadataListConstraints) ([]tableInfo, error) { + schema = s.effectiveMetadataSchema(schema) + pattern := "root.**" + if schema != "" { + pattern = quoteTreePath(schema) + ".**" + } + values, err := s.queryValues("SHOW DEVICES "+pattern, "", metadataQueryLimit, 0) + if err != nil { + return nil, err + } + prefix := strings.TrimSuffix(schema, ".") + if prefix != "" { + prefix += "." + } + result := make([]tableInfo, 0, len(values.Rows)) + for _, row := range values.Rows { + device := metadataRowString(values, row, "Device") + name := strings.TrimPrefix(device, prefix) + if name == "" || name == device && prefix != "" || !metadataNameMatches(name, constraints.Filter) { + continue + } + result = append(result, tableInfo{Name: name, TableType: "TABLE"}) + } + sort.Slice(result, func(left, right int) bool { return result[left].Name < result[right].Name }) + return applyMetadataWindow(result, constraints.Offset, constraints.Limit), nil +} + +func (s *server) listTableDialectTables(schema string, constraints metadataListConstraints) ([]tableInfo, error) { + schema = s.effectiveMetadataSchema(schema) + if schema == "" { + return []tableInfo{}, nil + } + values, err := s.queryValues("SHOW TABLES DETAILS FROM "+quoteTableIdentifier(schema), schema, metadataQueryLimit, 0) + if err != nil { + values, err = s.queryValues("SHOW TABLES FROM "+quoteTableIdentifier(schema), schema, metadataQueryLimit, 0) + if err != nil { + return nil, err + } + } + result := make([]tableInfo, 0, len(values.Rows)) + for _, row := range values.Rows { + name := firstMetadataRowString(values, row, "TableName", "table_name") + if name == "" || !metadataNameMatches(name, constraints.Filter) { + continue + } + tableType := firstMetadataRowString(values, row, "TableType", "table_type") + if tableType == "" || strings.EqualFold(tableType, "BASE TABLE") { + tableType = "TABLE" + } + result = append(result, tableInfo{ + Name: name, + TableType: tableType, + Comment: optionalMetadataString(firstMetadataRowString(values, row, "Comment", "comment")), + }) + } + sort.Slice(result, func(left, right int) bool { return result[left].Name < result[right].Name }) + return applyMetadataWindow(result, constraints.Offset, constraints.Limit), nil +} + +func (s *server) getTableComment(schema, table string) (*string, error) { + if s.config.Dialect != client.TableSqlDialect { + return nil, nil + } + tables, err := s.listTableDialectTables(schema, metadataListConstraints{}) + if err != nil { + return nil, err + } + for _, item := range tables { + if item.Name == table { + return item.Comment, nil + } + } + return nil, nil +} + +func (s *server) listObjects(schema string, constraints metadataListConstraints) ([]objectInfo, error) { + tables, err := s.listTables(schema, constraints) + if err != nil { + return nil, err + } + result := make([]objectInfo, 0, len(tables)) + for _, table := range tables { + result = append(result, objectInfo{ + Name: table.Name, + ObjectType: "TABLE", + Schema: schema, + Comment: table.Comment, + }) + } + return result, nil +} + +func (s *server) getColumns(schema, table string) ([]columnInfo, error) { + if strings.TrimSpace(table) == "" { + return []columnInfo{}, nil + } + if s.config.Dialect == client.TableSqlDialect { + return s.getTableDialectColumns(schema, table) + } + return s.getTreeDialectColumns(schema, table) +} + +func (s *server) getTreeDialectColumns(schema, table string) ([]columnInfo, error) { + device := treeDevicePath(s.effectiveMetadataSchema(schema), table) + if device == "" { + return []columnInfo{}, nil + } + values, err := s.queryValues("SHOW TIMESERIES "+quoteTreePath(device)+".**", "", metadataQueryLimit, 0) + if err != nil { + return nil, err + } + prefix := device + "." + result := make([]columnInfo, 0, len(values.Rows)) + for _, row := range values.Rows { + path := metadataRowString(values, row, "Timeseries") + name := strings.TrimPrefix(path, prefix) + if name == "" || name == path || strings.Contains(name, ".") { + continue + } + encoding := metadataRowString(values, row, "Encoding") + compression := metadataRowString(values, row, "Compression") + extra := make([]string, 0, 2) + if encoding != "" { + extra = append(extra, "encoding="+encoding) + } + if compression != "" { + extra = append(extra, "compressor="+compression) + } + result = append(result, columnInfo{ + Name: name, + DataType: metadataRowString(values, row, "DataType"), + IsNullable: true, + Extra: optionalMetadataString(strings.Join(extra, ", ")), + }) + } + return result, nil +} + +func (s *server) getTableDialectColumns(schema, table string) ([]columnInfo, error) { + schema = s.effectiveMetadataSchema(schema) + qualified := quoteTableIdentifier(schema) + "." + quoteTableIdentifier(table) + values, err := s.queryValues("DESC "+qualified+" DETAILS", schema, metadataQueryLimit, 0) + if err != nil { + values, err = s.queryValues("DESC "+qualified, schema, metadataQueryLimit, 0) + if err != nil { + return nil, err + } + } + result := make([]columnInfo, 0, len(values.Rows)) + for _, row := range values.Rows { + name := firstMetadataRowString(values, row, "ColumnName", "column_name") + if name == "" { + continue + } + category := strings.ToUpper(firstMetadataRowString(values, row, "Category", "category")) + isPrimary := category == "TIME" || category == "TAG" + result = append(result, columnInfo{ + Name: name, + DataType: strings.ToUpper(firstMetadataRowString(values, row, "DataType", "datatype")), + IsNullable: !isPrimary, + IsPrimaryKey: isPrimary, + Extra: optionalMetadataString(category), + Comment: optionalMetadataString(firstMetadataRowString(values, row, "Comment", "comment")), + }) + } + return result, nil +} + +func (s *server) listIndexes(_, _ string) ([]indexInfo, error) { + return []indexInfo{}, nil +} + +func (s *server) getTableDDL(schema, table string) (string, error) { + if strings.TrimSpace(table) == "" { + return "", nil + } + if s.config.Dialect == client.TableSqlDialect { + return s.getTableDialectDDL(schema, table) + } + return s.getTreeDialectDDL(schema, table) +} + +func (s *server) getTreeDialectDDL(schema, table string) (string, error) { + device := treeDevicePath(s.effectiveMetadataSchema(schema), table) + columns, err := s.getTreeDialectColumns(schema, table) + if err != nil { + return "", err + } + if len(columns) == 0 { + return "", nil + } + aligned := false + if values, queryErr := s.queryValues("SHOW DEVICES "+quoteTreePath(device), "", 1, 0); queryErr == nil && len(values.Rows) > 0 { + aligned = metadataRowBool(values, values.Rows[0], "IsAligned") + } + var ddl strings.Builder + for _, column := range columns { + ddl.WriteString("DELETE TIMESERIES ") + ddl.WriteString(quoteTreePath(device + "." + column.Name)) + ddl.WriteString(";\n") + } + ddl.WriteString("\n") + if aligned { + definitions := make([]string, 0, len(columns)) + for _, column := range columns { + definition := quoteTreeNode(column.Name) + " " + column.DataType + if column.Extra != nil && *column.Extra != "" { + definition += " " + strings.ReplaceAll(*column.Extra, ", ", " ") + } + definitions = append(definitions, definition) + } + ddl.WriteString("CREATE ALIGNED TIMESERIES ") + ddl.WriteString(quoteTreePath(device)) + ddl.WriteString("(") + ddl.WriteString(strings.Join(definitions, ", ")) + ddl.WriteString(");") + return ddl.String(), nil + } + for _, column := range columns { + ddl.WriteString("CREATE TIMESERIES ") + ddl.WriteString(quoteTreePath(device + "." + column.Name)) + ddl.WriteString(" WITH DATATYPE=") + ddl.WriteString(column.DataType) + if column.Extra != nil && *column.Extra != "" { + ddl.WriteString(", ") + ddl.WriteString(*column.Extra) + } + ddl.WriteString(";\n") + } + return strings.TrimSuffix(ddl.String(), "\n"), nil +} + +func (s *server) getTableDialectDDL(schema, table string) (string, error) { + schema = s.effectiveMetadataSchema(schema) + columns, err := s.getTableDialectColumns(schema, table) + if err != nil { + return "", err + } + if len(columns) == 0 { + return "", nil + } + comment, err := s.getTableComment(schema, table) + if err != nil { + return "", err + } + ttl := "INF" + if values, queryErr := s.queryValues("SHOW TABLES DETAILS FROM "+quoteTableIdentifier(schema), schema, metadataQueryLimit, 0); queryErr == nil { + for _, row := range values.Rows { + if firstMetadataRowString(values, row, "TableName", "table_name") == table { + if value := firstMetadataRowString(values, row, "TTL(ms)", "ttl(ms)"); value != "" { + ttl = value + } + break + } + } + } + definitions := make([]string, 0, len(columns)) + for _, column := range columns { + definition := " " + quoteTableIdentifier(column.Name) + " " + column.DataType + if column.Extra != nil && *column.Extra != "" { + definition += " " + *column.Extra + } + if column.Comment != nil && *column.Comment != "" { + definition += " COMMENT " + quoteSQLString(*column.Comment) + } + definitions = append(definitions, definition) + } + qualified := quoteTableIdentifier(schema) + "." + quoteTableIdentifier(table) + ddl := "DROP TABLE IF EXISTS " + qualified + ";\n\nCREATE TABLE " + qualified + " (\n" + strings.Join(definitions, ",\n") + "\n)" + if comment != nil && *comment != "" { + ddl += " COMMENT " + quoteSQLString(*comment) + } + if strings.EqualFold(ttl, "INF") { + ddl += " WITH (TTL='INF')" + } else if _, parseErr := strconv.ParseInt(ttl, 10, 64); parseErr == nil { + ddl += " WITH (TTL=" + ttl + ")" + } + return ddl + ";", nil +} + +func (s *server) completionAssistantSearch(input completionAssistantRequest) (completionAssistantResponse, error) { + limit := input.MaxResults + if limit <= 0 || limit > 1000 { + limit = 100 + } + candidates := make([]completionAssistantCandidate, 0, limit+1) + kinds := metadataStringSet(input.ObjectKinds) + if (kinds["column"] || kinds["property"] || len(kinds) == 0 && input.ParentName != "") && input.ParentName != "" { + schema := input.ParentSchema + if schema == "" { + schema = input.Schema + } + columns, err := s.getColumns(schema, input.ParentName) + if err != nil { + return completionAssistantResponse{}, err + } + for _, column := range columns { + if !completionNameMatches(column.Name, input) { + continue + } + dataType := column.DataType + candidates = append(candidates, completionAssistantCandidate{ + Name: column.Name, Kind: "COLUMN", Database: stringPtr(input.Database), Schema: stringPtr(schema), + ParentSchema: stringPtr(schema), ParentName: stringPtr(input.ParentName), DataType: &dataType, + }) + if len(candidates) > limit { + return completionAssistantResponse{Candidates: candidates[:limit], Incomplete: true}, nil + } + } + return completionAssistantResponse{Candidates: candidates}, nil + } + schemas := []string{input.Schema} + if input.GlobalSearch || input.Schema == "" { + var err error + schemas, err = s.listSchemas() + if err != nil { + return completionAssistantResponse{}, err + } + } + for _, schema := range schemas { + objects, err := s.listObjects(schema, metadataListConstraints{ObjectTypes: input.ObjectKinds}) + if err != nil { + return completionAssistantResponse{}, err + } + for _, object := range objects { + if !completionNameMatches(object.Name, input) { + continue + } + candidates = append(candidates, completionAssistantCandidate{ + Name: object.Name, Kind: object.ObjectType, Database: stringPtr(input.Database), + Schema: stringPtr(schema), Comment: object.Comment, + }) + if len(candidates) > limit { + return completionAssistantResponse{Candidates: candidates[:limit], Incomplete: true}, nil + } + } + } + return completionAssistantResponse{Candidates: candidates}, nil +} + +func (s *server) effectiveMetadataSchema(schema string) string { + if value := strings.TrimSpace(schema); value != "" { + return value + } + return strings.TrimSpace(s.config.Database) +} + +func metadataListConstraintsFromParams(params map[string]json.RawMessage) metadataListConstraints { + return metadataListConstraints{ + Filter: stringParam(params, "filter"), + Limit: intParam(params, "limit"), + Offset: intParam(params, "offset"), + ObjectTypes: stringSliceParam(params, "object_types"), + } +} + +func metadataNameMatches(name, filter string) bool { + return filter == "" || strings.Contains(strings.ToLower(name), strings.ToLower(filter)) +} + +func acceptsIoTDBTable(objectTypes []string) bool { + if len(objectTypes) == 0 { + return true + } + for _, objectType := range objectTypes { + switch strings.ToLower(strings.TrimSpace(objectType)) { + case "table", "base table", "device", "timeseries": + return true + } + } + return false +} + +func applyMetadataWindow[T any](values []T, offset, limit int) []T { + if offset < 0 { + offset = 0 + } + if offset >= len(values) { + return []T{} + } + values = values[offset:] + if limit > 0 && limit < len(values) { + values = values[:limit] + } + return values +} + +func completionNameMatches(name string, input completionAssistantRequest) bool { + mask := input.Mask + if mask == "" { + return true + } + if !input.CaseSensitive { + name = strings.ToLower(name) + mask = strings.ToLower(mask) + } + switch strings.ToLower(input.MatchMode) { + case "exact": + return name == mask + case "contains": + return strings.Contains(name, mask) + default: + return strings.HasPrefix(name, mask) + } +} + +func metadataStringSet(values []string) map[string]bool { + result := make(map[string]bool, len(values)) + for _, value := range values { + result[strings.ToLower(strings.TrimSpace(value))] = true + } + return result +} + +func metadataRowString(values queryResult, row []any, column string) string { + index := metadataColumnIndex(values.Columns, column) + if index < 0 || index >= len(row) || row[index] == nil { + return "" + } + return strings.TrimSpace(fmt.Sprint(row[index])) +} + +func firstMetadataRowString(values queryResult, row []any, columns ...string) string { + for _, column := range columns { + if value := metadataRowString(values, row, column); value != "" { + return value + } + } + return "" +} + +func metadataRowBool(values queryResult, row []any, column string) bool { + index := metadataColumnIndex(values.Columns, column) + if index < 0 || index >= len(row) || row[index] == nil { + return false + } + switch value := row[index].(type) { + case bool: + return value + case string: + parsed, _ := strconv.ParseBool(value) + return parsed + default: + parsed, _ := strconv.ParseBool(fmt.Sprint(value)) + return parsed + } +} + +func metadataColumnIndex(columns []string, target string) int { + for index, column := range columns { + if strings.EqualFold(strings.TrimSpace(column), strings.TrimSpace(target)) { + return index + } + } + return -1 +} + +func treeDevicePath(schema, table string) string { + schema = strings.TrimSuffix(strings.TrimSpace(schema), ".") + table = strings.TrimPrefix(strings.TrimSpace(table), ".") + if table == "" { + return "" + } + if schema == "" || table == schema || strings.HasPrefix(table, schema+".") { + return table + } + return schema + "." + table +} + +func quoteTreePath(value string) string { + nodes := strings.Split(strings.TrimSpace(value), ".") + for index, node := range nodes { + nodes[index] = quoteTreeNode(node) + } + return strings.Join(nodes, ".") +} + +func quoteTreeNode(value string) string { + if simpleTreeNodePattern.MatchString(value) { + return value + } + return "`" + strings.ReplaceAll(value, "`", "``") + "`" +} + +func quoteSQLString(value string) string { + return "'" + strings.ReplaceAll(value, "'", "''") + "'" +} + +func optionalMetadataString(value string) *string { + if strings.TrimSpace(value) == "" { + return nil + } + return &value +} + +func stringPtr(value string) *string { + return &value +} diff --git a/agents/drivers/iotdb/protocol_error.go b/agents/drivers/iotdb/protocol_error.go new file mode 100644 index 0000000000..b6020e0a3c --- /dev/null +++ b/agents/drivers/iotdb/protocol_error.go @@ -0,0 +1,138 @@ +package main + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "strings" + + "github.com/apache/iotdb-client-go/v2/client" +) + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` + Data *rpcErrorData `json:"data,omitempty"` +} + +type rpcErrorData struct { + Category string `json:"category"` + Retryable bool `json:"retryable"` + SessionDisposition string `json:"sessionDisposition"` + Stage string `json:"stage"` + ContractVersion int `json:"contractVersion"` + OperationOutcome string `json:"operationOutcome"` + SQLState string `json:"sqlState,omitempty"` + VendorCode int32 `json:"vendorCode,omitempty"` + ExceptionClass string `json:"exceptionClass,omitempty"` + AgentSessionID string `json:"agentSessionId,omitempty"` +} + +func classifyRPCError(method, agentSessionID string, err error) *rpcError { + stage := rpcErrorStage(method) + data := &rpcErrorData{ + Category: "protocol", + Retryable: false, + SessionDisposition: "keep", + Stage: stage, + ContractVersion: 1, + OperationOutcome: rpcOperationOutcome(stage), + ExceptionClass: safeRPCDiagnostic(fmt.Sprintf("%T", err), 160), + AgentSessionID: strings.TrimSpace(agentSessionID), + } + + var executionError *client.ExecutionError + if errors.As(err, &executionError) { + data.Category = "sql" + data.VendorCode = executionError.Code + data.SQLState = fmt.Sprintf("IOTDB-%d", executionError.Code) + lower := strings.ToLower(executionError.Message) + if strings.Contains(lower, "timeout") { + data.Category = "timeout" + data.Retryable = true + } + } else if errors.Is(err, context.Canceled) { + data.Category = "canceled" + data.SessionDisposition = "quarantine" + } else if errors.Is(err, context.DeadlineExceeded) || isTimeoutError(err) { + data.Category = "timeout" + data.SessionDisposition = "quarantine" + } else if isConnectionError(err) { + data.Category = "connection" + data.Retryable = stage == "connect" || stage == "validate" + if stage != "connect" { + data.SessionDisposition = "quarantine" + } + } + return &rpcError{Code: -1, Message: err.Error(), Data: data} +} + +func rpcErrorStage(method string) string { + switch method { + case "connect", "open_session", "test_connection": + return "connect" + case "validate_connection", "validate_session": + return "validate" + case "cancel_session": + return "cancel" + case "close_session", "disconnect", "close_query_session", "close_table_read_session", "shutdown": + return "close" + case "fetch_query_page", "fetch_table_read_page": + return "fetch" + case "handshake", "": + return "request" + default: + return "execute" + } +} + +func rpcOperationOutcome(stage string) string { + if stage == "request" || stage == "connect" || stage == "validate" { + return "not_started" + } + return "unknown" +} + +func isTimeoutError(err error) bool { + var timeout interface{ Timeout() bool } + return errors.As(err, &timeout) && timeout.Timeout() +} + +func isConnectionError(err error) bool { + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { + return true + } + var networkError *net.OpError + if errors.As(err, &networkError) { + return true + } + var executionError *client.ExecutionError + if errors.As(err, &executionError) { + return false + } + lower := strings.ToLower(err.Error()) + for _, marker := range []string{ + "connection refused", "connection reset", "broken pipe", "connection closed", "connection lost", + "unexpected eof", "no route to host", "no server can connect", "transport is not open", + } { + if strings.Contains(lower, marker) { + return true + } + } + return false +} + +func safeRPCDiagnostic(value string, maxLength int) string { + var result strings.Builder + for _, char := range value { + if result.Len() >= maxLength { + break + } + if char >= 0x21 && char <= 0x7e { + result.WriteRune(char) + } + } + return result.String() +} diff --git a/agents/drivers/iotdb/query.go b/agents/drivers/iotdb/query.go new file mode 100644 index 0000000000..59380d6676 --- /dev/null +++ b/agents/drivers/iotdb/query.go @@ -0,0 +1,515 @@ +package main + +import ( + "context" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math" + "strings" + "time" + + "github.com/apache/iotdb-client-go/v2/client" +) + +type querySession struct { + dataset *client.SessionDataSet + client *sessionClient + columns []string + columnTypes []string + prefetched []any + remaining int +} + +type queryPageRead struct { + state *querySession + rows [][]any + hasMore bool + truncated bool +} + +func (s *server) validateConnection() error { + values, err := s.queryValues("SHOW VERSION", "", 1, 0) + if err != nil { + return err + } + if len(values.Rows) == 0 { + return errors.New("IoTDB SHOW VERSION returned no rows") + } + return nil +} + +func (s *server) executeQuery(options queryOptions) (queryResult, error) { + started := time.Now() + sql := trimStatementSQL(options.SQL) + if sql == "" { + return queryResult{}, errors.New("sql is required") + } + if !isQueryStatement(sql) { + if err := s.executeNonQuery(sql, effectiveDatabase(options), options.TimeoutSecs); err != nil { + return queryResult{}, err + } + return queryResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + AffectedRows: 0, + ExecutionTimeMS: time.Since(started).Milliseconds(), + }, nil + } + + limit := options.MaxRows + if limit <= 0 { + limit = defaultMaxRows + } + values, err := s.queryValues(sql, effectiveDatabase(options), limit, options.TimeoutSecs) + if err != nil { + return queryResult{}, err + } + values.ExecutionTimeMS = time.Since(started).Milliseconds() + return values, nil +} + +func (s *server) queryValues(sql, database string, limit, timeoutSecs int) (queryResult, error) { + connected, err := s.ensureClient() + if err != nil { + return queryResult{}, err + } + timeout := timeoutMilliseconds(timeoutSecs) + return runCancelable(s, timeoutSecs, connected, func(ctx context.Context) (queryResult, error) { + if err := s.applyDatabaseContext(ctx, connected, database); err != nil { + return queryResult{}, err + } + dataset, err := connected.session.ExecuteQueryStatement(ctx, sql, timeout) + if err != nil { + return queryResult{}, err + } + defer dataset.Close() + columns := append([]string(nil), dataset.GetColumnNames()...) + columnTypes := normalizedColumnTypes(dataset.GetColumnTypes(), len(columns)) + rows, truncated, err := readDataset(ctx, dataset, connected.dialect, columns, columnTypes, limit) + if err != nil { + return queryResult{}, err + } + return queryResult{ + Columns: columns, + ColumnTypes: columnTypes, + Rows: rows, + AffectedRows: 0, + Truncated: truncated, + }, nil + }) +} + +func (s *server) executeNonQuery(sql, database string, timeoutSecs int) error { + connected, err := s.ensureClient() + if err != nil { + return err + } + _, err = runCancelable(s, timeoutSecs, connected, func(ctx context.Context) (struct{}, error) { + if err := s.applyDatabaseContext(ctx, connected, database); err != nil { + return struct{}{}, err + } + return struct{}{}, connected.session.ExecuteNonQueryStatement(ctx, sql) + }) + return err +} + +func (s *server) executeQueryPage(options queryOptions, requestedPageSize int) (queryPageResult, error) { + started := time.Now() + sql := trimStatementSQL(options.SQL) + if sql == "" { + return queryPageResult{}, errors.New("sql is required") + } + if !isQueryStatement(sql) { + if err := s.executeNonQuery(sql, effectiveDatabase(options), options.TimeoutSecs); err != nil { + return queryPageResult{}, err + } + return queryPageResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + ExecutionTimeMS: time.Since(started).Milliseconds(), + }, nil + } + + connected, err := s.ensureClient() + if err != nil { + return queryPageResult{}, err + } + timeout := timeoutMilliseconds(options.TimeoutSecs) + limit := options.MaxRows + if limit <= 0 { + limit = defaultMaxRows + } + pageSize := requestedPageSize + if pageSize <= 0 { + pageSize = options.FetchSize + } + if pageSize <= 0 { + pageSize = defaultPageSize + } + var dataset *client.SessionDataSet + page, err := runCancelable(s, options.TimeoutSecs, connected, func(ctx context.Context) (queryPageRead, error) { + if err := s.applyDatabaseContext(ctx, connected, effectiveDatabase(options)); err != nil { + return queryPageRead{}, err + } + dataset, err = connected.session.ExecuteQueryStatement(ctx, sql, timeout) + if err != nil { + return queryPageRead{}, err + } + queryState := &querySession{ + dataset: dataset, + client: connected, + columns: append([]string(nil), dataset.GetColumnNames()...), + columnTypes: normalizedColumnTypes(dataset.GetColumnTypes(), len(dataset.GetColumnNames())), + remaining: limit, + } + rows, hasMore, truncated, err := s.readQuerySessionPage(ctx, queryState, pageSize) + return queryPageRead{state: queryState, rows: rows, hasMore: hasMore, truncated: truncated}, err + }) + if err != nil { + if dataset != nil { + _ = dataset.Close() + } + return queryPageResult{}, err + } + result := queryPageResult{ + Columns: page.state.columns, + ColumnTypes: page.state.columnTypes, + Rows: page.rows, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: page.truncated, + HasMore: page.hasMore, + } + if page.hasMore { + s.nextSessionID++ + sessionID := fmt.Sprintf("iotdb-go-%d", s.nextSessionID) + s.querySessions[sessionID] = page.state + result.SessionID = &sessionID + } else { + _ = dataset.Close() + } + return result, nil +} + +func (s *server) fetchQueryPage(sessionID string, requestedPageSize int) (queryPageResult, error) { + started := time.Now() + queryState := s.querySessions[sessionID] + if queryState == nil { + return queryPageResult{}, fmt.Errorf("query session not found: %s", sessionID) + } + pageSize := requestedPageSize + if pageSize <= 0 { + pageSize = defaultPageSize + } + page, err := runCancelable(s, 0, queryState.client, func(ctx context.Context) (queryPageRead, error) { + rows, hasMore, truncated, err := s.readQuerySessionPage(ctx, queryState, pageSize) + return queryPageRead{state: queryState, rows: rows, hasMore: hasMore, truncated: truncated}, err + }) + if err != nil { + _ = s.closeQuerySession(sessionID) + return queryPageResult{}, err + } + result := queryPageResult{ + Columns: queryState.columns, + ColumnTypes: queryState.columnTypes, + Rows: page.rows, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: page.truncated, + HasMore: page.hasMore, + } + if page.hasMore { + result.SessionID = &sessionID + } else { + _ = s.closeQuerySession(sessionID) + } + return result, nil +} + +func (s *server) readQuerySessionPage(ctx context.Context, state *querySession, pageSize int) ([][]any, bool, bool, error) { + if pageSize <= 0 { + pageSize = defaultPageSize + } + target := min(pageSize, state.remaining) + rows := make([][]any, 0, target) + if state.prefetched != nil && target > 0 { + rows = append(rows, state.prefetched) + state.prefetched = nil + state.remaining-- + } + for len(rows) < target { + row, hasNext, err := s.nextDatasetRow(ctx, state) + if err != nil { + return nil, false, false, err + } + if !hasNext { + return rows, false, false, nil + } + rows = append(rows, row) + state.remaining-- + } + if state.remaining <= 0 { + _, hasNext, err := s.nextDatasetRow(ctx, state) + if err != nil { + return nil, false, false, err + } + return rows, false, hasNext, nil + } + next, hasNext, err := s.nextDatasetRow(ctx, state) + if err != nil { + return nil, false, false, err + } + if !hasNext { + return rows, false, false, nil + } + state.prefetched = next + return rows, true, false, nil +} + +func (s *server) nextDatasetRow(ctx context.Context, state *querySession) ([]any, bool, error) { + if err := ctx.Err(); err != nil { + return nil, false, err + } + hasNext, err := state.dataset.Next() + if err != nil || !hasNext { + return nil, false, err + } + if err := ctx.Err(); err != nil { + return nil, false, err + } + row, err := datasetRow(state.dataset, state.client.dialect, state.columns, state.columnTypes) + return row, true, err +} + +func (s *server) closeQuerySession(sessionID string) bool { + state := s.querySessions[sessionID] + delete(s.querySessions, sessionID) + if state == nil { + return false + } + _ = state.dataset.Close() + return true +} + +func (s *server) closeAllQuerySessions() error { + var firstErr error + for sessionID, state := range s.querySessions { + if err := state.dataset.Close(); err != nil && firstErr == nil { + firstErr = err + } + delete(s.querySessions, sessionID) + } + return firstErr +} + +func (s *server) executeStatements(params map[string]json.RawMessage, transaction bool) (queryResult, error) { + started := time.Now() + statements := stringSliceParam(params, "statements") + if len(statements) == 0 { + return queryResult{}, errors.New("statements are required") + } + database := strings.TrimSpace(stringParam(params, "schema")) + for _, statement := range statements { + sql := trimStatementSQL(statement) + if sql == "" { + continue + } + if isQueryStatement(sql) { + values, err := s.queryValues(sql, database, defaultMaxRows, 0) + if err != nil { + return queryResult{}, err + } + _ = values + } else if err := s.executeNonQuery(sql, database, 0); err != nil { + return queryResult{}, err + } + } + return queryResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + AffectedRows: 0, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: false, + }, nil +} + +func (s *server) applyDatabaseContext(ctx context.Context, connected *sessionClient, database string) error { + if connected.dialect != client.TableSqlDialect { + return nil + } + database = strings.TrimSpace(database) + if database == "" { + database = strings.TrimSpace(s.config.Database) + } + if database == "" || strings.EqualFold(database, "information_schema") { + return nil + } + return connected.session.ExecuteNonQueryStatement(ctx, "USE "+quoteTableIdentifier(database)) +} + +func runCancelable[T any](s *server, timeoutSecs int, connected *sessionClient, operation func(context.Context) (T, error)) (T, error) { + var zero T + ctx := context.Background() + var cancel context.CancelFunc + if timeoutSecs > 0 { + ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSecs)*time.Second) + } else { + ctx, cancel = context.WithCancel(ctx) + } + s.setActiveOperation(cancel) + defer s.clearActiveOperation(cancel) + value, err := operation(ctx) + if ctxErr := ctx.Err(); ctxErr != nil { + s.invalidateClient(connected) + return zero, ctxErr + } + if err != nil && isConnectionError(err) { + s.invalidateClient(connected) + } + return value, err +} + +func readDataset(ctx context.Context, dataset *client.SessionDataSet, dialect string, columns, columnTypes []string, limit int) ([][]any, bool, error) { + if limit <= 0 { + limit = defaultMaxRows + } + rows := make([][]any, 0, min(limit, defaultPageSize)) + for len(rows) < limit { + if err := ctx.Err(); err != nil { + return nil, false, err + } + hasNext, err := dataset.Next() + if err != nil { + return nil, false, err + } + if !hasNext { + return rows, false, nil + } + if err := ctx.Err(); err != nil { + return nil, false, err + } + row, err := datasetRow(dataset, dialect, columns, columnTypes) + if err != nil { + return nil, false, err + } + rows = append(rows, row) + } + if err := ctx.Err(); err != nil { + return nil, false, err + } + hasNext, err := dataset.Next() + return rows, hasNext, err +} + +func datasetRow(dataset *client.SessionDataSet, dialect string, columns, columnTypes []string) ([]any, error) { + row := make([]any, len(columns)) + for index := range columns { + columnIndex := int32(index + 1) + columnType := "" + if index < len(columnTypes) { + columnType = strings.ToUpper(columnTypes[index]) + } + if dialect == client.TreeSqlDialect && index == 0 && strings.EqualFold(columns[index], "Time") { + value, err := dataset.GetStringByIndex(columnIndex) + if err != nil { + return nil, err + } + row[index] = value + continue + } + value, err := dataset.GetObjectByIndex(columnIndex) + if err != nil { + return nil, err + } + row[index] = normalizeIoTDBValue(value, columnType) + } + return row, nil +} + +func normalizeIoTDBValue(value any, columnType string) any { + switch typed := value.(type) { + case nil: + return nil + case time.Time: + if columnType == "DATE" { + return typed.Format("2006-01-02") + } + return typed.Format(time.RFC3339Nano) + case []byte: + return hex.EncodeToString(typed) + case float32: + if math.IsNaN(float64(typed)) || math.IsInf(float64(typed), 0) { + return fmt.Sprint(typed) + } + return typed + case float64: + if math.IsNaN(typed) || math.IsInf(typed, 0) { + return fmt.Sprint(typed) + } + return typed + default: + return typed + } +} + +func normalizedColumnTypes(values []string, columnCount int) []string { + result := make([]string, columnCount) + for index := range result { + if index < len(values) { + result[index] = strings.ToUpper(values[index]) + } + } + return result +} + +func timeoutMilliseconds(timeoutSecs int) *int64 { + if timeoutSecs <= 0 { + return nil + } + value := int64(timeoutSecs) * 1000 + return &value +} + +func effectiveDatabase(options queryOptions) string { + if value := strings.TrimSpace(options.Schema); value != "" { + return value + } + return strings.TrimSpace(options.Database) +} + +func trimStatementSQL(sql string) string { + trimmed := strings.TrimSpace(sql) + for strings.HasSuffix(trimmed, ";") { + trimmed = strings.TrimSpace(strings.TrimSuffix(trimmed, ";")) + } + return trimmed +} + +func isQueryStatement(sql string) bool { + fields := strings.Fields(strings.TrimSpace(sql)) + if len(fields) == 0 { + return false + } + switch strings.ToUpper(fields[0]) { + case "SELECT", "SHOW", "DESC", "DESCRIBE", "EXPLAIN", "WITH": + return true + default: + return false + } +} + +func quoteTableIdentifier(value string) string { + return `"` + strings.ReplaceAll(value, `"`, `""`) + `"` +} + +func rawParams(values map[string]any) map[string]json.RawMessage { + result := make(map[string]json.RawMessage, len(values)) + for key, value := range values { + data, _ := json.Marshal(value) + result[key] = data + } + return result +} diff --git a/agents/drivers/iotdb/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java b/agents/drivers/iotdb/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java deleted file mode 100644 index cfe65870df..0000000000 --- a/agents/drivers/iotdb/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java +++ /dev/null @@ -1,46 +0,0 @@ -package com.dbx.agent.iotdb; - -import com.dbx.agent.ConfiguredJdbcAgent; -import com.dbx.agent.DatabaseInfo; -import com.dbx.agent.JdbcAgentProfile; -import com.dbx.agent.MultiSessionJsonRpcServer; -import java.sql.ResultSet; -import java.sql.Statement; -import java.util.ArrayList; -import java.util.List; - -public final class IoTDBAgent extends ConfiguredJdbcAgent { - public static final JdbcAgentProfile IOTDB_PROFILE = new JdbcAgentProfile( - "org.apache.iotdb.jdbc.IoTDBDriver", - "jdbc:iotdb://{host}:{port}/", - 6667, - true - ); - - public IoTDBAgent() { - super(IOTDB_PROFILE); - } - - // IoTDB 的 database 是 `root.xxx` 路径前缀；JDBC metadata 会把同一路径同时投影成 catalog/schema。 - // 这里使用 IoTDB 原生语句作为 database 列表来源，避免依赖驱动的关系型兼容层。 - @Override - public List listDatabases() { - return unchecked(() -> { - List result = new ArrayList<>(); - try (Statement stmt = requireConnection().createStatement(); - ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { - while (rs.next()) { - String name = rs.getString(1); - if (name != null && !name.trim().isEmpty()) { - result.add(new DatabaseInfo(name.trim())); - } - } - } - return result; - }); - } - - public static void main(String[] args) { - new MultiSessionJsonRpcServer(IoTDBAgent::new).run(); - } -} diff --git a/agents/drivers/iotdb/src/test/java/com/dbx/agent/iotdb/IoTDBAgentTest.java b/agents/drivers/iotdb/src/test/java/com/dbx/agent/iotdb/IoTDBAgentTest.java deleted file mode 100644 index acfc62f0d7..0000000000 --- a/agents/drivers/iotdb/src/test/java/com/dbx/agent/iotdb/IoTDBAgentTest.java +++ /dev/null @@ -1,32 +0,0 @@ -package com.dbx.agent.iotdb; - -import com.dbx.agent.ConnectParams; -import com.dbx.agent.DatabaseAgent; -import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -class IoTDBAgentTest extends JdbcFakeExecutionBehaviorTest { - @Override - protected DatabaseAgent createAgent() { - return new IoTDBAgent(); - } - - @Override - protected String resultSetSql() { - return "SELECT * FROM sample_table"; - } - - @Test - void buildsOfficialJdbcUrl() { - ConnectParams params = new ConnectParams(); - params.setHost("iotdb.example.com"); - params.setPort(6667); - params.setUrl_params("sql_dialect=table&useSSL=false"); - - Assertions.assertEquals( - "jdbc:iotdb://iotdb.example.com:6667/?sql_dialect=table&useSSL=false", - IoTDBAgent.IOTDB_PROFILE.buildUrl(params) - ); - } -} diff --git a/agents/metadata-constraint-coverage.tsv b/agents/metadata-constraint-coverage.tsv index ce4a321ad4..8c13e2d854 100644 --- a/agents/metadata-constraint-coverage.tsv +++ b/agents/metadata-constraint-coverage.tsv @@ -11,6 +11,7 @@ goldendb native-pushdown java-sql Uses information_schema with type, filter, sta h2 native-pushdown java-sql Uses INFORMATION_SCHEMA with type, filter, stable order, and LIMIT/OFFSET. hive intentional-fallback java-sql Uses Hive metadata paths where stable server-side fuzzy paging is not portable, common constraints filter locally. informix native-pushdown java-sql Uses systables/sysprocedures with type, filter, stable order, and SKIP/FIRST. +iotdb shared-fallback native-go Uses SHOW DEVICES or SHOW TABLES DETAILS for Tree/Table metadata, then applies stable filtering and paging in the native agent. kingbase-go shared-fallback native-go Uses sys_catalog/information_schema for compatibility-aware metadata, then applies stable filtering and paging in the native agent. vastbase-go shared-fallback native-go Uses pg_catalog/information_schema for Vastbase metadata, then applies stable filtering and paging in the native agent. kylin intentional-fallback java-jdbc-metadata Uses JDBC metadata from Kylin driver; no portable server-side paging API, common constraints filter locally. diff --git a/agents/scripts/driver_release_packages_test.py b/agents/scripts/driver_release_packages_test.py index 1de27a3a32..a9e1b78f62 100644 --- a/agents/scripts/driver_release_packages_test.py +++ b/agents/scripts/driver_release_packages_test.py @@ -35,6 +35,7 @@ class DriverReleasePackagesTest(unittest.TestCase): "oracle": "0.1.10", "xugu": "0.1.20", "kingbase": "0.1.34", + "iotdb": "0.1.30", "neo4j": "0.1.40", "vastbase": "0.1.37", "duckdb": "0.1.0", @@ -245,6 +246,21 @@ class DriverReleasePackagesTest(unittest.TestCase): self.assertFalse(source.exists()) self.assertEqual(versioned.read_bytes(), b"\xcf\xfa\xed\xfetest-neo4j-agent") + def test_versions_iotdb_native_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + release_dir = Path(temp_dir) + source = release_dir / "dbx-agent-iotdb-linux-x64" + source.write_bytes(b"\x7fELFtest-iotdb-agent") + versions = {driver: "0.1.0" for driver in NATIVE_DRIVERS} + versions["iotdb"] = "0.1.30" + + renamed = version_agent_artifacts(release_dir, versions) + versioned = release_dir / "dbx-agent-iotdb-0.1.30-linux-x64" + + self.assertEqual(renamed, [versioned]) + self.assertFalse(source.exists()) + self.assertEqual(versioned.read_bytes(), b"\x7fELFtest-iotdb-agent") + def test_full_offline_bundle_includes_supported_windows_artifacts(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: release_dir = Path(temp_dir) diff --git a/agents/scripts/validate_agents.py b/agents/scripts/validate_agents.py index a70ac1feda..8a77a431cc 100644 --- a/agents/scripts/validate_agents.py +++ b/agents/scripts/validate_agents.py @@ -17,6 +17,7 @@ NATIVE_ONLY_AGENT_MODULES = { "duckdb": "drivers/duckdb", "oracle": "drivers/oracle-go", "kingbase": "drivers/kingbase-go", + "iotdb": "drivers/iotdb", "neo4j": "drivers/neo4j-go", "vastbase": "drivers/vastbase-go", "tdengine": "drivers/tdengine", diff --git a/agents/scripts/validate_agents_test.py b/agents/scripts/validate_agents_test.py index 1a82f89707..849d464388 100644 --- a/agents/scripts/validate_agents_test.py +++ b/agents/scripts/validate_agents_test.py @@ -140,6 +140,7 @@ class ValidateAgentsTest(unittest.TestCase): "cassandra-go", "oracle-go", "kingbase-go", + "iotdb", "neo4j-go", "vastbase-go", "xugu", @@ -149,7 +150,7 @@ class ValidateAgentsTest(unittest.TestCase): ): (root / "drivers" / driver).mkdir(parents=True) (root / "versions.json").write_text( - json.dumps({"h2": "0.1.0", "cassandra": "0.1.0", "oracle": "0.1.0", "kingbase": "0.1.0", "neo4j": "0.1.0", "vastbase": "0.1.0", "xugu": "0.1.0", "rabbitmq": "0.1.0", "tdengine": "0.1.0"}), + json.dumps({"h2": "0.1.0", "cassandra": "0.1.0", "oracle": "0.1.0", "kingbase": "0.1.0", "iotdb": "0.1.0", "neo4j": "0.1.0", "vastbase": "0.1.0", "xugu": "0.1.0", "rabbitmq": "0.1.0", "tdengine": "0.1.0"}), encoding="utf-8", ) diff --git a/agents/scripts/version_agent_artifacts.py b/agents/scripts/version_agent_artifacts.py index 8e767c03b6..f272f6586a 100644 --- a/agents/scripts/version_agent_artifacts.py +++ b/agents/scripts/version_agent_artifacts.py @@ -4,7 +4,7 @@ import json from pathlib import Path -NATIVE_DRIVERS = ("cassandra", "oracle", "xugu", "kingbase", "neo4j", "vastbase", "duckdb", "rabbitmq", "tdengine") +NATIVE_DRIVERS = ("cassandra", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "duckdb", "rabbitmq", "tdengine") PLATFORMS = ( "macos-aarch64", "macos-x64", diff --git a/agents/settings.gradle b/agents/settings.gradle index 2cb648d850..ca1f6648c4 100644 --- a/agents/settings.gradle +++ b/agents/settings.gradle @@ -6,7 +6,7 @@ def driverModules = [ 'teradata', 'vertica', 'firebird', 'exasol', 'oceanbase-oracle', 'gbase8a', 'gbase8s', 'bigquery', 'kylin', 'sundb', 'h2', 'h2-legacy', 'snowflake', 'trino', 'hive', 'spark', 'db2', 'informix', 'mongodb', 'highgo', 'uxdb', 'yashandb', 'oscar', - 'iris', 'iotdb', 'etcd', 'zookeeper', 'kafka', 'rocketmq', 'sqlserver-legacy' + 'iris', 'etcd', 'zookeeper', 'kafka', 'rocketmq', 'sqlserver-legacy' ] include(*(infrastructureModules + driverModules))