diff --git a/.github/workflows/agents-release.yml b/.github/workflows/agents-release.yml new file mode 100644 index 0000000000..aae3f8dbe2 --- /dev/null +++ b/.github/workflows/agents-release.yml @@ -0,0 +1,624 @@ +name: Agents Release + +on: + push: + tags: ["agents-v*"] + +permissions: + contents: write + +jobs: + bump-versions: + runs-on: ubuntu-latest + outputs: + versions: ${{ steps.bump.outputs.versions }} + prev_versions: ${{ steps.bump.outputs.prev_versions }} + prev_tag: ${{ steps.bump.outputs.prev_tag }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Detect changes and bump versions + id: bump + shell: bash + run: | + PREV_TAG=$(git tag --sort=-creatordate | grep '^agents-v' | sed -n '2p') + if [ -z "$PREV_TAG" ]; then + echo "No previous tag found, all modules are new" + PREV_TAG=$(git rev-list --max-parents=0 HEAD) + fi + echo "Comparing $PREV_TAG..HEAD" + + VERSIONS=$(cat agents/versions.json) + PREV_VERSIONS="$VERSIONS" + CHANGED=false + + for dir in $(echo "$VERSIONS" | python3 -c "import sys,json; [print(k) for k in json.load(sys.stdin)]"); do + if [ -d "agents/drivers/$dir" ] || [ -d "agents/$dir" ] || [ "$dir" = "common" ]; then + if [ -d "agents/drivers/$dir" ]; then + CHECK_DIR="drivers/$dir" + else + CHECK_DIR="$dir" + fi + EXTRA_CHECK_DIR="" + if [ "$dir" = "oracle" ]; then + EXTRA_CHECK_DIR="drivers/oracle-go" + fi + DIFF_PATHS=("agents/$CHECK_DIR/") + if [ -n "$EXTRA_CHECK_DIR" ]; then + DIFF_PATHS+=("agents/$EXTRA_CHECK_DIR/") + fi + # common changes affect all modules + if git diff --quiet "$PREV_TAG"..HEAD -- agents/common/ 2>/dev/null; then + COMMON_CHANGED=false + else + COMMON_CHANGED=true + fi + + if git diff --quiet "$PREV_TAG"..HEAD -- "${DIFF_PATHS[@]}" 2>/dev/null && [ "$COMMON_CHANGED" = "false" ]; then + echo " $dir: no changes" + else + echo " $dir: CHANGED" + OLD_VER=$(echo "$VERSIONS" | python3 -c "import sys,json; print(json.load(sys.stdin).get('$dir','0.1.0'))") + IFS='.' read -r major minor patch <<< "$OLD_VER" + NEW_VER="$major.$minor.$((patch + 1))" + VERSIONS=$(echo "$VERSIONS" | python3 -c " + import sys, json + d = json.load(sys.stdin) + d['$dir'] = '$NEW_VER' + json.dump(d, sys.stdout, indent=2) + print() + ") + CHANGED=true + echo " $dir: $OLD_VER -> $NEW_VER" + fi + fi + done + + echo "$VERSIONS" > agents/versions.json + cat agents/versions.json + echo "versions=$(cat agents/versions.json | python3 -c 'import sys,json; print(json.dumps(json.load(sys.stdin)))')" >> "$GITHUB_OUTPUT" + echo "prev_versions=$(echo "$PREV_VERSIONS" | python3 -c 'import sys,json; print(json.dumps(json.load(sys.stdin)))')" >> "$GITHUB_OUTPUT" + echo "prev_tag=$PREV_TAG" >> "$GITHUB_OUTPUT" + + if [ "$CHANGED" = "true" ]; then + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + git add agents/versions.json + git commit -m "chore: bump module versions [skip ci]" + git push origin HEAD:main + fi + + build-agents: + needs: [bump-versions] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-java@v4 + with: + distribution: temurin + java-version: | + 8 + 21 + - uses: gradle/actions/setup-gradle@v4 + - run: ./gradlew shadowJar + working-directory: agents + - run: python3 scripts/validate_agent_jars.py + working-directory: agents + - uses: actions/upload-artifact@v4 + with: + name: agent-jars + path: "agents/drivers/*/build/libs/dbx-agent-*.jar" + + build-oracle-native: + needs: [bump-versions] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: "1.22.x" + - name: Test Oracle native agent + working-directory: agents/drivers/oracle-go + run: go test ./... + - name: Cross-compile Oracle native agent + shell: bash + run: | + mkdir -p release-native + cd agents/drivers/oracle-go + declare -A TARGETS=( + ["macos-aarch64"]="darwin/arm64" + ["macos-x64"]="darwin/amd64" + ["linux-aarch64"]="linux/arm64" + ["linux-x64"]="linux/amd64" + ["windows-aarch64"]="windows/arm64" + ["windows-x64"]="windows/amd64" + ) + for platform in "${!TARGETS[@]}"; do + IFS=/ read -r goos goarch <<< "${TARGETS[$platform]}" + output="../../../release-native/dbx-agent-oracle-${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: oracle-native + path: "release-native/dbx-agent-oracle-*" + + build-xugu-native: + needs: [bump-versions] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: "1.22.x" + - name: Test Xugu native agent + working-directory: agents/drivers/xugu + run: GONOSUMDB=gitee.com/XuguDB/go-xugu-driver go test ./... + - name: Cross-compile Xugu native agent + shell: bash + run: | + mkdir -p release-native + cd agents/drivers/xugu + 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-xugu-${platform}" + if [[ "$goos" == "windows" ]]; then + output="${output}.exe" + fi + echo "Building $platform ($goos/$goarch)" + GONOSUMDB=gitee.com/XuguDB/go-xugu-driver 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: xugu-native + path: "release-native/dbx-agent-xugu-*" + + build-jre: + runs-on: ubuntu-latest + strategy: + matrix: + include: + - jre-key: "21" + java-version: "21" + modules: "java.base,java.sql,java.sql.rowset,java.naming,java.management,java.desktop,java.security.jgss,java.security.sasl,jdk.charsets,java.scripting" + steps: + - uses: actions/setup-java@v4 + with: + distribution: temurin + java-version: ${{ matrix.java-version }} + - name: Build JRE for all platforms + shell: bash + run: | + MODULES="${{ matrix.modules }}" + JLINK_OPTS="--strip-debug --no-header-files --no-man-pages --compress=zip-6" + JRE_KEY="${{ matrix.jre-key }}" + + jlink --add-modules $MODULES $JLINK_OPTS --output dbx-jre + tar czf dbx-jre-${JRE_KEY}-linux-x64.tar.gz dbx-jre + rm -rf dbx-jre + + ADOPTIUM="https://api.adoptium.net/v3/binary/latest/${{ matrix.java-version }}/ga" + declare -A PLATFORMS=( + ["macos-aarch64"]="mac/aarch64" + ["macos-x64"]="mac/x64" + ["linux-aarch64"]="linux/aarch64" + ["windows-x64"]="windows/x64" + ) + for platform in "${!PLATFORMS[@]}"; do + os_arch="${PLATFORMS[$platform]}" + echo "=== Downloading JDK for $platform ($os_arch) ===" + if [[ "$platform" == windows-* ]]; then + curl -L -o jdk-$platform.zip "$ADOPTIUM/$os_arch/jdk/hotspot/normal/eclipse?project=jdk" + mkdir -p jdk-$platform + unzip -q jdk-$platform.zip -d jdk-$platform-tmp + mv jdk-$platform-tmp/*/* jdk-$platform/ || mv jdk-$platform-tmp/* jdk-$platform/ + rm -rf jdk-$platform-tmp jdk-$platform.zip + else + curl -L -o jdk-$platform.tar.gz "$ADOPTIUM/$os_arch/jdk/hotspot/normal/eclipse?project=jdk" + mkdir -p jdk-$platform + tar xzf jdk-$platform.tar.gz -C jdk-$platform --strip-components=1 + rm -f jdk-$platform.tar.gz + fi + JAVA_HOME_CROSS="jdk-$platform" + if [ -d "$JAVA_HOME_CROSS/Contents/Home/jmods" ]; then + JMODS="$JAVA_HOME_CROSS/Contents/Home/jmods" + else + JMODS="$JAVA_HOME_CROSS/jmods" + fi + LLVM_OBJCOPY=$(ls /usr/bin/llvm-objcopy-* 2>/dev/null | sort -V | tail -1) + if [ -z "$LLVM_OBJCOPY" ]; then LLVM_OBJCOPY=$(which llvm-objcopy 2>/dev/null || true); fi + CROSS_OPTS="" + if [ -n "$LLVM_OBJCOPY" ]; then + CROSS_OPTS="--strip-native-debug-symbols=objcopy=$LLVM_OBJCOPY" + else + CROSS_OPTS="--disable-plugin strip-native-debug-symbols" + fi + jlink --module-path "$JMODS" --add-modules $MODULES $JLINK_OPTS $CROSS_OPTS --output dbx-jre + tar czf dbx-jre-${JRE_KEY}-$platform.tar.gz dbx-jre + rm -rf dbx-jre jdk-$platform + done + + ls -lh dbx-jre-*.tar.gz + - uses: actions/upload-artifact@v4 + with: + name: jre-${{ matrix.jre-key }} + path: "dbx-jre-*.tar.gz" + + build-jre8: + runs-on: ubuntu-latest + steps: + - name: Download Adoptium JRE 8 for all platforms + shell: bash + run: | + ADOPTIUM="https://api.adoptium.net/v3/binary/latest/8/ga" + declare -A PLATFORMS=( + ["linux-x64"]="linux/x64" + ["linux-aarch64"]="linux/aarch64" + ["macos-x64"]="mac/x64" + ["windows-x64"]="windows/x64" + ) + for platform in "${!PLATFORMS[@]}"; do + os_arch="${PLATFORMS[$platform]}" + echo "=== Downloading JRE 8 for $platform ($os_arch) ===" + if [[ "$platform" == windows-* ]]; then + curl -L -f -o jre8-$platform.zip "$ADOPTIUM/$os_arch/jre/hotspot/normal/eclipse?project=jdk" + mkdir -p jre8-extract + unzip -q jre8-$platform.zip -d jre8-extract + mv jre8-extract/*/ dbx-jre || mv jre8-extract/* dbx-jre + rm -rf jre8-extract jre8-$platform.zip + else + curl -L -f -o jre8-$platform.tar.gz "$ADOPTIUM/$os_arch/jre/hotspot/normal/eclipse?project=jdk" + mkdir -p dbx-jre + tar xzf jre8-$platform.tar.gz -C dbx-jre --strip-components=1 + rm -f jre8-$platform.tar.gz + fi + tar czf dbx-jre-8-$platform.tar.gz dbx-jre + rm -rf dbx-jre + done + + if [ -f dbx-jre-8-macos-x64.tar.gz ]; then + cp dbx-jre-8-macos-x64.tar.gz dbx-jre-8-macos-aarch64.tar.gz + fi + + ls -lh dbx-jre-8-*.tar.gz + - uses: actions/upload-artifact@v4 + with: + name: jre-8 + path: "dbx-jre-8-*.tar.gz" + + release: + needs: [bump-versions, build-agents, build-oracle-native, build-xugu-native, build-jre, build-jre8] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/download-artifact@v4 + with: + path: artifacts + + - name: Flatten artifacts + run: | + mkdir -p release + find artifacts/agent-jars -name '*.jar' -exec cp {} release/ \; + find artifacts/oracle-native -type f -name 'dbx-agent-oracle-*' -exec cp {} release/ \; + find artifacts/xugu-native -type f -name 'dbx-agent-xugu-*' -exec cp {} release/ \; + find artifacts -name 'dbx-jre-*.tar.gz' -exec cp {} release/ \; + ls -lh release/ + + - name: Generate agent-registry.json + env: + MODULE_VERSIONS: ${{ needs.bump-versions.outputs.versions }} + run: | + TAG="${GITHUB_REF_NAME}" + REPO="${GITHUB_REPOSITORY}" + RELEASE_VERSION="${TAG#agents-v}" + + get_module_version() { + local name="$1" + echo "$MODULE_VERSIONS" | python3 -c "import sys,json; print(json.load(sys.stdin).get('$name','${RELEASE_VERSION}'))" + } + + generate_jar_entry() { + local name="$1" label="$2" file="$3" jre_key="$4" version="$5" external_driver="$6" native_json="$7" + local sha256=$(sha256sum "$file" | cut -d' ' -f1) + local size=$(stat -c%s "$file") + local url="https://github.com/${REPO}/releases/download/${TAG}/$(basename $file)" + local native_line="" + if [ -n "$native_json" ]; then + native_line=" \"native\": {${native_json} + }," + fi + cat < release/agent-registry.json < /dev/null + echo "=== agent-registry.json ===" + cat release/agent-registry.json + + - name: Build offline ZIP bundles + run: bash agents/scripts/build_offline_zip.sh release + + - name: Generate release notes + env: + MODULE_VERSIONS: ${{ needs.bump-versions.outputs.versions }} + PREV_VERSIONS: ${{ needs.bump-versions.outputs.prev_versions }} + PREV_TAG: ${{ needs.bump-versions.outputs.prev_tag }} + run: | + NOTES="" + DRIVER_NAMES=() + for f in release/dbx-agent-*.jar; do + DRIVER_NAMES+=("$(basename "$f" .jar | sed 's/dbx-agent-//')") + done + for name in xugu; do + [ -f "release/dbx-agent-${name}.jar" ] && continue + compgen -G "release/dbx-agent-${name}-*" > /dev/null && DRIVER_NAMES+=("$name") + done + + native_only_label() { + local name="$1" + case "$name" in + xugu) echo "虚谷 XuguDB" ;; + *) echo "$name" ;; + esac + } + + for name in "${DRIVER_NAMES[@]}"; do + old_ver=$(echo "$PREV_VERSIONS" | python3 -c "import sys,json; print(json.load(sys.stdin).get('$name',''))") + new_ver=$(echo "$MODULE_VERSIONS" | python3 -c "import sys,json; print(json.load(sys.stdin).get('$name',''))") + [ "$old_ver" = "$new_ver" ] && continue + + jar_file="release/dbx-agent-${name}.jar" + if [ -f "$jar_file" ]; then + label=$(unzip -p "$jar_file" META-INF/MANIFEST.MF | awk -F': ' 'BEGIN{IGNORECASE=1} /^Agent-Label:/ {sub(/\r$/, "", $2); print $2; exit}' || echo "$name") + [ -z "$label" ] && label="$name" + else + label=$(native_only_label "$name") + fi + + NOTES="${NOTES}### ${label} (${old_ver} → ${new_ver})"$'\n' + if [ -d "agents/drivers/$name" ]; then + LOG_PATH="agents/drivers/$name/" + else + LOG_PATH="agents/$name/" + fi + LOG_PATHS=("$LOG_PATH") + if [ "$name" = "oracle" ]; then + LOG_PATHS+=("agents/drivers/oracle-go/") + fi + while IFS= read -r line; do + [ -n "$line" ] && NOTES="${NOTES}- ${line}"$'\n' + done < <(git log --oneline "$PREV_TAG"..HEAD -- "${LOG_PATHS[@]}" | sed 's/^[0-9a-f]* //') + + # include common changes if any + while IFS= read -r line; do + [ -n "$line" ] && NOTES="${NOTES}- ${line}"$'\n' + done < <(git log --oneline "$PREV_TAG"..HEAD -- agents/common/ | sed 's/^[0-9a-f]* //') + NOTES="${NOTES}"$'\n' + done + + if [ -z "$NOTES" ]; then + NOTES="Routine rebuild, no agent-specific changes."$'\n' + fi + + printf "## 更新内容\n\n%s" "$NOTES" > RELEASE_NOTES.md + echo "=== Release Notes ===" + cat RELEASE_NOTES.md + + - name: Create GitHub Release + uses: softprops/action-gh-release@v2 + with: + files: release/* + body_path: RELEASE_NOTES.md + make_latest: false + + - name: Update agents-latest GitHub release + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + git tag -f agents-latest "$GITHUB_SHA" + git push origin refs/tags/agents-latest --force + + if gh release view agents-latest --repo "$GITHUB_REPOSITORY" >/dev/null 2>&1; then + gh release upload agents-latest release/* --repo "$GITHUB_REPOSITORY" --clobber + gh release edit agents-latest \ + --repo "$GITHUB_REPOSITORY" \ + --title "Agents latest" \ + --notes "Moving release for the latest DBX agent registry, driver artifacts, managed JRE archives, and offline bundles." \ + --latest=false + else + gh release create agents-latest release/* \ + --repo "$GITHUB_REPOSITORY" \ + --title "Agents latest" \ + --notes "Moving release for the latest DBX agent registry, driver artifacts, managed JRE archives, and offline bundles." \ + --latest=false + fi + + - name: Upload to R2 + env: + MODULE_VERSIONS: ${{ needs.bump-versions.outputs.versions }} + PREV_VERSIONS: ${{ needs.bump-versions.outputs.prev_versions }} + AWS_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} + R2_ENDPOINT: "https://${{ secrets.R2_ACCOUNT_ID }}.r2.cloudflarestorage.com" + R2_BUCKET: "s3://${{ secrets.R2_BUCKET_NAME }}" + run: | + # Always upload agent-registry.json + aws s3 cp release/agent-registry.json \ + "$R2_BUCKET/agents/agent-registry.json" \ + --endpoint-url "$R2_ENDPOINT" + echo "Uploaded agent-registry.json" + + # Upload only changed driver JARs + UPLOADED=0 + SKIPPED=0 + for f in release/dbx-agent-*.jar; do + [ ! -f "$f" ] && continue + name=$(basename "$f" .jar | sed 's/dbx-agent-//') + old_ver=$(echo "$PREV_VERSIONS" | python3 -c "import sys,json; print(json.load(sys.stdin).get('$name',''))" 2>/dev/null) + new_ver=$(echo "$MODULE_VERSIONS" | python3 -c "import sys,json; print(json.load(sys.stdin).get('$name',''))" 2>/dev/null) + if [ "$old_ver" = "$new_ver" ] && [ -n "$old_ver" ]; then + echo "Skip $(basename $f) (unchanged $old_ver)" + SKIPPED=$((SKIPPED + 1)) + continue + fi + aws s3 cp "$f" \ + "$R2_BUCKET/agents/drivers/$(basename $f)" \ + --endpoint-url "$R2_ENDPOINT" + echo "Uploaded $(basename $f) ($old_ver → $new_ver)" + UPLOADED=$((UPLOADED + 1)) + done + echo "Drivers: $UPLOADED uploaded, $SKIPPED skipped" + + for f in release/dbx-agent-*; do + [ -f "$f" ] || continue + [[ "$f" != *.jar ]] || continue + aws s3 cp "$f" \ + "$R2_BUCKET/agents/drivers/$(basename $f)" \ + --endpoint-url "$R2_ENDPOINT" + echo "Uploaded $(basename $f)" + done + + # Upload JRE — always overwrite to ensure latest modules + for f in release/dbx-jre-*.tar.gz; do + [ ! -f "$f" ] && continue + key="agents/jre/$(basename $f)" + aws s3 cp "$f" "$R2_BUCKET/$key" --endpoint-url "$R2_ENDPOINT" + echo "Uploaded $(basename $f)" + done diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 91a05a88ae..8d288fbb98 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -103,3 +103,86 @@ jobs: - name: JDBC plugin package check run: ./plugins/jdbc/package.sh + + changes: + runs-on: ubuntu-22.04 + outputs: + agents: ${{ steps.filter.outputs.agents }} + steps: + - uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Detect agent changes + id: filter + shell: bash + env: + BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.event.before }} + run: | + if [ -z "$BASE_SHA" ] || echo "$BASE_SHA" | grep -Eq '^0+$' || ! git rev-parse --verify "$BASE_SHA^{commit}" >/dev/null 2>&1; then + if git rev-parse --verify HEAD^ >/dev/null 2>&1; then + BASE_SHA="HEAD^" + else + BASE_SHA="$(git hash-object -t tree /dev/null)" + fi + fi + + if git diff --name-only "$BASE_SHA" HEAD -- agents/ | grep -q .; then + echo "agents=true" >> "$GITHUB_OUTPUT" + else + echo "agents=false" >> "$GITHUB_OUTPUT" + fi + + agents: + needs: changes + if: needs.changes.outputs.agents == 'true' + runs-on: ubuntu-22.04 + defaults: + run: + working-directory: agents + steps: + - uses: actions/checkout@v5 + + - name: Setup Java + uses: actions/setup-java@v5 + with: + distribution: temurin + java-version: | + 8 + 21 + + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version: "1.22.x" + + - name: Setup Gradle + uses: gradle/actions/setup-gradle@v4 + + - name: Agent script tests + run: python3 -m unittest discover -s scripts -p '*_test.py' + + - name: Agent validation + run: python3 scripts/validate_agents.py + + - name: Oracle native agent tests + run: go test ./... + working-directory: agents/drivers/oracle-go + + - name: Xugu native agent tests + run: GONOSUMDB=gitee.com/XuguDB/go-xugu-driver go test ./... + working-directory: agents/drivers/xugu + + - 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 + + - name: Xugu native agent build + run: GONOSUMDB=gitee.com/XuguDB/go-xugu-driver CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-xugu-linux-x64 . + working-directory: agents/drivers/xugu + + - name: Java agent tests and packages + run: ./gradlew test shadowJar --continue + + - name: Agent jar validation + run: python3 scripts/validate_agent_jars.py diff --git a/README.md b/README.md index 3a0638261e..20eba2e8a0 100644 --- a/README.md +++ b/README.md @@ -273,16 +273,14 @@ pnpm dev:web # frontend pnpm dev:backend # backend ``` -[dbx-agents](https://github.com/t8y2/dbx-agents) is a separate repository containing JDBC agent driver development projects. For local development, clone it alongside `dbx/` under the same workspace directory and open the parent folder in your IDE: +JDBC agent driver development projects live in `agents/`: ```bash -mkdir dbx-workspace && cd dbx-workspace -git clone https://github.com/t8y2/dbx.git -git clone https://github.com/t8y2/dbx-agents.git -# Open dbx-workspace/ in your IDE to work on both projects together +cd agents +./gradlew test ``` -This keeps the two repositories independent (separate git histories) while making it easy to navigate between them during development. +Build artifacts from `agents/drivers//build/libs/` are picked up by local driver install flows when available. ### Build diff --git a/README.zh-CN.md b/README.zh-CN.md index bebe9f9463..1bcca8f93e 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -272,15 +272,15 @@ pnpm dev:web # 前端 pnpm dev:backend # 后端 ``` -[dbx-agents](https://github.com/t8y2/dbx-agents) 是独立仓库，存放 JDBC Agent 驱动开发工程。本地开发时建议克隆到 `dbx/` 同级目录下，用 IDE 打开父文件夹即可同时开发两个项目： +JDBC Agent 驱动开发工程位于 `agents/` 目录： ```bash -mkdir dbx-workspace && cd dbx-workspace -git clone https://github.com/t8y2/dbx.git -git clone https://github.com/t8y2/dbx-agents.git -# 用 IDE 打开 dbx-workspace/ 即可同时开发两个项目 +cd agents +./gradlew test ``` +本地驱动安装流程会优先查找 `agents/drivers//build/libs/` 下的构建产物。 + ### 构建 ```bash diff --git a/agents/.gitignore b/agents/.gitignore new file mode 100644 index 0000000000..f797c2c638 --- /dev/null +++ b/agents/.gitignore @@ -0,0 +1,7 @@ +*/libs/*.jar +gradle.properties +build/ +*/bin/ +.gradle/ +.DS_Store +openspec diff --git a/agents/README.md b/agents/README.md new file mode 100644 index 0000000000..b10576742f --- /dev/null +++ b/agents/README.md @@ -0,0 +1,83 @@ +# DBX Agents + +Agent drivers for [DBX](https://github.com/t8y2/dbx) — database support via JDBC and native database drivers. + +Each agent runs as a standalone process and communicates with DBX via stdin/stdout JSON-RPC 2.0. + +## Supported Databases + +| Agent | Database | JDBC Driver | +|-------|----------|-------------| +| access | Microsoft Access | UCanAccess | +| dameng | 达梦 DM8 | DM JDBC | +| kingbase | 人大金仓 KingbaseES | KingbaseES JDBC | +| vastbase | Vastbase | Vastbase JDBC | +| goldendb | GoldenDB | MySQL Connector/J | +| databend | Databend | Databend JDBC | +| databricks | Databricks SQL | Databricks JDBC | +| saphana | SAP HANA | SAP HANA JDBC | +| teradata | Teradata | Teradata JDBC | +| vertica | Vertica | Vertica JDBC | +| firebird | Firebird | Jaybird JDBC | +| exasol | Exasol | Exasol JDBC | +| oceanbase-oracle | OceanBase Oracle Mode | OceanBase JDBC | +| gbase8a | GBase 8a | External GBase 8a JDBC | +| gbase8s | GBase 8s | External GBase 8s JDBC | +| oracle | Oracle 10g+ | go-ora native agent | +| oracle-legacy | Oracle 11g/12c/18c/19c | ojdbc8 (compatibility fallback) | +| oracle-10g | Oracle 10g | ojdbc6 (compatibility fallback, JRE 8) | +| h2 | H2 | H2 JDBC | +| snowflake | Snowflake | Snowflake JDBC | +| trino | Trino (Presto) | Trino JDBC | +| hive | Apache Hive | Hive JDBC | +| db2 | IBM DB2 | DB2 JDBC | +| informix | IBM Informix | Informix JDBC | +| neo4j | Neo4j | Neo4j JDBC | +| cassandra | Apache Cassandra | Cassandra JDBC | +| bigquery | Google BigQuery | BigQuery JDBC | +| kylin | Apache Kylin | Kylin JDBC | +| sundb | SunDB | SunDB JDBC | +| gaussdb | GaussDB | GaussDB JDBC | +| tdengine | TDengine | taos-jdbcdriver (WebSocket, REST fallback) | +| yashandb | 崖山 YashanDB | YashanDB JDBC | +| xugu | 虚谷 XuguDB | XuguDB Go native agent | +| iotdb | Apache IoTDB | IoTDB JDBC | +| etcd | etcd | jetcd | + +## Multi-JRE Support + +Most Java agents target JRE 21. Native agents, such as `oracle` and `xugu`, do not require a JRE. Agents that still require legacy Java runtimes (e.g. compatibility fallback `oracle-10g` uses JRE 8) declare their JRE version in the registry. DBX downloads and manages multiple JRE installations automatically. + +## Build + +Requires JDK 8 and 21 (Gradle toolchain auto-downloads if needed). + +```bash +./gradlew shadowJar +(cd drivers/oracle-go && go build -o agent .) +(cd drivers/xugu && go build -o agent .) +``` + +Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/oracle-go` and `drivers/xugu`. + +## Development + +- Agent authoring guide: [docs/agent-authoring.md](docs/agent-authoring.md) +- JDBC agent template: [docs/examples/jdbc-agent-template](docs/examples/jdbc-agent-template) +- Release checklist: [docs/release-checklist.md](docs/release-checklist.md) + +## Architecture + +``` +DBX Main Process (Rust/Tauri) + │ stdin/stdout (JSON-RPC 2.0) + ▼ +agent / java -jar dbx-agent-{type}.jar + │ + ▼ +Native driver / JDBC → Database +``` + +## License + +[AGPL-3.0](https://github.com/t8y2/dbx/blob/main/LICENSE) diff --git a/agents/build.gradle b/agents/build.gradle new file mode 100644 index 0000000000..4891f9e63f --- /dev/null +++ b/agents/build.gradle @@ -0,0 +1,54 @@ +plugins { + id 'com.gradleup.shadow' version '9.0.0-beta12' apply false +} + +def java8Projects = ['common', 'test-support', 'oracle-10g'] as Set +def infrastructureProjects = ['common', 'test-support'] as Set +def legacyStandaloneProjects = ['mongodb'] as Set +def agentProjects = subprojects.findAll { !infrastructureProjects.contains(it.name) } +def jdbcAgentProjects = agentProjects.findAll { !legacyStandaloneProjects.contains(it.name) } + +allprojects { + repositories { + mavenCentral() + } +} + +subprojects { + if (name == 'test-support') { + apply plugin: 'java-library' + } else { + apply plugin: 'java' + } + + java { + toolchain { + languageVersion = JavaLanguageVersion.of(java8Projects.contains(name) ? 8 : 21) + } + } + + dependencies { + testImplementation 'org.junit.jupiter:junit-jupiter:5.11.4' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher:1.11.4' + } + + tasks.withType(Test).configureEach { + useJUnitPlatform() + } +} + +configure(agentProjects) { + apply plugin: 'com.gradleup.shadow' + + tasks.named('shadowJar') { + archiveBaseName = "dbx-agent-${project.name}" + archiveClassifier = '' + } +} + +configure(jdbcAgentProjects) { + dependencies { + implementation project(':common') + testImplementation project(':test-support') + } +} diff --git a/agents/common/build.gradle b/agents/common/build.gradle new file mode 100644 index 0000000000..ee0d25b0ff --- /dev/null +++ b/agents/common/build.gradle @@ -0,0 +1,3 @@ +dependencies { + implementation 'com.google.code.gson:gson:2.12.1' +} diff --git a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java new file mode 100644 index 0000000000..8cd794e497 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java @@ -0,0 +1,205 @@ +package com.dbx.agent; + +import java.net.URL; +import java.net.URLClassLoader; +import java.nio.file.Paths; +import java.sql.Connection; +import java.sql.Driver; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Types; +import java.util.ArrayList; +import java.util.List; + +public abstract class AbstractJdbcAgent extends BaseDatabaseAgent { + private Connection connection; + private String configuredDatabase = ""; + + @Override + public final Connection getConnection() { + return connection; + } + + @Override + public final void connect(ConnectParams params) { + uncheckedVoid(() -> { + loadDriver(params); + connection = openConnection(params); + configuredDatabase = params.getDatabase(); + afterConnect(params, connection); + }); + } + + @Override + public final boolean testConnection(ConnectParams params) { + return unchecked(() -> { + loadDriver(params); + try (Connection conn = openConnection(params)) { + return conn.isValid(5); + } + }); + } + + private void loadDriver(ConnectParams params) throws Exception { + List driverPaths = params.getJdbc_driver_paths(); + String driverClass = params.getJdbc_driver_class(); + if (driverClass == null || driverClass.isEmpty()) { + driverClass = driverClass(); + } + if (driverPaths != null && !driverPaths.isEmpty()) { + List urls = new ArrayList<>(); + for (String path : driverPaths) { + urls.add(Paths.get(path).toUri().toURL()); + } + URLClassLoader loader = new URLClassLoader(urls.toArray(new URL[0]), getClass().getClassLoader()); + Driver driver = (Driver) Class.forName(driverClass, true, loader).getDeclaredConstructor().newInstance(); + DriverManager.registerDriver(new DriverShim(driver)); + } else { + Class.forName(driverClass); + } + } + + private static final class DriverShim implements Driver { + private final Driver driver; + + DriverShim(Driver driver) { + this.driver = driver; + } + + @Override + public Connection connect(String url, java.util.Properties info) throws java.sql.SQLException { + return driver.connect(url, info); + } + + @Override + public boolean acceptsURL(String url) throws java.sql.SQLException { + return driver.acceptsURL(url); + } + + @Override + public java.sql.DriverPropertyInfo[] getPropertyInfo(String url, java.util.Properties info) throws java.sql.SQLException { + return driver.getPropertyInfo(url, info); + } + + @Override + public int getMajorVersion() { + return driver.getMajorVersion(); + } + + @Override + public int getMinorVersion() { + return driver.getMinorVersion(); + } + + @Override + public boolean jdbcCompliant() { + return driver.jdbcCompliant(); + } + + @Override + public java.util.logging.Logger getParentLogger() throws java.sql.SQLFeatureNotSupportedException { + return driver.getParentLogger(); + } + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::resultValue + ); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + return JdbcExecutor.INSTANCE.executePage( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options, + this::resultValue + ); + } + + @Override + public QueryPageResult fetchQueryPage(String sessionId, int pageSize) { + return JdbcExecutor.INSTANCE.fetchPage(sessionId, pageSize); + } + + @Override + public boolean closeQuerySession(String sessionId) { + return JdbcExecutor.INSTANCE.closeQuerySession(sessionId); + } + + @Override + public QueryResult executeTransaction(List statements, String schema) { + return TransactionExecutor.executeUpdateStatements(requireConnected(), statements, schema, this::setSchemaSQL); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + protected abstract String driverClass(); + + protected abstract String buildJdbcUrl(ConnectParams params); + + protected Connection openConnection(ConnectParams params) throws Exception { + return DriverManager.getConnection(buildJdbcUrl(params), params.getUsername(), params.getPassword()); + } + + protected void afterConnect(ConnectParams params, Connection connection) throws Exception { + } + + protected String getConfiguredDatabase() { + return configuredDatabase; + } + + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value; + switch (sqlType) { + case Types.BIGINT: + value = rs.getLong(index); + break; + case Types.INTEGER: + case Types.SMALLINT: + case Types.TINYINT: + value = rs.getInt(index); + break; + case Types.FLOAT: + case Types.REAL: + value = rs.getFloat(index); + break; + case Types.DOUBLE: + value = rs.getDouble(index); + break; + case Types.DECIMAL: + case Types.NUMERIC: + value = rs.getBigDecimal(index); + break; + case Types.BOOLEAN: + case Types.BIT: + value = rs.getBoolean(index); + break; + default: + value = rs.getString(index); + break; + } + return rs.wasNull() ? null : value; + }); + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java new file mode 100644 index 0000000000..29513d54ba --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -0,0 +1,134 @@ +package com.dbx.agent; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +public final class AgentProtocol { + public static final int PROTOCOL_VERSION = 1; + + public static final String METHOD_HANDSHAKE = "handshake"; + public static final String METHOD_CONNECT = "connect"; + public static final String METHOD_TEST_CONNECTION = "test_connection"; + public static final String METHOD_VALIDATE_CONNECTION = "validate_connection"; + public static final String METHOD_LIST_DATABASES = "list_databases"; + public static final String METHOD_LIST_SCHEMAS = "list_schemas"; + public static final String METHOD_LIST_TABLES = "list_tables"; + public static final String METHOD_LIST_OBJECTS = "list_objects"; + public static final String METHOD_GET_OBJECT_SOURCE = "get_object_source"; + public static final String METHOD_GET_TABLE_DDL = "get_table_ddl"; + public static final String METHOD_GET_COLUMNS = "get_columns"; + public static final String METHOD_LIST_INDEXES = "list_indexes"; + public static final String METHOD_LIST_FOREIGN_KEYS = "list_foreign_keys"; + public static final String METHOD_LIST_TRIGGERS = "list_triggers"; + public static final String METHOD_EXECUTE_QUERY = "execute_query"; + public static final String METHOD_EXECUTE_QUERY_PAGE = "execute_query_page"; + public static final String METHOD_FETCH_QUERY_PAGE = "fetch_query_page"; + public static final String METHOD_CLOSE_QUERY_SESSION = "close_query_session"; + public static final String METHOD_GET_EXPLAIN_INFO = "get_explain_info"; + public static final String METHOD_EXECUTE_TRANSACTION = "execute_transaction"; + public static final String METHOD_DISCONNECT = "disconnect"; + public static final String METHOD_SHUTDOWN = "shutdown"; + + public static final String MONGO_METHOD_LIST_DATABASES = "list_databases"; + public static final String MONGO_METHOD_LIST_COLLECTIONS = "list_collections"; + public static final String MONGO_METHOD_FIND_DOCUMENTS = "find_documents"; + public static final String MONGO_METHOD_INSERT_DOCUMENT = "insert_document"; + public static final String MONGO_METHOD_UPDATE_DOCUMENT = "update_document"; + public static final String MONGO_METHOD_DELETE_DOCUMENT = "delete_document"; + + public static final String KV_METHOD_LIST_PREFIX = "kv_list_prefix"; + public static final String KV_METHOD_GET = "kv_get"; + public static final String KV_METHOD_PUT = "kv_put"; + public static final String KV_METHOD_DELETE = "kv_delete"; + + public static final String CAPABILITY_CONNECT = "connect"; + public static final String CAPABILITY_TEST_CONNECTION = "test_connection"; + public static final String CAPABILITY_METADATA = "metadata"; + public static final String CAPABILITY_QUERY = "query"; + public static final String CAPABILITY_PAGED_QUERY = "paged_query"; + public static final String CAPABILITY_TRANSACTION = "transaction"; + public static final String CAPABILITY_DDL = "ddl"; + public static final String CAPABILITY_KV = "kv"; + + public static final List CAPABILITIES = Collections.unmodifiableList(Arrays.asList( + CAPABILITY_CONNECT, + CAPABILITY_TEST_CONNECTION, + CAPABILITY_METADATA, + CAPABILITY_QUERY, + CAPABILITY_PAGED_QUERY, + CAPABILITY_TRANSACTION, + CAPABILITY_DDL + )); + + public static final List ALL_CAPABILITIES = Collections.unmodifiableList(Arrays.asList( + CAPABILITY_CONNECT, + CAPABILITY_TEST_CONNECTION, + CAPABILITY_METADATA, + CAPABILITY_QUERY, + CAPABILITY_PAGED_QUERY, + CAPABILITY_TRANSACTION, + CAPABILITY_DDL, + CAPABILITY_KV + )); + + public static final List COMMON_METHODS = Collections.unmodifiableList(Arrays.asList( + METHOD_HANDSHAKE, + METHOD_CONNECT, + METHOD_TEST_CONNECTION, + METHOD_VALIDATE_CONNECTION, + METHOD_LIST_DATABASES, + METHOD_LIST_SCHEMAS, + METHOD_LIST_TABLES, + METHOD_LIST_OBJECTS, + METHOD_GET_OBJECT_SOURCE, + METHOD_GET_TABLE_DDL, + METHOD_GET_COLUMNS, + METHOD_LIST_INDEXES, + METHOD_LIST_FOREIGN_KEYS, + METHOD_LIST_TRIGGERS, + METHOD_EXECUTE_QUERY, + METHOD_EXECUTE_QUERY_PAGE, + METHOD_FETCH_QUERY_PAGE, + METHOD_CLOSE_QUERY_SESSION, + METHOD_GET_EXPLAIN_INFO, + METHOD_EXECUTE_TRANSACTION, + METHOD_DISCONNECT, + METHOD_SHUTDOWN + )); + + public static final List MONGO_LEGACY_METHODS = Collections.unmodifiableList(Arrays.asList( + MONGO_METHOD_LIST_DATABASES, + MONGO_METHOD_LIST_COLLECTIONS, + MONGO_METHOD_FIND_DOCUMENTS, + MONGO_METHOD_INSERT_DOCUMENT, + MONGO_METHOD_UPDATE_DOCUMENT, + MONGO_METHOD_DELETE_DOCUMENT + )); + + public static final List KV_METHODS = Collections.unmodifiableList(Arrays.asList( + KV_METHOD_LIST_PREFIX, + KV_METHOD_GET, + KV_METHOD_PUT, + KV_METHOD_DELETE + )); + + private AgentProtocol() { + } + + public static HandshakeResult handshakeResult() { + return new HandshakeResult(PROTOCOL_VERSION, PROTOCOL_VERSION, CAPABILITIES); + } + + public static final class HandshakeResult { + private final int protocolVersion; + private final int agentProtocolVersion; + private final List capabilities; + + private HandshakeResult(int protocolVersion, int agentProtocolVersion, List capabilities) { + this.protocolVersion = protocolVersion; + this.agentProtocolVersion = agentProtocolVersion; + this.capabilities = capabilities; + } + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java new file mode 100644 index 0000000000..d904fe779c --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java @@ -0,0 +1,102 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +public abstract class BaseDatabaseAgent implements DatabaseAgent { + @Override + public List listObjects(String schema) { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + return result; + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + throw new UnsupportedOperationException("Object source is not supported"); + } + + @Override + public String getTableDdl(String schema, String table) { + List indexes; + try { + indexes = listIndexes(schema, table); + } catch (RuntimeException e) { + indexes = Collections.emptyList(); + } + + List foreignKeys; + try { + foreignKeys = listForeignKeys(schema, table); + } catch (RuntimeException e) { + foreignKeys = Collections.emptyList(); + } + + return DatabaseAgent.buildTableDdl(schema, table, getColumns(schema, table), indexes, foreignKeys); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + Connection conn = requireConnected(); + return JdbcExecutor.INSTANCE.executePage( + conn, + sql, + schema, + this::setSchemaSQL, + options, + JdbcExecutor.INSTANCE::defaultResultValue + ); + } + + @Override + public QueryPageResult fetchQueryPage(String sessionId, int pageSize) { + return JdbcExecutor.INSTANCE.fetchPage(sessionId, pageSize); + } + + @Override + public boolean closeQuerySession(String sessionId) { + return JdbcExecutor.INSTANCE.closeQuerySession(sessionId); + } + + @Override + public QueryResult executeTransaction(List statements, String schema) { + return TransactionExecutor.executeUpdateStatements(requireConnected(), statements, schema, this::setSchemaSQL); + } + + protected Connection requireConnected() { + Connection conn = getConnection(); + if (conn == null) { + throw new IllegalStateException("Not connected"); + } + return conn; + } + + protected static T unchecked(ThrowingSupplier supplier) { + try { + return supplier.get(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + protected static void uncheckedVoid(ThrowingRunnable runnable) { + unchecked(() -> { + runnable.run(); + return null; + }); + } + + protected interface ThrowingSupplier { + T get() throws Exception; + } + + protected interface ThrowingRunnable { + void run() throws Exception; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/ColumnInfo.java b/agents/common/src/main/java/com/dbx/agent/ColumnInfo.java new file mode 100644 index 0000000000..d65df3b5e0 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ColumnInfo.java @@ -0,0 +1,176 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class ColumnInfo { + private String name; + private String data_type; + private boolean is_nullable; + private String column_default; + private boolean is_primary_key; + private String extra; + private String comment; + private Integer numeric_precision; + private Integer numeric_scale; + private Integer character_maximum_length; + + public ColumnInfo() { + this("", "", true, null, false); + } + + public ColumnInfo(String name, String data_type, boolean is_nullable, String column_default, boolean is_primary_key) { + this(name, data_type, is_nullable, column_default, is_primary_key, null, null, null, null, null); + } + + public ColumnInfo( + String name, + String data_type, + boolean is_nullable, + String column_default, + boolean is_primary_key, + String extra, + String comment, + Integer numeric_precision, + Integer numeric_scale, + Integer character_maximum_length + ) { + this.name = name; + this.data_type = data_type; + this.is_nullable = is_nullable; + this.column_default = column_default; + this.is_primary_key = is_primary_key; + this.extra = extra; + this.comment = comment; + this.numeric_precision = numeric_precision; + this.numeric_scale = numeric_scale; + this.character_maximum_length = character_maximum_length; + } + + public String getName() { + return name; + } + + public String getData_type() { + return data_type; + } + + public boolean getIs_nullable() { + return is_nullable; + } + + public String getColumn_default() { + return column_default; + } + + public boolean getIs_primary_key() { + return is_primary_key; + } + + public String getExtra() { + return extra; + } + + public String getComment() { + return comment; + } + + public Integer getNumeric_precision() { + return numeric_precision; + } + + public Integer getNumeric_scale() { + return numeric_scale; + } + + public Integer getCharacter_maximum_length() { + return character_maximum_length; + } + + public void setName(String name) { + this.name = name; + } + + public void setData_type(String data_type) { + this.data_type = data_type; + } + + public void setIs_nullable(boolean is_nullable) { + this.is_nullable = is_nullable; + } + + public void setColumn_default(String column_default) { + this.column_default = column_default; + } + + public void setIs_primary_key(boolean is_primary_key) { + this.is_primary_key = is_primary_key; + } + + public void setExtra(String extra) { + this.extra = extra; + } + + public void setComment(String comment) { + this.comment = comment; + } + + public void setNumeric_precision(Integer numeric_precision) { + this.numeric_precision = numeric_precision; + } + + public void setNumeric_scale(Integer numeric_scale) { + this.numeric_scale = numeric_scale; + } + + public void setCharacter_maximum_length(Integer character_maximum_length) { + this.character_maximum_length = character_maximum_length; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ColumnInfo)) return false; + ColumnInfo that = (ColumnInfo) other; + return is_nullable == that.is_nullable + && is_primary_key == that.is_primary_key + && Objects.equals(name, that.name) + && Objects.equals(data_type, that.data_type) + && Objects.equals(column_default, that.column_default) + && Objects.equals(extra, that.extra) + && Objects.equals(comment, that.comment) + && Objects.equals(numeric_precision, that.numeric_precision) + && Objects.equals(numeric_scale, that.numeric_scale) + && Objects.equals(character_maximum_length, that.character_maximum_length); + } + + @Override + public int hashCode() { + return Objects.hash( + name, + data_type, + is_nullable, + column_default, + is_primary_key, + extra, + comment, + numeric_precision, + numeric_scale, + character_maximum_length + ); + } + + @Override + public String toString() { + return "ColumnInfo(name=" + name + + ", data_type=" + data_type + + ", is_nullable=" + is_nullable + + ", column_default=" + column_default + + ", is_primary_key=" + is_primary_key + + ", extra=" + extra + + ", comment=" + comment + + ", numeric_precision=" + numeric_precision + + ", numeric_scale=" + numeric_scale + + ", character_maximum_length=" + character_maximum_length + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java new file mode 100644 index 0000000000..c779c3b05d --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java @@ -0,0 +1,179 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public abstract class ConfiguredJdbcAgent extends AbstractJdbcAgent { + private final JdbcAgentProfile profile; + private String configuredDatabase = ""; + + protected ConfiguredJdbcAgent(JdbcAgentProfile profile) { + this.profile = profile; + } + + public JdbcAgentProfile getProfile() { + return profile; + } + + @Override + protected String driverClass() { + return profile.getDriverClass(); + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return profile.buildUrl(params); + } + + @Override + protected void afterConnect(ConnectParams params, Connection connection) { + configuredDatabase = params.getDatabase(); + } + + @Override + public List listDatabases() { + return StandardJdbcMetadata.INSTANCE.listDatabases(requireConnection(), configuredDatabase); + } + + @Override + public List listSchemas() { + return StandardJdbcMetadata.INSTANCE.listSchemas(requireConnection(), profile); + } + + @Override + public List listTables(String schema) { + return StandardJdbcMetadata.INSTANCE.listTables(requireConnection(), profile, configuredDatabase, schema); + } + + @Override + public List listObjects(String schema) { + return StandardJdbcMetadata.INSTANCE.listObjects(listTables(schema), schema); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + throw new UnsupportedOperationException("Object source is not supported"); + } + + @Override + public List getColumns(String schema, String table) { + return StandardJdbcMetadata.INSTANCE.getColumns(requireConnection(), profile, configuredDatabase, schema, table); + } + + @Override + public List listIndexes(String schema, String table) { + return StandardJdbcMetadata.INSTANCE.listIndexes(requireConnection(), schema, table); + } + + @Override + public List listForeignKeys(String schema, String table) { + return StandardJdbcMetadata.INSTANCE.listForeignKeys(requireConnection(), schema, table); + } + + @Override + public List listTriggers(String schema, String table) { + return StandardJdbcMetadata.INSTANCE.listTriggers(schema, table); + } + + @Override + public String getTableDdl(String schema, String table) { + List indexes; + try { + indexes = listIndexes(schema, table); + } catch (RuntimeException e) { + indexes = Collections.emptyList(); + } + + List foreignKeys; + try { + foreignKeys = listForeignKeys(schema, table); + } catch (RuntimeException e) { + foreignKeys = Collections.emptyList(); + } + + return DatabaseAgent.buildTableDdl(schema, table, getColumns(schema, table), indexes, foreignKeys); + } + + @Override + public String setSchemaSQL(String schema) { + if (profile.getSkipExecutionContext()) { + return ""; + } + return StandardJdbcMetadata.schemaSwitchSql(requireConnection(), profile, schema); + } + + protected Connection requireConnection() { + return requireConnected(); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return super.resultValue(rs, index, sqlType); + } + + private static String normalizeTableType(String type) { + if (type == null || type.trim().isEmpty()) { + return "TABLE"; + } + if ("BASE TABLE".equals(type.toUpperCase(Locale.ROOT))) { + return "TABLE"; + } + return type; + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static Integer characterLength(ResultSet rs) throws Exception { + String typeName = rs.getString("TYPE_NAME"); + if (typeName == null) { + return null; + } + String normalized = typeName.toLowerCase(Locale.ROOT); + if (!normalized.contains("char") && !normalized.contains("text")) { + return null; + } + return intOrNull(rs, "COLUMN_SIZE"); + } + + private static String blankToNull(String value) { + return value == null || value.trim().isEmpty() ? null : value; + } + + private static void addNonBlank(Set values, String value) { + if (value != null && !value.trim().isEmpty()) { + values.add(value); + } + } + + private static final class IndexColumn { + private final int ordinal; + private final String name; + + private IndexColumn(int ordinal, String name) { + this.ordinal = ordinal; + this.name = name; + } + + private int getOrdinal() { + return ordinal; + } + + private String getName() { + return name; + } + } + +} diff --git a/agents/common/src/main/java/com/dbx/agent/ConnectParams.java b/agents/common/src/main/java/com/dbx/agent/ConnectParams.java new file mode 100644 index 0000000000..3e4639a426 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ConnectParams.java @@ -0,0 +1,243 @@ +package com.dbx.agent; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +public final class ConnectParams { + private String host; + private int port; + private String database; + private String username; + private String password; + private String url_params; + private String connection_string; + private boolean mysql_compat_mode; + private String jdbc_driver_class; + private List jdbc_driver_paths; + private boolean ssl; + private String ca_cert_path; + private String client_cert_path; + private String client_key_path; + private String gbase_server; + private String informix_server; + + public ConnectParams() { + this("", 0, "", "", "", "", "", false, "", Collections.emptyList()); + } + + public ConnectParams( + String host, + int port, + String database, + String username, + String password, + String url_params, + String connection_string, + boolean mysql_compat_mode + ) { + this(host, port, database, username, password, url_params, connection_string, mysql_compat_mode, "", Collections.emptyList()); + } + + public ConnectParams( + String host, + int port, + String database, + String username, + String password, + String url_params, + String connection_string, + boolean mysql_compat_mode, + String jdbc_driver_class, + List jdbc_driver_paths + ) { + this.host = host; + this.port = port; + this.database = database; + this.username = username; + this.password = password; + this.url_params = url_params; + this.connection_string = connection_string; + this.mysql_compat_mode = mysql_compat_mode; + this.jdbc_driver_class = jdbc_driver_class; + this.jdbc_driver_paths = jdbc_driver_paths; + } + + public String getHost() { + return host; + } + + public int getPort() { + return port; + } + + public String getDatabase() { + return database; + } + + public String getUsername() { + return username; + } + + public String getPassword() { + return password; + } + + public String getUrl_params() { + return url_params; + } + + public String getConnection_string() { + return connection_string; + } + + public void setHost(String host) { + this.host = host; + } + + public void setPort(int port) { + this.port = port; + } + + public void setDatabase(String database) { + this.database = database; + } + + public void setUsername(String username) { + this.username = username; + } + + public void setPassword(String password) { + this.password = password; + } + + public void setUrl_params(String url_params) { + this.url_params = url_params; + } + + public void setConnection_string(String connection_string) { + this.connection_string = connection_string; + } + + public boolean isMysql_compat_mode() { + return mysql_compat_mode; + } + + public void setMysql_compat_mode(boolean mysql_compat_mode) { + this.mysql_compat_mode = mysql_compat_mode; + } + + public String getJdbc_driver_class() { + return jdbc_driver_class; + } + + public void setJdbc_driver_class(String jdbc_driver_class) { + this.jdbc_driver_class = jdbc_driver_class; + } + + public List getJdbc_driver_paths() { + return jdbc_driver_paths; + } + + public void setJdbc_driver_paths(List jdbc_driver_paths) { + this.jdbc_driver_paths = jdbc_driver_paths; + } + + public boolean isSsl() { + return ssl; + } + + public void setSsl(boolean ssl) { + this.ssl = ssl; + } + + public String getCa_cert_path() { + return ca_cert_path; + } + + public void setCa_cert_path(String ca_cert_path) { + this.ca_cert_path = ca_cert_path; + } + + public String getClient_cert_path() { + return client_cert_path; + } + + public void setClient_cert_path(String client_cert_path) { + this.client_cert_path = client_cert_path; + } + + public String getClient_key_path() { + return client_key_path; + } + + public void setClient_key_path(String client_key_path) { + this.client_key_path = client_key_path; + } + + public String getGbase_server() { + return gbase_server; + } + + public void setGbase_server(String gbase_server) { + this.gbase_server = gbase_server; + } + + public String getInformix_server() { + return informix_server; + } + + public void setInformix_server(String informix_server) { + this.informix_server = informix_server; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ConnectParams)) return false; + ConnectParams that = (ConnectParams) other; + return port == that.port + && Objects.equals(host, that.host) + && Objects.equals(database, that.database) + && Objects.equals(username, that.username) + && Objects.equals(password, that.password) + && Objects.equals(url_params, that.url_params) + && Objects.equals(connection_string, that.connection_string) + && mysql_compat_mode == that.mysql_compat_mode + && Objects.equals(jdbc_driver_class, that.jdbc_driver_class) + && Objects.equals(jdbc_driver_paths, that.jdbc_driver_paths) + && ssl == that.ssl + && Objects.equals(ca_cert_path, that.ca_cert_path) + && Objects.equals(client_cert_path, that.client_cert_path) + && Objects.equals(client_key_path, that.client_key_path) + && Objects.equals(gbase_server, that.gbase_server) + && Objects.equals(informix_server, that.informix_server); + } + + @Override + public int hashCode() { + return Objects.hash(host, port, database, username, password, url_params, connection_string, + mysql_compat_mode, jdbc_driver_class, jdbc_driver_paths, ssl, ca_cert_path, client_cert_path, client_key_path, gbase_server, informix_server); + } + + @Override + public String toString() { + return "ConnectParams(host=" + host + + ", port=" + port + + ", database=" + database + + ", username=" + username + + ", password=" + password + + ", url_params=" + url_params + + ", connection_string=" + connection_string + + ", mysql_compat_mode=" + mysql_compat_mode + + ", jdbc_driver_class=" + jdbc_driver_class + + ", jdbc_driver_paths=" + jdbc_driver_paths + + ", ssl=" + ssl + + ", ca_cert_path=" + ca_cert_path + + ", client_cert_path=" + client_cert_path + + ", client_key_path=" + client_key_path + + ", gbase_server=" + gbase_server + + ", informix_server=" + informix_server + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java new file mode 100644 index 0000000000..7c4958203f --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java @@ -0,0 +1,147 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +public interface DatabaseAgent { + void connect(ConnectParams params); + + boolean testConnection(ConnectParams params); + + List listDatabases(); + + List listSchemas(); + + List listTables(String schema); + + default List listObjects(String schema) { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + return result; + } + + List getColumns(String schema, String table); + + default ObjectSource getObjectSource(String schema, String name, String objectType) { + throw new UnsupportedOperationException("Object source is not supported"); + } + + default String getTableDdl(String schema, String table) { + List indexes; + try { + indexes = listIndexes(schema, table); + } catch (RuntimeException e) { + indexes = Collections.emptyList(); + } + + List foreignKeys; + try { + foreignKeys = listForeignKeys(schema, table); + } catch (RuntimeException e) { + foreignKeys = Collections.emptyList(); + } + + return buildTableDdl(schema, table, getColumns(schema, table), indexes, foreignKeys); + } + + List listIndexes(String schema, String table); + + List listForeignKeys(String schema, String table); + + List listTriggers(String schema, String table); + + default QueryResult executeQuery(String sql, String schema) { + return executeQuery(sql, schema, new ExecuteQueryOptions()); + } + + QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options); + + default QueryPageResult executeQueryPage(String sql, String schema) { + return executeQueryPage(sql, schema, new QueryPageOptions()); + } + + default QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + Connection conn = getConnection(); + if (conn == null) { + throw new IllegalStateException("Not connected"); + } + return JdbcExecutor.INSTANCE.executePage( + conn, + sql, + schema, + this::setSchemaSQL, + options, + JdbcExecutor.INSTANCE::defaultResultValue + ); + } + + default QueryPageResult fetchQueryPage(String sessionId, int pageSize) { + return JdbcExecutor.INSTANCE.fetchPage(sessionId, pageSize); + } + + default boolean closeQuerySession(String sessionId) { + return JdbcExecutor.INSTANCE.closeQuerySession(sessionId); + } + + /** + * Get DM execution plan. Supports two modes: + * mode="explain" (default) — direct plan, no execution + * mode="autotrace" — enable MONITOR_SQL_EXEC, execute SQL, then get plan with actual stats + * @return plan text + */ + default String getExplainInfo(String sql, String database, String schema, int timeoutSecs, String mode) { + throw new UnsupportedOperationException("getExplainInfo is not supported by this agent"); + } + + void disconnect(); + + Connection getConnection(); + + default QueryResult executeTransaction(List statements, String schema) { + Connection conn = getConnection(); + if (conn == null) { + throw new IllegalStateException("Not connected"); + } + return TransactionExecutor.executeUpdateStatements(conn, statements, schema, this::setSchemaSQL); + } + + default String setSchemaSQL(String schema) { + return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + static String buildTableDdl( + String schema, + String table, + List columns, + List indexes, + List foreignKeys + ) { + return DdlBuilder.buildTableDdl(schema, table, columns, indexes, foreignKeys); + } + + static String trimSql(String sql) { + String trimmed = sql.trim(); + while (trimmed.endsWith(";")) { + trimmed = trimmed.substring(0, trimmed.length() - 1).trim(); + } + return trimmed; + } + + static T unchecked(ThrowingSupplier supplier) { + try { + return supplier.get(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + interface ThrowingSupplier { + T get() throws Exception; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/DatabaseInfo.java b/agents/common/src/main/java/com/dbx/agent/DatabaseInfo.java new file mode 100644 index 0000000000..e2c2f12761 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/DatabaseInfo.java @@ -0,0 +1,41 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class DatabaseInfo { + private String name; + + public DatabaseInfo() { + this(""); + } + + public DatabaseInfo(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public void setName(String name) { + this.name = name; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof DatabaseInfo)) return false; + DatabaseInfo that = (DatabaseInfo) other; + return Objects.equals(name, that.name); + } + + @Override + public int hashCode() { + return Objects.hash(name); + } + + @Override + public String toString() { + return "DatabaseInfo(name=" + name + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/DdlBuilder.java b/agents/common/src/main/java/com/dbx/agent/DdlBuilder.java new file mode 100644 index 0000000000..e7549edc01 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/DdlBuilder.java @@ -0,0 +1,195 @@ +package com.dbx.agent; + +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; + +public final class DdlBuilder { + private DdlBuilder() { + } + + public static String buildTableDdl( + String schema, + String table, + List columns, + List indexes, + List foreignKeys + ) { + return buildTableDdl(schema, table, columns, indexes, foreignKeys, false, false); + } + + public static String buildTableDdl( + String schema, + String table, + List columns, + List indexes, + List foreignKeys, + boolean useBacktick + ) { + return buildTableDdl(schema, table, columns, indexes, foreignKeys, useBacktick, false); + } + + public static String buildTableDdl( + String schema, + String table, + List columns, + List indexes, + List foreignKeys, + boolean useBacktick, + boolean includeColumnComments + ) { + String tableRef = qualifiedName(schema, table, useBacktick); + List columnLines = new ArrayList<>(); + for (ColumnInfo column : columns) { + StringBuilder line = new StringBuilder(); + line.append(" "); + line.append(quoteIdent(column.getName(), useBacktick)); + line.append(" "); + line.append(columnTypeSql(column)); + if (!column.getIs_nullable()) { + line.append(" NOT NULL"); + } + if (notBlank(column.getColumn_default())) { + line.append(" DEFAULT "); + line.append(column.getColumn_default()); + } + columnLines.add(line.toString()); + } + + List primaryKeys = new ArrayList<>(); + for (ColumnInfo column : columns) { + if (column.getIs_primary_key()) { + primaryKeys.add(quoteIdent(column.getName(), useBacktick)); + } + } + if (!primaryKeys.isEmpty()) { + columnLines.add(" PRIMARY KEY (" + join(primaryKeys, ", ") + ")"); + } + + for (ForeignKeyInfo fk : foreignKeys) { + columnLines.add( + " CONSTRAINT " + quoteIdent(fk.getName(), useBacktick) + + " FOREIGN KEY (" + quoteIdent(fk.getColumn(), useBacktick) + ") " + + "REFERENCES " + quoteIdent(fk.getRef_table(), useBacktick) + "(" + quoteIdent(fk.getRef_column(), useBacktick) + ")" + ); + } + + StringBuilder ddl = new StringBuilder(); + ddl.append("CREATE TABLE "); + ddl.append(tableRef); + ddl.append(" (\n"); + ddl.append(join(columnLines, ",\n")); + ddl.append("\n);\n"); + + if (includeColumnComments) { + for (ColumnInfo column : columns) { + if (!notBlank(column.getComment())) { + continue; + } + ddl.append("\nCOMMENT ON COLUMN "); + ddl.append(tableRef); + ddl.append("."); + ddl.append(quoteIdent(column.getName(), useBacktick)); + ddl.append(" IS '"); + ddl.append(sqlStringBody(column.getComment())); + ddl.append("';"); + } + } + + for (IndexInfo index : indexes) { + if (index.getIs_primary()) { + continue; + } + String unique = index.getIs_unique() ? "UNIQUE " : ""; + String using = notBlank(index.getIndex_type()) ? " USING " + index.getIndex_type() : ""; + List quotedColumns = new ArrayList<>(); + for (String column : index.getColumns()) { + quotedColumns.add(quoteIdent(column, useBacktick)); + } + String filter = notBlank(index.getFilter()) ? " WHERE " + index.getFilter() : ""; + ddl.append("\nCREATE "); + ddl.append(unique); + ddl.append("INDEX "); + ddl.append(quoteIdent(index.getName(), useBacktick)); + ddl.append(" ON "); + ddl.append(tableRef); + ddl.append(using); + ddl.append(" ("); + ddl.append(join(quotedColumns, ", ")); + ddl.append(")"); + ddl.append(filter); + ddl.append(";"); + if (notBlank(index.getComment())) { + ddl.append("\nCOMMENT ON INDEX "); + if (notBlank(schema)) { + ddl.append(quoteIdent(schema, useBacktick)); + ddl.append("."); + } + ddl.append(quoteIdent(index.getName(), useBacktick)); + ddl.append(" IS '"); + ddl.append(sqlStringBody(index.getComment())); + ddl.append("';"); + } + } + + return ddl.toString(); + } + + private static String quoteIdent(String identifier, boolean useBacktick) { + return useBacktick + ? JdbcIdentifiers.INSTANCE.backtick(identifier) + : JdbcIdentifiers.INSTANCE.doubleQuote(identifier); + } + + private static String qualifiedName(String schema, String name, boolean useBacktick) { + if (!notBlank(schema)) { + return quoteIdent(name, useBacktick); + } + return quoteIdent(schema, useBacktick) + "." + quoteIdent(name, useBacktick); + } + + private static String columnTypeSql(ColumnInfo column) { + String type = column.getData_type(); + String normalized = type.toLowerCase(Locale.ROOT); + if (isCharacterType(normalized) && column.getCharacter_maximum_length() != null) { + return type + "(" + column.getCharacter_maximum_length() + ")"; + } + if (isNumericType(normalized) && column.getNumeric_precision() != null) { + if (column.getNumeric_scale() != null) { + return type + "(" + column.getNumeric_precision() + ", " + column.getNumeric_scale() + ")"; + } + return type + "(" + column.getNumeric_precision() + ")"; + } + return type; + } + + private static boolean isCharacterType(String normalized) { + return "character varying".equals(normalized) + || "varchar".equals(normalized) + || "char".equals(normalized) + || "character".equals(normalized); + } + + private static boolean isNumericType(String normalized) { + return "numeric".equals(normalized) || "decimal".equals(normalized); + } + + private static boolean notBlank(String value) { + return value != null && !value.trim().isEmpty(); + } + + private static String sqlStringBody(String value) { + return value.replace("'", "''"); + } + + private static String join(List values, String separator) { + StringBuilder builder = new StringBuilder(); + for (int i = 0; i < values.size(); i++) { + if (i > 0) { + builder.append(separator); + } + builder.append(values.get(i)); + } + return builder.toString(); + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/EwkbWktDecoder.java b/agents/common/src/main/java/com/dbx/agent/EwkbWktDecoder.java new file mode 100644 index 0000000000..96af67a6f6 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/EwkbWktDecoder.java @@ -0,0 +1,671 @@ +package com.dbx.agent; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.List; + +/** + * Decode PostGIS EWKB / WKB byte streams (and PGgeometry / hex strings) into + * OGC Well-Known Text. Behaves identically to the Rust implementation in + * {@code crates/dbx-core/src/db/postgres.rs::ewkb_to_wkt}, including the + * {@code Z} / {@code M} / {@code ZM} dimension suffixes and {@code EMPTY} + * geometries. Used by the JDBC agent layer so that PostgresLike databases + * such as HighGo / KingBase / Vastbase / openGauss / GaussDB return WKT for + * {@code geometry}/{@code geography} columns instead of opaque hex blobs. + * + *

The decoder is permissive: any unrecognized payload (truncated bytes, + * unsupported geometry kind, non-hex strings, etc.) falls back to a hex + * representation prefixed with {@code 0x} so the cell remains visible. + */ +public final class EwkbWktDecoder { + /** Reflectively-detected PGgeometry classes that expose getValue() returning EWKT text. */ + private static final String[] PGGEOMETRY_CLASS_NAMES = new String[]{ + "org.postgis.PGgeometry", // legacy PostGIS-JDBC + "net.postgis.jdbc.PGgeometry", // modern PostGIS-JDBC (>= 2024 release) + "com.highgo.jdbc.geometry.PGgeometry" + }; + + /** + * Generic {@code org.postgresql.util.PGobject}-shaped classes. Their + * {@code getValue()} returns the textual representation the driver got off + * the wire. For PostGIS columns on HighGo / KingBase / Vastbase that text + * is the EWKB hex (e.g. {@code 0101000020E61000...}) — we still need to run + * it through the EWKB hex parser. + */ + private static final String[] PGOBJECT_CLASS_NAMES = new String[]{ + "org.postgresql.util.PGobject", + "com.highgo.jdbc.util.PGobject", + "org.postgresql.geometric.PGobject" + }; + + private EwkbWktDecoder() { + } + + /** + * Decode a JDBC-returned object from a {@code geometry} / {@code geography} + * column into a textual representation, preferring WKT. + * + *

Accepted inputs: + *

    + *
  • {@code null} — returned as-is.
    • + *
  • {@code byte[]} — interpreted as raw EWKB.
    • + *
  • A {@code PGgeometry}-like object exposing {@code getValue()} — the + * returned string is taken verbatim (already EWKT/WKT).
    • + *
  • {@link String} — if it looks like hex EWKB (with optional + * {@code 0x}/{@code \\x} prefix and a leading endian byte of + * {@code 00}/{@code 01}), it is decoded; otherwise returned as-is.
    • + *
+ * + * @param value raw column value + * @return decoded WKT, or a {@code 0x}-prefixed hex fallback, or {@code null} + */ + public static Object decode(Object value) { + if (value == null) { + return null; + } + if (value instanceof byte[]) { + return decodeBytes((byte[]) value); + } + if (value instanceof String) { + return decodeString((String) value); + } + // PGgeometry / PGobject and friends — pull out their getValue() text + // and run it through the same string path. For PostGIS-JDBC types + // getValue() yields EWKT (passes through unchanged); for the generic + // PGobject driver returns EWKB hex, which decodeString() decodes. + String pgValue = tryReflectiveGetValue(value); + if (pgValue != null) { + return decodeString(pgValue); + } + // Last resort: rely on toString() / underlying value. + return value.toString(); + } + + /** Decode a raw EWKB byte sequence; returns hex {@code 0x..} on failure. */ + static String decodeBytes(byte[] raw) { + if (raw == null) { + return null; + } + String wkt = parseEwkb(raw); + if (wkt != null) { + return wkt; + } + return bytesToHexWithPrefix(raw); + } + + /** + * Decode a string column. If the text matches an EWKB hex layout the bytes + * are decoded; otherwise (already WKT, ST_AsText output, etc.) returned + * unchanged. + */ + static String decodeString(String value) { + if (value == null) { + return null; + } + byte[] bytes = tryParseHex(value); + if (bytes == null) { + return value; + } + String wkt = parseEwkb(bytes); + if (wkt != null) { + return wkt; + } + // looked like hex but couldn't be parsed — preserve hex form. + return ensureHexPrefix(value); + } + + /** + * Try to extract textual value from PGgeometry / PGobject-like objects via + * {@code getValue()}. Returns {@code null} if the class doesn't match any + * known PostGIS / PostgreSQL JDBC type (or the call fails reflectively). + */ + private static String tryReflectiveGetValue(Object value) { + Class cls = value.getClass(); + if (!matchesAnyClassName(cls, PGGEOMETRY_CLASS_NAMES) + && !matchesAnyClassName(cls, PGOBJECT_CLASS_NAMES)) { + return null; + } + try { + Method getValue = cls.getMethod("getValue"); + Object result = getValue.invoke(value); + if (result instanceof String) { + return (String) result; + } + } catch (ReflectiveOperationException ignored) { + // fall through + } + return null; + } + + private static boolean matchesAnyClassName(Class cls, String[] names) { + Class current = cls; + while (current != null) { + for (String name : names) { + if (current.getName().equals(name)) { + return true; + } + for (Class iface : current.getInterfaces()) { + if (iface.getName().equals(name)) { + return true; + } + } + } + current = current.getSuperclass(); + } + return false; + } + + /** @return parsed bytes, or {@code null} if {@code value} is not hex EWKB. */ + static byte[] tryParseHex(String value) { + String s = value; + if (s.length() >= 2 && (s.startsWith("0x") || s.startsWith("0X"))) { + s = s.substring(2); + } else if (s.length() >= 2 && (s.startsWith("\\x") || s.startsWith("\\X"))) { + s = s.substring(2); + } + if (s.length() < 10 || (s.length() & 1) != 0) { + return null; + } + // First byte must be the WKB endian marker (00 = big, 01 = little) + char e0 = s.charAt(0); + char e1 = s.charAt(1); + if (!((e0 == '0') && (e1 == '0' || e1 == '1'))) { + return null; + } + for (int i = 0; i < s.length(); i++) { + char c = s.charAt(i); + if (!((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'))) { + return null; + } + } + byte[] out = new byte[s.length() / 2]; + for (int i = 0; i < out.length; i++) { + int hi = Character.digit(s.charAt(i * 2), 16); + int lo = Character.digit(s.charAt(i * 2 + 1), 16); + out[i] = (byte) ((hi << 4) | lo); + } + return out; + } + + private static String ensureHexPrefix(String value) { + if (value.startsWith("0x") || value.startsWith("0X")) { + return value; + } + return "0x" + value; + } + + static String bytesToHexWithPrefix(byte[] bytes) { + StringBuilder sb = new StringBuilder(bytes.length * 2 + 2); + sb.append("0x"); + for (byte b : bytes) { + sb.append(Character.forDigit((b >> 4) & 0xF, 16)); + sb.append(Character.forDigit(b & 0xF, 16)); + } + return sb.toString(); + } + + // ── EWKB parser ──────────────────────────────────────────────────────── + + private static String parseEwkb(byte[] raw) { + if (raw == null || raw.length == 0) { + return null; + } + Reader reader = new Reader(raw); + Geometry g = parseGeometry(reader); + if (g == null || reader.pos != raw.length) { + return null; + } + return geometryToWkt(g); + } + + private static Geometry parseGeometry(Reader reader) { + Integer endian = reader.readU8(); + if (endian == null) { + return null; + } + boolean little; + if (endian == 0) { + little = false; + } else if (endian == 1) { + little = true; + } else { + return null; + } + Long typeWordBoxed = reader.readU32(little); + if (typeWordBoxed == null) { + return null; + } + long typeWord = typeWordBoxed; + Dims parsed = parseDimensions(typeWord); + long baseType = parsed.baseType; + if (parsed.hasSrid) { + if (reader.readU32(little) == null) { + return null; + } + } + Dimensions dims = parsed.dims; + switch ((int) baseType) { + case 1: { // Point + int len = dims.coordinateLen(); + List coords = new ArrayList<>(len); + for (int i = 0; i < len; i++) { + Double v = reader.readF64(little); + if (v == null) { + return null; + } + coords.add(v); + } + boolean allNan = true; + for (Double v : coords) { + if (!Double.isNaN(v)) { + allNan = false; + break; + } + } + return new Geometry(GeomKind.POINT, dims, allNan ? null : coords, null, null, null, null, null); + } + case 2: { // LineString + List> points = parsePoints(reader, little, dims); + if (points == null) return null; + return new Geometry(GeomKind.LINESTRING, dims, null, points, null, null, null, null); + } + case 3: { // Polygon + Long ringCount = reader.readU32(little); + if (ringCount == null) return null; + List>> rings = new ArrayList<>(); + for (int i = 0; i < ringCount; i++) { + List> ring = parsePoints(reader, little, dims); + if (ring == null) return null; + rings.add(ring); + } + return new Geometry(GeomKind.POLYGON, dims, null, null, rings, null, null, null); + } + case 4: { // MultiPoint + Long count = reader.readU32(little); + if (count == null) return null; + List> mpoints = new ArrayList<>(); + for (int i = 0; i < count; i++) { + List coords = readPointCoords(reader, dims); + if (coords == null) { + // sub-point parse failure (returning null != EMPTY); abort. + return null; + } + boolean allNan = true; + for (Double v : coords) { + if (!Double.isNaN(v)) { + allNan = false; + break; + } + } + mpoints.add(allNan ? null : coords); + } + return new Geometry(GeomKind.MULTIPOINT, dims, null, null, null, mpoints, null, null); + } + case 5: { // MultiLineString + Long count = reader.readU32(little); + if (count == null) return null; + List>> lines = new ArrayList<>(); + for (int i = 0; i < count; i++) { + Geometry sub = parseGeometry(reader); + if (sub == null || sub.kind != GeomKind.LINESTRING) return null; + lines.add(sub.points); + } + return new Geometry(GeomKind.MULTILINESTRING, dims, null, null, lines, null, null, null); + } + case 6: { // MultiPolygon + Long count = reader.readU32(little); + if (count == null) return null; + List>>> polys = new ArrayList<>(); + for (int i = 0; i < count; i++) { + Geometry sub = parseGeometry(reader); + if (sub == null || sub.kind != GeomKind.POLYGON) return null; + polys.add(sub.rings); + } + return new Geometry(GeomKind.MULTIPOLYGON, dims, null, null, null, null, polys, null); + } + case 7: { // GeometryCollection + Long count = reader.readU32(little); + if (count == null) return null; + List children = new ArrayList<>(); + for (int i = 0; i < count; i++) { + Geometry sub = parseGeometry(reader); + if (sub == null) return null; + children.add(sub); + } + return new Geometry(GeomKind.GEOMETRYCOLLECTION, dims, null, null, null, null, null, children); + } + default: + return null; + } + } + + private static List> parsePoints(Reader reader, boolean little, Dimensions dims) { + Long count = reader.readU32(little); + if (count == null) return null; + int dim = dims.coordinateLen(); + List> result = new ArrayList<>(); + for (int i = 0; i < count; i++) { + List coords = new ArrayList<>(dim); + for (int j = 0; j < dim; j++) { + Double v = reader.readF64(little); + if (v == null) return null; + coords.add(v); + } + result.add(coords); + } + return result; + } + + private static List readPointCoords(Reader reader, Dimensions parentDims) { + Integer endian = reader.readU8(); + if (endian == null) return null; + boolean little; + if (endian == 0) { + little = false; + } else if (endian == 1) { + little = true; + } else { + return null; + } + Long typeWord = reader.readU32(little); + if (typeWord == null) return null; + Dims parsed = parseDimensions(typeWord); + if (parsed.baseType != 1 || !parsed.dims.equals(parentDims)) { + return null; + } + if (parsed.hasSrid) { + if (reader.readU32(little) == null) return null; + } + int dim = parsed.dims.coordinateLen(); + List coords = new ArrayList<>(dim); + for (int i = 0; i < dim; i++) { + Double v = reader.readF64(little); + if (v == null) return null; + coords.add(v); + } + return coords; + } + + private static Dims parseDimensions(long typeWord) { + long baseType = typeWord & 0x1FFF_FFFFL; + boolean hasZ = (typeWord & 0x8000_0000L) != 0; + boolean hasM = (typeWord & 0x4000_0000L) != 0; + boolean hasSrid = (typeWord & 0x2000_0000L) != 0; + if (baseType >= 3000) { + hasZ = true; + hasM = true; + baseType -= 3000; + } else if (baseType >= 2000) { + hasM = true; + baseType -= 2000; + } else if (baseType >= 1000) { + hasZ = true; + baseType -= 1000; + } + return new Dims(baseType, new Dimensions(hasZ, hasM), hasSrid); + } + + // ── WKT formatting ───────────────────────────────────────────────────── + + private static String geometryToWkt(Geometry g) { + String suffix = g.dims.suffix(); + switch (g.kind) { + case POINT: + if (g.coords == null) { + return "POINT" + suffix + " EMPTY"; + } + return "POINT" + suffix + "(" + formatCoord(g.coords) + ")"; + case LINESTRING: + if (g.points == null || g.points.isEmpty()) { + return "LINESTRING" + suffix + " EMPTY"; + } + return "LINESTRING" + suffix + "(" + formatCoordSeq(g.points) + ")"; + case POLYGON: + if (g.rings == null || g.rings.isEmpty()) { + return "POLYGON" + suffix + " EMPTY"; + } + return "POLYGON" + suffix + "(" + joinRings(g.rings) + ")"; + case MULTIPOINT: + if (g.multiPoints == null || g.multiPoints.isEmpty()) { + return "MULTIPOINT" + suffix + " EMPTY"; + } + return "MULTIPOINT" + suffix + "(" + joinMultiPoints(g.multiPoints) + ")"; + case MULTILINESTRING: + if (g.rings == null || g.rings.isEmpty()) { + return "MULTILINESTRING" + suffix + " EMPTY"; + } + return "MULTILINESTRING" + suffix + "(" + joinRings(g.rings) + ")"; + case MULTIPOLYGON: + if (g.polygons == null || g.polygons.isEmpty()) { + return "MULTIPOLYGON" + suffix + " EMPTY"; + } + return "MULTIPOLYGON" + suffix + "(" + joinPolygons(g.polygons) + ")"; + case GEOMETRYCOLLECTION: + if (g.children == null || g.children.isEmpty()) { + return "GEOMETRYCOLLECTION" + suffix + " EMPTY"; + } + StringBuilder sb = new StringBuilder("GEOMETRYCOLLECTION").append(suffix).append("("); + for (int i = 0; i < g.children.size(); i++) { + if (i > 0) sb.append(','); + sb.append(geometryToWkt(g.children.get(i))); + } + sb.append(')'); + return sb.toString(); + default: + throw new IllegalStateException("unknown geometry kind: " + g.kind); + } + } + + private static String formatCoord(List coords) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < coords.size(); i++) { + if (i > 0) sb.append(' '); + sb.append(formatDouble(coords.get(i))); + } + return sb.toString(); + } + + private static String formatCoordSeq(List> points) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < points.size(); i++) { + if (i > 0) sb.append(','); + sb.append(formatCoord(points.get(i))); + } + return sb.toString(); + } + + private static String joinRings(List>> rings) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < rings.size(); i++) { + if (i > 0) sb.append(','); + sb.append('(').append(formatCoordSeq(rings.get(i))).append(')'); + } + return sb.toString(); + } + + private static String joinMultiPoints(List> mpoints) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < mpoints.size(); i++) { + if (i > 0) sb.append(','); + List p = mpoints.get(i); + if (p == null) { + sb.append("EMPTY"); + } else { + sb.append('(').append(formatCoord(p)).append(')'); + } + } + return sb.toString(); + } + + private static String joinPolygons(List>>> polys) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < polys.size(); i++) { + if (i > 0) sb.append(','); + sb.append('(').append(joinRings(polys.get(i))).append(')'); + } + return sb.toString(); + } + + /** + * Format a double the same way Rust's {@code f64::Display} does — shortest + * round-trippable decimal, trailing {@code .0} stripped (so {@code 116.0} + * becomes {@code "116"}, not {@code "116.0"}). + */ + static String formatDouble(double value) { + if (Double.isNaN(value)) { + return "NaN"; + } + if (Double.isInfinite(value)) { + return value > 0 ? "inf" : "-inf"; + } + String s = Double.toString(value); + // Strip trailing ".0" for whole-number doubles so output matches Rust. + if (s.endsWith(".0")) { + return s.substring(0, s.length() - 2); + } + return s; + } + + // ── helper structs ───────────────────────────────────────────────────── + + private enum GeomKind { + POINT, LINESTRING, POLYGON, MULTIPOINT, MULTILINESTRING, MULTIPOLYGON, GEOMETRYCOLLECTION + } + + private static final class Geometry { + final GeomKind kind; + final Dimensions dims; + final List coords; // POINT + final List> points; // LINESTRING + final List>> rings; // POLYGON / MULTILINESTRING (lines == rings shape-wise) + final List> multiPoints; // MULTIPOINT (null entry = EMPTY) + final List>>> polygons; // MULTIPOLYGON + final List children; // GEOMETRYCOLLECTION + + Geometry( + GeomKind kind, + Dimensions dims, + List coords, + List> points, + List>> rings, + List> multiPoints, + List>>> polygons, + List children + ) { + this.kind = kind; + this.dims = dims; + this.coords = coords; + this.points = points; + this.rings = rings; + this.multiPoints = multiPoints; + this.polygons = polygons; + this.children = children; + } + } + + private static final class Dimensions { + final boolean hasZ; + final boolean hasM; + + Dimensions(boolean hasZ, boolean hasM) { + this.hasZ = hasZ; + this.hasM = hasM; + } + + String suffix() { + if (hasZ && hasM) return " ZM"; + if (hasZ) return " Z"; + if (hasM) return " M"; + return ""; + } + + int coordinateLen() { + return 2 + (hasZ ? 1 : 0) + (hasM ? 1 : 0); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof Dimensions)) return false; + Dimensions d = (Dimensions) o; + return hasZ == d.hasZ && hasM == d.hasM; + } + + @Override + public int hashCode() { + return (hasZ ? 1 : 0) * 31 + (hasM ? 1 : 0); + } + } + + private static final class Dims { + final long baseType; + final Dimensions dims; + final boolean hasSrid; + + Dims(long baseType, Dimensions dims, boolean hasSrid) { + this.baseType = baseType; + this.dims = dims; + this.hasSrid = hasSrid; + } + } + + private static final class Reader { + final byte[] raw; + int pos; + + Reader(byte[] raw) { + this.raw = raw; + } + + Integer readU8() { + if (pos >= raw.length) return null; + int v = raw[pos] & 0xFF; + pos += 1; + return v; + } + + Long readU32(boolean little) { + if (pos + 4 > raw.length) return null; + long v; + if (little) { + v = ((long) (raw[pos] & 0xFF)) + | ((long) (raw[pos + 1] & 0xFF) << 8) + | ((long) (raw[pos + 2] & 0xFF) << 16) + | ((long) (raw[pos + 3] & 0xFF) << 24); + } else { + v = ((long) (raw[pos] & 0xFF) << 24) + | ((long) (raw[pos + 1] & 0xFF) << 16) + | ((long) (raw[pos + 2] & 0xFF) << 8) + | ((long) (raw[pos + 3] & 0xFF)); + } + pos += 4; + return v & 0xFFFF_FFFFL; + } + + Double readF64(boolean little) { + if (pos + 8 > raw.length) return null; + long bits; + if (little) { + bits = ((long) (raw[pos] & 0xFF)) + | ((long) (raw[pos + 1] & 0xFF) << 8) + | ((long) (raw[pos + 2] & 0xFF) << 16) + | ((long) (raw[pos + 3] & 0xFF) << 24) + | ((long) (raw[pos + 4] & 0xFF) << 32) + | ((long) (raw[pos + 5] & 0xFF) << 40) + | ((long) (raw[pos + 6] & 0xFF) << 48) + | ((long) (raw[pos + 7] & 0xFF) << 56); + } else { + bits = ((long) (raw[pos] & 0xFF) << 56) + | ((long) (raw[pos + 1] & 0xFF) << 48) + | ((long) (raw[pos + 2] & 0xFF) << 40) + | ((long) (raw[pos + 3] & 0xFF) << 32) + | ((long) (raw[pos + 4] & 0xFF) << 24) + | ((long) (raw[pos + 5] & 0xFF) << 16) + | ((long) (raw[pos + 6] & 0xFF) << 8) + | ((long) (raw[pos + 7] & 0xFF)); + } + pos += 8; + return Double.longBitsToDouble(bits); + } + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/ExecuteQueryOptions.java b/agents/common/src/main/java/com/dbx/agent/ExecuteQueryOptions.java new file mode 100644 index 0000000000..ae33923cd2 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ExecuteQueryOptions.java @@ -0,0 +1,67 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class ExecuteQueryOptions { + private int maxRows; + private Integer fetchSize; + private int timeoutSecs; + + public ExecuteQueryOptions() { + this(JdbcExecutor.DEFAULT_MAX_ROWS, null, 0); + } + + public ExecuteQueryOptions(int maxRows, Integer fetchSize) { + this(maxRows, fetchSize, 0); + } + + public ExecuteQueryOptions(int maxRows, Integer fetchSize, int timeoutSecs) { + this.maxRows = maxRows; + this.fetchSize = fetchSize; + this.timeoutSecs = timeoutSecs; + } + + public int getMaxRows() { + return maxRows; + } + + public Integer getFetchSize() { + return fetchSize; + } + + public int getTimeoutSecs() { + return timeoutSecs; + } + + public void setMaxRows(int maxRows) { + this.maxRows = maxRows; + } + + public void setFetchSize(Integer fetchSize) { + this.fetchSize = fetchSize; + } + + public void setTimeoutSecs(int timeoutSecs) { + this.timeoutSecs = timeoutSecs; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ExecuteQueryOptions)) return false; + ExecuteQueryOptions that = (ExecuteQueryOptions) other; + return maxRows == that.maxRows + && timeoutSecs == that.timeoutSecs + && Objects.equals(fetchSize, that.fetchSize); + } + + @Override + public int hashCode() { + return Objects.hash(maxRows, fetchSize, timeoutSecs); + } + + @Override + public String toString() { + return "ExecuteQueryOptions(maxRows=" + maxRows + ", fetchSize=" + fetchSize + ", timeoutSecs=" + timeoutSecs + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/ExecuteQueryParams.java b/agents/common/src/main/java/com/dbx/agent/ExecuteQueryParams.java new file mode 100644 index 0000000000..0f96201889 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ExecuteQueryParams.java @@ -0,0 +1,78 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class ExecuteQueryParams { + private String sql; + private String schema; + private Integer maxRows; + private Integer fetchSize; + + public ExecuteQueryParams() { + this("", null, null, null); + } + + public ExecuteQueryParams(String sql, String schema, Integer maxRows, Integer fetchSize) { + this.sql = sql; + this.schema = schema; + this.maxRows = maxRows; + this.fetchSize = fetchSize; + } + + public String getSql() { + return sql; + } + + public String getSchema() { + return schema; + } + + public Integer getMaxRows() { + return maxRows; + } + + public Integer getFetchSize() { + return fetchSize; + } + + public void setSql(String sql) { + this.sql = sql; + } + + public void setSchema(String schema) { + this.schema = schema; + } + + public void setMaxRows(Integer maxRows) { + this.maxRows = maxRows; + } + + public void setFetchSize(Integer fetchSize) { + this.fetchSize = fetchSize; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ExecuteQueryParams)) return false; + ExecuteQueryParams that = (ExecuteQueryParams) other; + return Objects.equals(sql, that.sql) + && Objects.equals(schema, that.schema) + && Objects.equals(maxRows, that.maxRows) + && Objects.equals(fetchSize, that.fetchSize); + } + + @Override + public int hashCode() { + return Objects.hash(sql, schema, maxRows, fetchSize); + } + + @Override + public String toString() { + return "ExecuteQueryParams(sql=" + sql + + ", schema=" + schema + + ", maxRows=" + maxRows + + ", fetchSize=" + fetchSize + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/ForeignKeyInfo.java b/agents/common/src/main/java/com/dbx/agent/ForeignKeyInfo.java new file mode 100644 index 0000000000..de018df94c --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ForeignKeyInfo.java @@ -0,0 +1,78 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class ForeignKeyInfo { + private String name; + private String column; + private String ref_table; + private String ref_column; + + public ForeignKeyInfo() { + this("", "", "", ""); + } + + public ForeignKeyInfo(String name, String column, String ref_table, String ref_column) { + this.name = name; + this.column = column; + this.ref_table = ref_table; + this.ref_column = ref_column; + } + + public String getName() { + return name; + } + + public String getColumn() { + return column; + } + + public String getRef_table() { + return ref_table; + } + + public String getRef_column() { + return ref_column; + } + + public void setName(String name) { + this.name = name; + } + + public void setColumn(String column) { + this.column = column; + } + + public void setRef_table(String ref_table) { + this.ref_table = ref_table; + } + + public void setRef_column(String ref_column) { + this.ref_column = ref_column; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ForeignKeyInfo)) return false; + ForeignKeyInfo that = (ForeignKeyInfo) other; + return Objects.equals(name, that.name) + && Objects.equals(column, that.column) + && Objects.equals(ref_table, that.ref_table) + && Objects.equals(ref_column, that.ref_column); + } + + @Override + public int hashCode() { + return Objects.hash(name, column, ref_table, ref_column); + } + + @Override + public String toString() { + return "ForeignKeyInfo(name=" + name + + ", column=" + column + + ", ref_table=" + ref_table + + ", ref_column=" + ref_column + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/IndexInfo.java b/agents/common/src/main/java/com/dbx/agent/IndexInfo.java new file mode 100644 index 0000000000..ea38747f26 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/IndexInfo.java @@ -0,0 +1,141 @@ +package com.dbx.agent; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +public final class IndexInfo { + private String name; + private List columns; + private boolean is_unique; + private boolean is_primary; + private String filter; + private String index_type; + private List included_columns; + private String comment; + + public IndexInfo() { + this("", Collections.emptyList(), false, false); + } + + public IndexInfo(String name, List columns, boolean is_unique, boolean is_primary) { + this(name, columns, is_unique, is_primary, null, null, null, null); + } + + public IndexInfo( + String name, + List columns, + boolean is_unique, + boolean is_primary, + String filter, + String index_type, + List included_columns, + String comment + ) { + this.name = name; + this.columns = columns == null ? Collections.emptyList() : columns; + this.is_unique = is_unique; + this.is_primary = is_primary; + this.filter = filter; + this.index_type = index_type; + this.included_columns = included_columns; + this.comment = comment; + } + + public String getName() { + return name; + } + + public List getColumns() { + return columns; + } + + public boolean getIs_unique() { + return is_unique; + } + + public boolean getIs_primary() { + return is_primary; + } + + public String getFilter() { + return filter; + } + + public String getIndex_type() { + return index_type; + } + + public List getIncluded_columns() { + return included_columns; + } + + public String getComment() { + return comment; + } + + public void setName(String name) { + this.name = name; + } + + public void setColumns(List columns) { + this.columns = columns; + } + + public void setIs_unique(boolean is_unique) { + this.is_unique = is_unique; + } + + public void setIs_primary(boolean is_primary) { + this.is_primary = is_primary; + } + + public void setFilter(String filter) { + this.filter = filter; + } + + public void setIndex_type(String index_type) { + this.index_type = index_type; + } + + public void setIncluded_columns(List included_columns) { + this.included_columns = included_columns; + } + + public void setComment(String comment) { + this.comment = comment; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof IndexInfo)) return false; + IndexInfo that = (IndexInfo) other; + return is_unique == that.is_unique + && is_primary == that.is_primary + && Objects.equals(name, that.name) + && Objects.equals(columns, that.columns) + && Objects.equals(filter, that.filter) + && Objects.equals(index_type, that.index_type) + && Objects.equals(included_columns, that.included_columns) + && Objects.equals(comment, that.comment); + } + + @Override + public int hashCode() { + return Objects.hash(name, columns, is_unique, is_primary, filter, index_type, included_columns, comment); + } + + @Override + public String toString() { + return "IndexInfo(name=" + name + + ", columns=" + columns + + ", is_unique=" + is_unique + + ", is_primary=" + is_primary + + ", filter=" + filter + + ", index_type=" + index_type + + ", included_columns=" + included_columns + + ", comment=" + comment + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcAgentProfile.java b/agents/common/src/main/java/com/dbx/agent/JdbcAgentProfile.java new file mode 100644 index 0000000000..433f817ca5 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/JdbcAgentProfile.java @@ -0,0 +1,223 @@ +package com.dbx.agent; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Set; + +public class JdbcAgentProfile { + private final String driverClass; + private final String urlTemplate; + private final int defaultPort; + private final boolean skipExecutionContext; + private final Set excludedSchemas; + private final List tableTypes; + private final String identifierQuote; + private final String schemaSwitchPrefix; + private final boolean catalogFallbackEnabled; + private final boolean nativeTableDdlSupported; + private final boolean objectSourceSupported; + private final boolean triggersSupported; + + public JdbcAgentProfile(String driverClass, String urlTemplate) { + this(driverClass, urlTemplate, 0); + } + + public JdbcAgentProfile(String driverClass, String urlTemplate, int defaultPort) { + this(driverClass, urlTemplate, defaultPort, false); + } + + public JdbcAgentProfile(String driverClass, String urlTemplate, int defaultPort, boolean skipExecutionContext) { + this( + driverClass, + urlTemplate, + defaultPort, + skipExecutionContext, + Collections.emptySet(), + Arrays.asList("TABLE", "VIEW", "BASE TABLE", "MATERIALIZED VIEW", "SYSTEM TABLE", "SYSTEM VIEW") + ); + } + + public JdbcAgentProfile( + String driverClass, + String urlTemplate, + int defaultPort, + boolean skipExecutionContext, + Set excludedSchemas, + List tableTypes + ) { + this( + driverClass, + urlTemplate, + defaultPort, + skipExecutionContext, + excludedSchemas, + tableTypes, + "\"", + "SET SCHEMA", + true, + false, + false, + false + ); + } + + public JdbcAgentProfile( + String driverClass, + String urlTemplate, + int defaultPort, + boolean skipExecutionContext, + Set excludedSchemas, + List tableTypes, + String identifierQuote, + String schemaSwitchPrefix, + boolean catalogFallbackEnabled, + boolean nativeTableDdlSupported, + boolean objectSourceSupported, + boolean triggersSupported + ) { + this.driverClass = driverClass; + this.urlTemplate = urlTemplate; + this.defaultPort = defaultPort; + this.skipExecutionContext = skipExecutionContext; + this.excludedSchemas = excludedSchemas; + this.tableTypes = tableTypes; + this.identifierQuote = identifierQuote; + this.schemaSwitchPrefix = schemaSwitchPrefix; + this.catalogFallbackEnabled = catalogFallbackEnabled; + this.nativeTableDdlSupported = nativeTableDdlSupported; + this.objectSourceSupported = objectSourceSupported; + this.triggersSupported = triggersSupported; + } + + public String getDriverClass() { + return driverClass; + } + + public String getUrlTemplate() { + return urlTemplate; + } + + public int getDefaultPort() { + return defaultPort; + } + + public boolean getSkipExecutionContext() { + return skipExecutionContext; + } + + public Set getExcludedSchemas() { + return excludedSchemas; + } + + public List getTableTypes() { + return tableTypes; + } + + public String getIdentifierQuote() { + return identifierQuote; + } + + public String getSchemaSwitchPrefix() { + return schemaSwitchPrefix; + } + + public boolean getCatalogFallbackEnabled() { + return catalogFallbackEnabled; + } + + public boolean getNativeTableDdlSupported() { + return nativeTableDdlSupported; + } + + public boolean getObjectSourceSupported() { + return objectSourceSupported; + } + + public boolean getTriggersSupported() { + return triggersSupported; + } + + public String buildUrl(ConnectParams params) { + if (!params.getConnection_string().trim().isEmpty()) { + return params.getConnection_string(); + } + int port = params.getPort() > 0 ? params.getPort() : defaultPort; + String base = urlTemplate + .replace("{host}", params.getHost()) + .replace("{port}", Integer.toString(port)) + .replace("{database}", params.getDatabase()); + return appendUrlParams(base, params.getUrl_params()); + } + + public String quoteIdentifier(String identifier) { + return identifierQuote + identifier.replace(identifierQuote, identifierQuote + identifierQuote) + identifierQuote; + } + + public String schemaSwitchSql(String schema) { + return schemaSwitchSql(schema, identifierQuote); + } + + public String schemaSwitchSql(String schema, String quote) { + if (skipExecutionContext) { + return ""; + } + return schemaSwitchPrefix + " " + quote + schema.replace(quote, quote + quote) + quote; + } + + private static String appendUrlParams(String url, String urlParams) { + String params = trimUrlParams(urlParams); + if (params.isEmpty()) { + return url; + } + if (usesColonProperties(url) && !params.endsWith(";")) { + params = params + ";"; + } + String separator = urlParamSeparator(url); + return url + separator + params; + } + + private static String trimUrlParams(String urlParams) { + String value = urlParams == null ? "" : urlParams.trim(); + while (value.startsWith("?") || value.startsWith("&") || value.startsWith(";") || value.startsWith(":") || value.startsWith(",")) { + value = value.substring(1); + } + return value; + } + + private static String urlParamSeparator(String url) { + if (startsWithIgnoreCase(url, "jdbc:sqlserver:") || startsWithIgnoreCase(url, "jdbc:exa:")) { + return url.endsWith(";") ? "" : ";"; + } + if (startsWithIgnoreCase(url, "jdbc:teradata:")) { + return url.endsWith(",") ? "" : ","; + } + if (usesColonProperties(url)) { + if (url.endsWith(":") || url.endsWith(";")) { + return ""; + } + return hasColonProperties(url) ? ";" : ":"; + } + return url.contains("?") ? "&" : "?"; + } + + private static boolean usesColonProperties(String url) { + return startsWithIgnoreCase(url, "jdbc:db2:") || startsWithIgnoreCase(url, "jdbc:informix-sqli:"); + } + + private static boolean hasColonProperties(String url) { + int schemeEnd = url.indexOf("://"); + if (schemeEnd < 0) { + return false; + } + int pathStart = url.indexOf('/', schemeEnd + 3); + if (pathStart < 0) { + return false; + } + return url.indexOf(':', pathStart + 1) >= 0; + } + + private static boolean startsWithIgnoreCase(String value, String prefix) { + return value.regionMatches(true, 0, prefix, 0, prefix.length()); + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java new file mode 100644 index 0000000000..73f7c36cf3 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java @@ -0,0 +1,652 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.sql.Blob; +import java.sql.Clob; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.SQLXML; +import java.sql.Statement; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Function; + +public final class JdbcExecutor { + public static final JdbcExecutor INSTANCE = new JdbcExecutor(); + public static final int DEFAULT_MAX_ROWS = 10000; + public static final long QUERY_SESSION_IDLE_TIMEOUT_MILLIS = 10 * 60 * 1000L; + + private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + + private JdbcExecutor() { + } + + public QueryResult execute(Connection conn, String sql, String schema, Function setSchemaSql) { + return execute(conn, sql, schema, setSchemaSql, DEFAULT_MAX_ROWS, null, this::defaultResultValue); + } + + public QueryResult execute( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + int maxRows, + Integer fetchSize + ) { + return execute(conn, sql, schema, setSchemaSql, maxRows, fetchSize, this::defaultResultValue); + } + + public QueryResult execute( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + int maxRows, + Integer fetchSize, + ResultValueReader valueReader + ) { + return execute(conn, sql, schema, setSchemaSql, maxRows, fetchSize, 0, valueReader); + } + + public QueryResult execute( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + int maxRows, + Integer fetchSize, + int timeoutSecs, + ResultValueReader valueReader + ) { + return unchecked(() -> { + String trimmedSql = trimSql(sql); + String upperSql = trimmedSql.toUpperCase(Locale.ROOT).trim(); + long start = System.currentTimeMillis(); + + if ("BEGIN".equals(upperSql) || "BEGIN TRANSACTION".equals(upperSql)) { + conn.setAutoCommit(false); + return emptyQueryResult(start); + } + if ("COMMIT".equals(upperSql)) { + conn.commit(); + conn.setAutoCommit(true); + return emptyQueryResult(start); + } + if ("ROLLBACK".equals(upperSql)) { + conn.rollback(); + conn.setAutoCommit(true); + return emptyQueryResult(start); + } + + applySchema(conn, schema, setSchemaSql); + + try (Statement stmt = conn.createStatement()) { + int effectiveMaxRows = Math.max(maxRows, 1); + stmt.setMaxRows(effectiveMaxRows + 1); + applyQueryTimeout(stmt, timeoutSecs); + if (fetchSize != null && fetchSize > 0) { + stmt.setFetchSize(fetchSize); + } + boolean hasResultSet = stmt.execute(trimmedSql); + long elapsed = System.currentTimeMillis() - start; + if (hasResultSet) { + try (ResultSet rs = stmt.getResultSet()) { + return readResultSet(rs, elapsed, effectiveMaxRows, valueReader); + } + } + + int updateCount = stmt.getUpdateCount(); + return new QueryResult( + Collections.emptyList(), + Collections.emptyList(), + updateCount >= 0 ? updateCount : 0, + elapsed, + false + ); + } + }); + } + + public QueryResult readResultSet(ResultSet rs, long executionTimeMs) { + return readResultSet(rs, executionTimeMs, DEFAULT_MAX_ROWS, this::defaultResultValue); + } + + public QueryResult readResultSet( + ResultSet rs, + long executionTimeMs, + int maxRows, + ResultValueReader valueReader + ) { + return unchecked(() -> { + ResultSetMetaData meta = rs.getMetaData(); + int colCount = meta.getColumnCount(); + List columns = new ArrayList<>(); + List columnTypes = new ArrayList<>(); + String[] typeNameByIndex = new String[colCount]; + for (int i = 1; i <= colCount; i++) { + columns.add(meta.getColumnLabel(i)); + String typeName = safeColumnTypeName(meta, i); + columnTypes.add(typeName); + typeNameByIndex[i - 1] = typeName; + } + + List> rows = new ArrayList<>(); + boolean truncated = false; + while (rs.next()) { + if (rows.size() >= maxRows) { + truncated = true; + break; + } + rows.add(rowValues(rs, valueReader, typeNameByIndex)); + } + + return new QueryResult(columns, columnTypes, rows, 0L, executionTimeMs, truncated); + }); + } + + public QueryPageResult executePage( + Connection conn, + String sql, + String schema, + Function setSchemaSql + ) { + return executePage(conn, sql, schema, setSchemaSql, new QueryPageOptions(), this::defaultResultValue); + } + + public QueryPageResult executePage( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + QueryPageOptions options + ) { + return executePage(conn, sql, schema, setSchemaSql, options, this::defaultResultValue); + } + + public QueryPageResult executePage( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + QueryPageOptions options, + ResultValueReader valueReader + ) { + return unchecked(() -> { + expireIdleQuerySessions(); + String trimmedSql = trimSql(sql); + String upperSql = trimmedSql.toUpperCase(Locale.ROOT).trim(); + long start = System.currentTimeMillis(); + + if ("BEGIN".equals(upperSql) || "BEGIN TRANSACTION".equals(upperSql)) { + conn.setAutoCommit(false); + return emptyQueryPageResult(start); + } + if ("COMMIT".equals(upperSql)) { + conn.commit(); + conn.setAutoCommit(true); + return emptyQueryPageResult(start); + } + if ("ROLLBACK".equals(upperSql)) { + conn.rollback(); + conn.setAutoCommit(true); + return emptyQueryPageResult(start); + } + + applySchema(conn, schema, setSchemaSql); + + Statement stmt = conn.createStatement(); + try { + applyQueryTimeout(stmt, options.getTimeoutSecs()); + if (options.getFetchSize() != null && options.getFetchSize() > 0) { + stmt.setFetchSize(options.getFetchSize()); + } + boolean hasResultSet = stmt.execute(trimmedSql); + long elapsed = System.currentTimeMillis() - start; + if (!hasResultSet) { + int updateCount = stmt.getUpdateCount(); + stmt.close(); + return new QueryPageResult( + Collections.emptyList(), + Collections.emptyList(), + updateCount >= 0 ? updateCount : 0, + elapsed + ); + } + + ResultSet rs = stmt.getResultSet(); + ResultSetMetaData meta = rs.getMetaData(); + String sessionId = UUID.randomUUID().toString(); + int colCount = meta.getColumnCount(); + List columns = new ArrayList<>(colCount); + List columnTypes = new ArrayList<>(colCount); + String[] typeNameByIndex = new String[colCount]; + for (int i = 1; i <= colCount; i++) { + columns.add(meta.getColumnLabel(i)); + String typeName = safeColumnTypeName(meta, i); + columnTypes.add(typeName); + typeNameByIndex[i - 1] = typeName; + } + QuerySession session = new QuerySession( + sessionId, + stmt, + rs, + columns, + columnTypes, + typeNameByIndex, + Math.max(options.getMaxRows(), 1), + valueReader + ); + sessions.put(sessionId, session); + return readSessionPage(session, options.getPageSize(), elapsed); + } catch (Exception e) { + try { + stmt.close(); + } catch (Exception ignored) { + } + throw e; + } + }); + } + + public QueryPageResult fetchPage(String sessionId, int pageSize) { + expireIdleQuerySessions(); + QuerySession session = sessions.get(sessionId); + if (session == null) { + throw new IllegalArgumentException("Query session not found"); + } + synchronized (session) { + return readSessionPage(session, pageSize, 0L); + } + } + + public boolean closeQuerySession(String sessionId) { + return closeSession(sessionId); + } + + public void closeAllQuerySessions() { + List ids = new ArrayList<>(sessions.keySet()); + for (String id : ids) { + closeSession(id); + } + } + + public int expireIdleQuerySessions() { + return expireIdleQuerySessions(System.currentTimeMillis(), QUERY_SESSION_IDLE_TIMEOUT_MILLIS); + } + + public int expireIdleQuerySessions(long nowMillis, long idleTimeoutMillis) { + if (idleTimeoutMillis < 0) { + return 0; + } + int closed = 0; + List snapshot = new ArrayList<>(sessions.values()); + for (QuerySession session : snapshot) { + boolean expired; + synchronized (session) { + expired = nowMillis - session.lastAccessedAtMillis >= idleTimeoutMillis; + } + if (expired && closeSession(session.id)) { + closed += 1; + } + } + return closed; + } + + public Object defaultResultValue(ResultSet rs, int index, int sqlType) throws SQLException { + Object value; + switch (sqlType) { + case Types.BIGINT: + value = rs.getLong(index); + break; + case Types.INTEGER: + case Types.SMALLINT: + case Types.TINYINT: + value = rs.getInt(index); + break; + case Types.FLOAT: + case Types.REAL: + value = rs.getFloat(index); + break; + case Types.DOUBLE: + value = rs.getDouble(index); + break; + case Types.DECIMAL: + case Types.NUMERIC: + value = rs.getBigDecimal(index); + break; + case Types.BOOLEAN: + case Types.BIT: + value = rs.getBoolean(index); + break; + case Types.CHAR: + case Types.VARCHAR: + case Types.LONGVARCHAR: + case Types.NCHAR: + case Types.NVARCHAR: + case Types.LONGNVARCHAR: + case Types.CLOB: + case Types.NCLOB: + value = rs.getString(index); + break; + case Types.BINARY: + case Types.VARBINARY: + case Types.LONGVARBINARY: + case Types.BLOB: + value = bytesToHex(rs.getBytes(index)); + break; + case Types.SQLXML: + value = sqlXmlToString(rs.getSQLXML(index)); + break; + default: + value = normalizeResultValue(rs.getObject(index)); + break; + } + return rs.wasNull() ? null : value; + } + + public static Object stringResultValue(ResultSet rs, int index, int sqlType) throws SQLException { + Object value; + switch (sqlType) { + case Types.BINARY: + case Types.VARBINARY: + case Types.LONGVARBINARY: + case Types.BLOB: + value = bytesToHex(rs.getBytes(index)); + break; + case Types.SQLXML: + value = sqlXmlToString(rs.getSQLXML(index)); + break; + default: + value = rs.getString(index); + break; + } + return rs.wasNull() ? null : value; + } + + public static Object normalizeResultValue(Object value) throws SQLException { + if (value == null) { + return null; + } + if (value instanceof Clob) { + Clob clob = (Clob) value; + return clob.getSubString(1, Math.toIntExact(clob.length())); + } + if (value instanceof Blob) { + Blob blob = (Blob) value; + return bytesToHex(blob.getBytes(1, Math.toIntExact(blob.length()))); + } + if (value instanceof SQLXML) { + SQLXML sqlxml = (SQLXML) value; + return sqlxml.getString(); + } + if (value instanceof byte[]) { + byte[] bytes = (byte[]) value; + return bytesToHex(bytes); + } + return value instanceof Number || value instanceof Boolean ? value : value.toString(); + } + + public static String bytesToHex(byte[] bytes) { + if (bytes == null) { + return null; + } + StringBuilder result = new StringBuilder(bytes.length * 2 + 2); + result.append("0x"); + for (byte b : bytes) { + result.append(Character.forDigit((b >> 4) & 0xF, 16)); + result.append(Character.forDigit(b & 0xF, 16)); + } + return result.toString(); + } + + private static String sqlXmlToString(SQLXML value) throws SQLException { + return value == null ? null : value.getString(); + } + + private QueryPageResult readSessionPage(QuerySession session, int pageSize, long executionTimeMs) { + return unchecked(() -> { + session.lastAccessedAtMillis = System.currentTimeMillis(); + int effectivePageSize = Math.max(pageSize, 1); + List> rows = new ArrayList<>(); + + if (session.pendingRow != null) { + rows.add(session.pendingRow); + session.pendingRow = null; + } + + while (rows.size() < effectivePageSize && session.rowsRead < session.maxRows) { + if (!session.resultSet.next()) { + closeSession(session.id); + return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, false, null, false); + } + rows.add(rowValues(session.resultSet, session.valueReader, session.typeNameByIndex)); + session.rowsRead += 1; + } + + if (session.rowsRead >= session.maxRows) { + boolean truncated = session.resultSet.next(); + closeSession(session.id); + return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, truncated, null, false); + } + + boolean hasMore = session.resultSet.next(); + if (!hasMore) { + closeSession(session.id); + return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, false, null, false); + } + + session.pendingRow = rowValues(session.resultSet, session.valueReader, session.typeNameByIndex); + session.rowsRead += 1; + return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, false, session.id, true); + }); + } + + private boolean closeSession(String sessionId) { + QuerySession session = sessions.remove(sessionId); + if (session == null) { + return false; + } + synchronized (session) { + try { + session.resultSet.close(); + } catch (Exception ignored) { + } + try { + session.statement.close(); + } catch (Exception ignored) { + } + } + return true; + } + + static String safeColumnTypeName(ResultSetMetaData meta, int columnIndex) { + try { + String name = meta.getColumnTypeName(columnIndex); + return name == null ? "" : name; + } catch (SQLException ignored) { + return ""; + } + } + + private List rowValues(ResultSet rs, ResultValueReader valueReader, String[] typeNameByIndex) throws SQLException { + ResultSetMetaData meta = rs.getMetaData(); + int colCount = meta.getColumnCount(); + List row = new ArrayList<>(colCount); + for (int i = 1; i <= colCount; i++) { + int sqlType = meta.getColumnType(i); + Object value; + if (valueReader instanceof ColumnAwareResultValueReader) { + String typeName = typeNameByIndex != null && i - 1 < typeNameByIndex.length + ? typeNameByIndex[i - 1] + : safeColumnTypeName(meta, i); + value = ((ColumnAwareResultValueReader) valueReader).read(rs, i, sqlType, typeName); + } else { + value = valueReader.read(rs, i, sqlType); + } + row.add(value); + } + return row; + } + + private void applySchema(Connection conn, String schema, Function setSchemaSql) throws SQLException { + if (schema == null || schema.trim().isEmpty()) { + return; + } + // Prefer JDBC standard APIs over database-specific SQL. + // setSchema (JDBC 4.1) and setCatalog (JDBC 1.0) work universally + // across all JDBC drivers without needing to know the database dialect. + try { + conn.setSchema(schema); + return; + } catch (SQLException | AbstractMethodError ignored) { + // setSchema not supported by this driver + } + try { + conn.setCatalog(schema); + return; + } catch (SQLException | AbstractMethodError ignored) { + // setCatalog not supported either + } + // Fallback: execute database-specific SQL (e.g. USE, SET SCHEMA, etc.) + String schemaSql = setSchemaSql.apply(schema); + if (schemaSql == null || schemaSql.trim().isEmpty()) { + return; + } + try (Statement stmt = conn.createStatement()) { + stmt.execute(schemaSql); + } + } + + private static void applyQueryTimeout(Statement stmt, int timeoutSecs) throws SQLException { + if (timeoutSecs > 0) { + stmt.setQueryTimeout(timeoutSecs); + } + } + + private QueryResult emptyQueryResult(long start) { + return new QueryResult( + Collections.emptyList(), + Collections.emptyList(), + 0L, + System.currentTimeMillis() - start + ); + } + + private QueryPageResult emptyQueryPageResult(long start) { + return new QueryPageResult( + Collections.emptyList(), + Collections.emptyList(), + 0L, + System.currentTimeMillis() - start + ); + } + + static String trimSql(String sql) { + String trimmed = stripTrailingSlashDelimiter(sql.trim()); + if (isPlSqlBlock(trimmed)) { + return trimmed; + } + while (trimmed.endsWith(";")) { + trimmed = trimmed.substring(0, trimmed.length() - 1).trim(); + } + return trimmed; + } + + private static String stripTrailingSlashDelimiter(String sql) { + String trimmed = sql.trim(); + if (!trimmed.endsWith("/")) { + return trimmed; + } + int slashStart = trimmed.length() - 1; + int lineStart = trimmed.lastIndexOf('\n', slashStart - 1) + 1; + if (!trimmed.substring(lineStart, slashStart).trim().isEmpty()) { + return trimmed; + } + String beforeSlash = trimmed.substring(0, lineStart).trim(); + return isPlSqlBlock(beforeSlash) ? beforeSlash : trimmed; + } + + private static boolean isPlSqlBlock(String sql) { + String upperSql = sql.toUpperCase(Locale.ROOT).trim(); + if (!upperSql.startsWith("DECLARE") && !upperSql.startsWith("BEGIN")) { + return false; + } + return upperSql.matches("(?s).*\\bEND\\s+(?!IF\\b|LOOP\\b|CASE\\b)[A-Z0-9_$#]+\\s*;\\s*$") + || upperSql.matches("(?s).*\\bEND\\s*;\\s*$"); + } + + private static T unchecked(ThrowingSupplier supplier) { + try { + return supplier.get(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @FunctionalInterface + public interface ResultValueReader { + Object read(ResultSet rs, int index, int sqlType) throws SQLException; + } + + /** + * Optional extension of {@link ResultValueReader} that exposes the JDBC + * {@code getColumnTypeName} alongside the SQL type code, allowing per-driver + * agents to convert vendor-specific column types (e.g. PostGIS + * {@code geometry}) without re-querying the metadata. + */ + public interface ColumnAwareResultValueReader extends ResultValueReader { + Object read(ResultSet rs, int index, int sqlType, String columnTypeName) throws SQLException; + + @Override + default Object read(ResultSet rs, int index, int sqlType) throws SQLException { + return read(rs, index, sqlType, safeColumnTypeName(rs.getMetaData(), index)); + } + } + + private interface ThrowingSupplier { + T get() throws Exception; + } + + private static final class QuerySession { + private final String id; + private final Statement statement; + private final ResultSet resultSet; + private final List columns; + private final List columnTypes; + private final String[] typeNameByIndex; + private final int maxRows; + private final ResultValueReader valueReader; + private int rowsRead; + private List pendingRow; + private long lastAccessedAtMillis; + + private QuerySession( + String id, + Statement statement, + ResultSet resultSet, + List columns, + List columnTypes, + String[] typeNameByIndex, + int maxRows, + ResultValueReader valueReader + ) { + this.id = id; + this.statement = statement; + this.resultSet = resultSet; + this.columns = columns; + this.columnTypes = columnTypes; + this.typeNameByIndex = typeNameByIndex; + this.maxRows = maxRows; + this.valueReader = valueReader; + this.lastAccessedAtMillis = System.currentTimeMillis(); + } + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcIdentifiers.java b/agents/common/src/main/java/com/dbx/agent/JdbcIdentifiers.java new file mode 100644 index 0000000000..92f34ecf5f --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/JdbcIdentifiers.java @@ -0,0 +1,16 @@ +package com.dbx.agent; + +public final class JdbcIdentifiers { + public static final JdbcIdentifiers INSTANCE = new JdbcIdentifiers(); + + private JdbcIdentifiers() { + } + + public String doubleQuote(String identifier) { + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } + + public String backtick(String identifier) { + return "`" + identifier.replace("`", "``") + "`"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java new file mode 100644 index 0000000000..7c0da9e9b4 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java @@ -0,0 +1,301 @@ +package com.dbx.agent; + +import com.google.gson.Gson; +import com.google.gson.JsonElement; +import com.google.gson.JsonNull; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import com.google.gson.reflect.TypeToken; + +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.lang.reflect.Type; +import java.sql.Connection; +import java.util.Collections; +import java.util.List; + +public final class JsonRpcServer { + private static final long CONNECTION_VALIDATION_INTERVAL_MILLIS = 5_000L; + + private final DatabaseAgent agent; + private final Gson gson = new Gson(); + private ConnectParams lastConnectParams; + private long lastConnectionValidationTimeMillis; + + public JsonRpcServer(DatabaseAgent agent) { + this.agent = agent; + } + + public void run() { + System.out.println("{\"ready\":true}"); + System.out.flush(); + + try { + BufferedReader reader = new BufferedReader(new InputStreamReader(System.in)); + while (true) { + String line = reader.readLine(); + if (line == null) { + break; + } + String response = handleRequest(line); + System.out.println(response); + System.out.flush(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + String handleRequest(String line) { + JsonObject req = JsonParser.parseString(line).getAsJsonObject(); + JsonElement id = req.get("id"); + String method = req.get("method").getAsString(); + JsonObject params = paramsObject(req); + + JsonObject response = new JsonObject(); + response.addProperty("jsonrpc", "2.0"); + response.add("id", id); + + try { + Object result = dispatch(method, params); + response.add("result", gson.toJsonTree(result)); + return gson.toJson(response); + } catch (Throwable e) { + JsonObject error = new JsonObject(); + error.addProperty("code", -1); + error.addProperty("message", e.getMessage() == null ? e.toString() : e.getMessage()); + response.add("error", error); + return gson.toJson(response); + } + } + + private Object dispatch(String method, JsonObject params) throws Exception { + if (AgentProtocol.METHOD_HANDSHAKE.equals(method)) { + return AgentProtocol.handshakeResult(); + } + if (AgentProtocol.METHOD_CONNECT.equals(method)) { + lastConnectParams = gson.fromJson(params, ConnectParams.class); + agent.connect(lastConnectParams); + lastConnectionValidationTimeMillis = 0L; + return Collections.singletonMap("ok", true); + } + if (AgentProtocol.METHOD_TEST_CONNECTION.equals(method)) { + if (!agent.testConnection(gson.fromJson(params, ConnectParams.class))) { + throw new RuntimeException("Connection failed"); + } + return Collections.singletonMap("ok", true); + } + if (AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method)) { + Connection conn = agent.getConnection(); + boolean valid = false; + if (conn != null) { + try { + valid = !conn.isClosed() && conn.isValid(2); + } catch (Exception ignored) { + } + } + if (!valid) { + throw new IllegalStateException("Connection is not valid"); + } + return Collections.singletonMap("ok", true); + } + ensureLiveConnection(method); + if (AgentProtocol.METHOD_LIST_DATABASES.equals(method)) { + return agent.listDatabases(); + } + if (AgentProtocol.METHOD_LIST_SCHEMAS.equals(method)) { + switchCatalog(params); + return agent.listSchemas(); + } + if (AgentProtocol.METHOD_LIST_TABLES.equals(method)) { + switchCatalog(params); + return agent.listTables(params.get("schema").getAsString()); + } + if (AgentProtocol.METHOD_LIST_OBJECTS.equals(method)) { + switchCatalog(params); + return agent.listObjects(params.get("schema").getAsString()); + } + if (AgentProtocol.METHOD_GET_OBJECT_SOURCE.equals(method)) { + switchCatalog(params); + return agent.getObjectSource( + params.get("schema").getAsString(), + params.get("name").getAsString(), + params.get("object_type").getAsString() + ); + } + if (AgentProtocol.METHOD_GET_TABLE_DDL.equals(method)) { + switchCatalog(params); + return agent.getTableDdl(params.get("schema").getAsString(), params.get("table").getAsString()); + } + if (AgentProtocol.METHOD_GET_COLUMNS.equals(method)) { + switchCatalog(params); + return agent.getColumns(params.get("schema").getAsString(), params.get("table").getAsString()); + } + if (AgentProtocol.METHOD_LIST_INDEXES.equals(method)) { + switchCatalog(params); + return agent.listIndexes(params.get("schema").getAsString(), params.get("table").getAsString()); + } + if (AgentProtocol.METHOD_LIST_FOREIGN_KEYS.equals(method)) { + switchCatalog(params); + return agent.listForeignKeys(params.get("schema").getAsString(), params.get("table").getAsString()); + } + if (AgentProtocol.METHOD_LIST_TRIGGERS.equals(method)) { + switchCatalog(params); + return agent.listTriggers(params.get("schema").getAsString(), params.get("table").getAsString()); + } + if (AgentProtocol.METHOD_EXECUTE_QUERY.equals(method)) { + return agent.executeQuery( + params.get("sql").getAsString(), + stringOrNull(params, "schema"), + new ExecuteQueryOptions( + intOrDefault(params, "maxRows", JdbcExecutor.DEFAULT_MAX_ROWS), + intOrNull(params, "fetchSize"), + intOrDefault(params, "timeoutSecs", 0) + ) + ); + } + if (AgentProtocol.METHOD_EXECUTE_QUERY_PAGE.equals(method)) { + return agent.executeQueryPage( + params.get("sql").getAsString(), + stringOrNull(params, "schema"), + new QueryPageOptions( + intOrDefault(params, "pageSize", 100), + intOrNull(params, "fetchSize"), + intOrDefault(params, "maxRows", JdbcExecutor.DEFAULT_MAX_ROWS), + intOrDefault(params, "timeoutSecs", 0) + ) + ); + } + if (AgentProtocol.METHOD_FETCH_QUERY_PAGE.equals(method)) { + return agent.fetchQueryPage( + params.get("sessionId").getAsString(), + intOrDefault(params, "pageSize", 100) + ); + } + if (AgentProtocol.METHOD_CLOSE_QUERY_SESSION.equals(method)) { + return agent.closeQuerySession(params.get("sessionId").getAsString()); + } + if (AgentProtocol.METHOD_GET_EXPLAIN_INFO.equals(method)) { + String plan = agent.getExplainInfo( + params.get("sql").getAsString(), + stringOrNull(params, "database"), + stringOrNull(params, "schema"), + intOrDefault(params, "timeoutSecs", -1), + stringOrNull(params, "mode") + ); + java.util.HashMap result = new java.util.HashMap<>(); + result.put("plan", plan); + result.put("has_actual_stats", "autotrace".equals(stringOrNull(params, "mode"))); + return result; + } + if (AgentProtocol.METHOD_EXECUTE_TRANSACTION.equals(method)) { + Type statementsType = new TypeToken>() {}.getType(); + List statements = gson.fromJson(params.get("statements"), statementsType); + return agent.executeTransaction(statements, stringOrNull(params, "schema")); + } + if (AgentProtocol.METHOD_DISCONNECT.equals(method)) { + JdbcExecutor.INSTANCE.closeAllQuerySessions(); + agent.disconnect(); + lastConnectParams = null; + return Collections.singletonMap("ok", true); + } + if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) { + JdbcExecutor.INSTANCE.closeAllQuerySessions(); + agent.disconnect(); + lastConnectParams = null; + System.exit(0); + return Collections.singletonMap("ok", true); + } + throw new IllegalArgumentException("Unknown method: " + method); + } + + private void switchCatalog(JsonObject params) throws Exception { + String database = stringOrNull(params, "database"); + if (database != null && !database.trim().isEmpty() && agent.getConnection() != null) { + String currentCatalog = null; + try { + currentCatalog = agent.getConnection().getCatalog(); + } catch (Exception ignored) { + } + if (currentCatalog == null || !currentCatalog.trim().equalsIgnoreCase(database.trim())) { + try { + agent.getConnection().setCatalog(database); + } catch (Exception ignored) { + } + } + } + } + + private void ensureLiveConnection(String method) { + if (lastConnectParams == null || !shouldValidateConnection(method)) { + return; + } + Connection conn = agent.getConnection(); + if (conn == null) { + return; + } + long now = System.currentTimeMillis(); + if (lastConnectionValidationTimeMillis > 0 + && now - lastConnectionValidationTimeMillis < CONNECTION_VALIDATION_INTERVAL_MILLIS) { + return; + } + boolean valid = false; + try { + valid = !conn.isClosed() && conn.isValid(2); + } catch (Exception ignored) { + } + if (valid) { + lastConnectionValidationTimeMillis = now; + return; + } + + JdbcExecutor.INSTANCE.closeAllQuerySessions(); + try { + agent.disconnect(); + } catch (Exception ignored) { + } + agent.connect(lastConnectParams); + lastConnectionValidationTimeMillis = System.currentTimeMillis(); + } + + private static boolean shouldValidateConnection(String method) { + return !AgentProtocol.METHOD_HANDSHAKE.equals(method) + && !AgentProtocol.METHOD_CONNECT.equals(method) + && !AgentProtocol.METHOD_TEST_CONNECTION.equals(method) + && !AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method) + && !AgentProtocol.METHOD_FETCH_QUERY_PAGE.equals(method) + && !AgentProtocol.METHOD_CLOSE_QUERY_SESSION.equals(method) + && !AgentProtocol.METHOD_DISCONNECT.equals(method) + && !AgentProtocol.METHOD_SHUTDOWN.equals(method); + } + + private static JsonObject paramsObject(JsonObject req) { + JsonElement params = req.get("params"); + if (params == null || params instanceof JsonNull || !params.isJsonObject()) { + return new JsonObject(); + } + return params.getAsJsonObject(); + } + + private static String stringOrNull(JsonObject object, String key) { + JsonElement element = object.get(key); + if (element == null || element instanceof JsonNull) { + return null; + } + return element.getAsString(); + } + + private static Integer intOrNull(JsonObject object, String key) { + JsonElement element = object.get(key); + if (element == null || element instanceof JsonNull) { + return null; + } + return element.getAsInt(); + } + + private static int intOrDefault(JsonObject object, String key, int defaultValue) { + Integer value = intOrNull(object, key); + return value == null ? defaultValue : value; + } + +} diff --git a/agents/common/src/main/java/com/dbx/agent/ObjectInfo.java b/agents/common/src/main/java/com/dbx/agent/ObjectInfo.java new file mode 100644 index 0000000000..00b26392b2 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ObjectInfo.java @@ -0,0 +1,82 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class ObjectInfo { + private String name; + private String object_type; + private String schema; + private String comment; + + public ObjectInfo() { + this("", "", null, null); + } + + public ObjectInfo(String name, String object_type) { + this(name, object_type, null, null); + } + + public ObjectInfo(String name, String object_type, String schema, String comment) { + this.name = name; + this.object_type = object_type; + this.schema = schema; + this.comment = comment; + } + + public String getName() { + return name; + } + + public String getObject_type() { + return object_type; + } + + public String getSchema() { + return schema; + } + + public String getComment() { + return comment; + } + + public void setName(String name) { + this.name = name; + } + + public void setObject_type(String object_type) { + this.object_type = object_type; + } + + public void setSchema(String schema) { + this.schema = schema; + } + + public void setComment(String comment) { + this.comment = comment; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ObjectInfo)) return false; + ObjectInfo that = (ObjectInfo) other; + return Objects.equals(name, that.name) + && Objects.equals(object_type, that.object_type) + && Objects.equals(schema, that.schema) + && Objects.equals(comment, that.comment); + } + + @Override + public int hashCode() { + return Objects.hash(name, object_type, schema, comment); + } + + @Override + public String toString() { + return "ObjectInfo(name=" + name + + ", object_type=" + object_type + + ", schema=" + schema + + ", comment=" + comment + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/ObjectSource.java b/agents/common/src/main/java/com/dbx/agent/ObjectSource.java new file mode 100644 index 0000000000..8ada3b610d --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/ObjectSource.java @@ -0,0 +1,82 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class ObjectSource { + private String name; + private String object_type; + private String schema; + private String source; + + public ObjectSource() { + this("", "", null, ""); + } + + public ObjectSource(String name, String object_type, String source) { + this(name, object_type, null, source); + } + + public ObjectSource(String name, String object_type, String schema, String source) { + this.name = name; + this.object_type = object_type; + this.schema = schema; + this.source = source; + } + + public String getName() { + return name; + } + + public String getObject_type() { + return object_type; + } + + public String getSchema() { + return schema; + } + + public String getSource() { + return source; + } + + public void setName(String name) { + this.name = name; + } + + public void setObject_type(String object_type) { + this.object_type = object_type; + } + + public void setSchema(String schema) { + this.schema = schema; + } + + public void setSource(String source) { + this.source = source; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof ObjectSource)) return false; + ObjectSource that = (ObjectSource) other; + return Objects.equals(name, that.name) + && Objects.equals(object_type, that.object_type) + && Objects.equals(schema, that.schema) + && Objects.equals(source, that.source); + } + + @Override + public int hashCode() { + return Objects.hash(name, object_type, schema, source); + } + + @Override + public String toString() { + return "ObjectSource(name=" + name + + ", object_type=" + object_type + + ", schema=" + schema + + ", source=" + source + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java new file mode 100644 index 0000000000..43d6e37f51 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java @@ -0,0 +1,499 @@ +package com.dbx.agent; + +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public abstract class PostgresLikeAgent extends AbstractJdbcAgent { + private final PostgresLikeAgentProfile profile; + private boolean mysqlCompatMode = false; + + protected PostgresLikeAgent(PostgresLikeAgentProfile profile) { + this.profile = profile; + } + + public PostgresLikeAgentProfile getProfile() { + return profile; + } + + public void setMysqlCompatMode(boolean mysqlCompatMode) { + this.mysqlCompatMode = mysqlCompatMode; + } + + public boolean isMysqlCompatMode() { + return mysqlCompatMode; + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + geometryAwareResolver() + ); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + return JdbcExecutor.INSTANCE.executePage( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options, + geometryAwareResolver() + ); + } + + /** + * Wrap {@link AbstractJdbcAgent#resultValue} so PostGIS-style {@code geometry} + * and {@code geography} columns are decoded into WKT (matching the native + * tokio_postgres path in {@code crates/dbx-core/src/db/postgres.rs}). + */ + private JdbcExecutor.ColumnAwareResultValueReader geometryAwareResolver() { + return (rs, index, sqlType, columnTypeName) -> { + if (isPostgisGeometryTypeName(columnTypeName)) { + Object raw = rs.getObject(index); + if (rs.wasNull() || raw == null) { + return null; + } + return EwkbWktDecoder.decode(raw); + } + return resultValue(rs, index, sqlType, columnTypeName); + }; + } + + protected Object resultValue(ResultSet rs, int index, int sqlType, String columnTypeName) { + return resultValue(rs, index, sqlType); + } + + static boolean isPostgisGeometryTypeName(String columnTypeName) { + if (columnTypeName == null) { + return false; + } + String trimmed = columnTypeName.trim().toLowerCase(Locale.ROOT); + if (trimmed.isEmpty()) { + return false; + } + // Strip optional schema prefix (e.g. "public.geometry") or "(srid,type)" suffix. + int dot = trimmed.lastIndexOf('.'); + if (dot >= 0 && dot < trimmed.length() - 1) { + trimmed = trimmed.substring(dot + 1); + } + int paren = trimmed.indexOf('('); + if (paren >= 0) { + trimmed = trimmed.substring(0, paren); + } + return "geometry".equals(trimmed) || "geography".equals(trimmed); + } + + private String quoteIdentifier(String identifier) { + return mysqlCompatMode + ? JdbcIdentifiers.INSTANCE.backtick(identifier) + : JdbcIdentifiers.INSTANCE.doubleQuote(identifier); + } + + @Override + protected String driverClass() { + return profile.getDriverClass(); + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return profile.buildUrl(params); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname"); + ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString("datname"))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement( + "SELECT n.nspname AS schema_name " + + "FROM pg_catalog.pg_namespace n " + + "WHERE n.nspname NOT IN ('pg_catalog','information_schema','pg_toast') " + + "AND n.nspname NOT LIKE 'pg_toast_temp_%' " + + "AND n.nspname NOT LIKE 'pg_temp_%' " + + "ORDER BY n.nspname" + ); ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(rs.getString("schema_name")); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement( + "SELECT c.relname AS table_name, " + + "CASE c.relkind " + + "WHEN 'r' THEN 'TABLE' " + + "WHEN 'p' THEN 'TABLE' " + + "WHEN 'v' THEN 'VIEW' " + + "WHEN 'm' THEN 'MATERIALIZED VIEW' " + + "WHEN 'f' THEN 'FOREIGN TABLE' " + + "ELSE 'TABLE' END AS table_type, " + + "obj_description(c.oid) AS table_comment " + + "FROM pg_catalog.pg_class c " + + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + + "WHERE n.nspname = ? AND c.relkind IN ('r','p','v','m','f') " + + "ORDER BY c.relname" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String tableType = normalizeTableType(rs.getString("table_type")); + result.add(new TableInfo(rs.getString("table_name"), tableType, rs.getString("table_comment"))); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement( + "SELECT p.proname AS routine_name, " + + "CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS routine_type " + + "FROM pg_catalog.pg_proc p " + + "JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace " + + "WHERE n.nspname = ? AND p.prokind IN ('p','f') " + + "ORDER BY p.proname" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement( + "SELECT c.relname AS sequence_name, 'SEQUENCE' AS object_type " + + "FROM pg_catalog.pg_class c " + + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + + "WHERE n.nspname = ? AND c.relkind = 'S' " + + "ORDER BY c.relname" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public String getTableDdl(String schema, String table) { + List indexes; + try { + indexes = listIndexes(schema, table); + } catch (RuntimeException e) { + indexes = Collections.emptyList(); + } + + List foreignKeys; + try { + foreignKeys = listForeignKeys(schema, table); + } catch (RuntimeException e) { + foreignKeys = Collections.emptyList(); + } + + return DdlBuilder.buildTableDdl(schema, table, getColumns(schema, table), indexes, foreignKeys, mysqlCompatMode, true); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String upperType = objectType.toUpperCase(); + String sql; + if ("VIEW".equals(upperType) || "MATERIALIZED VIEW".equals(upperType)) { + sql = "SELECT pg_get_viewdef(to_regclass(?), true)"; + } else if ("FUNCTION".equals(upperType)) { + sql = "SELECT pg_get_functiondef(p.oid)\n" + + "FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace\n" + + "WHERE n.nspname = ? AND p.proname = ? AND p.prokind = 'f'\n" + + "ORDER BY p.oid LIMIT 1"; + } else if ("PROCEDURE".equals(upperType)) { + sql = "SELECT pg_get_functiondef(p.oid)\n" + + "FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace\n" + + "WHERE n.nspname = ? AND p.proname = ? AND p.prokind = 'p'\n" + + "ORDER BY p.oid LIMIT 1"; + } else { + throw new IllegalArgumentException("Unsupported object type: " + objectType); + } + + String source; + if ("VIEW".equals(upperType) || "MATERIALIZED VIEW".equals(upperType)) { + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, quoteQualifiedIdentifier(schema, name)); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + } else { + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, name); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = primaryKeys(schema, table); + List result = new ArrayList<>(); + String sql = "SELECT a.attname AS column_name, " + + "format_type(a.atttypid, a.atttypmod) AS data_type, " + + "NOT a.attnotnull AS is_nullable, " + + "pg_get_expr(ad.adbin, ad.adrelid) AS column_default, " + + "col_description(a.attrelid, a.attnum) AS column_comment, " + + "CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 " + + "THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, " + + "CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 " + + "THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, " + + "CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 " + + "THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length " + + "FROM pg_catalog.pg_attribute a " + + "JOIN pg_catalog.pg_type t ON t.oid = a.atttypid " + + "JOIN pg_catalog.pg_class c ON c.oid = a.attrelid " + + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + + "LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum " + + "WHERE n.nspname = ? AND c.relname = ? " + + "AND a.attnum > 0 AND NOT a.attisdropped " + + "ORDER BY a.attnum"; + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String colName = rs.getString("column_name"); + result.add(new ColumnInfo( + colName, + rs.getString("data_type"), + "YES".equals(rs.getString("is_nullable")), + rs.getString("column_default"), + primaryKeys.contains(colName), + null, + rs.getString("column_comment"), + intObject(rs, "numeric_precision"), + intObject(rs, "numeric_scale"), + intObject(rs, "character_maximum_length") + )); + } + } + } + return result; + }); + } + + private Set primaryKeys(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = new LinkedHashSet<>(); + String sql = "SELECT a.attname AS column_name " + + "FROM pg_catalog.pg_constraint co " + + "JOIN pg_catalog.pg_class c ON c.oid = co.conrelid " + + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + + "JOIN LATERAL (SELECT unnest(co.conkey) AS attnum, generate_series(1, array_length(co.conkey, 1)) AS ord) AS pk_cols ON true " + + "JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum = pk_cols.attnum " + + "WHERE co.contype = 'p' " + + "AND n.nspname = ? " + + "AND c.relname = ? " + + "ORDER BY pk_cols.ord"; + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + primaryKeys.add(rs.getString("column_name")); + } + } + } + return primaryKeys; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT i.relname AS index_name, am.amname AS index_type, " + + "ix.indisunique AS is_unique, ix.indisprimary AS is_primary, " + + "array_agg(a.attname ORDER BY k.n) AS columns " + + "FROM pg_index ix " + + "JOIN pg_class t ON t.oid = ix.indrelid " + + "JOIN pg_class i ON i.oid = ix.indexrelid " + + "JOIN pg_namespace n ON n.oid = t.relnamespace " + + "JOIN pg_am am ON am.oid = i.relam " + + "JOIN LATERAL (SELECT unnest(ix.indkey) AS attnum, generate_series(1, array_length(ix.indkey, 1)) AS n) AS k ON true " + + "JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum " + + "WHERE n.nspname = ? AND t.relname = ? " + + "GROUP BY i.relname, am.amname, ix.indisunique, ix.indisprimary " + + "ORDER BY i.relname"; + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + Object[] columnArray = (Object[]) rs.getArray("columns").getArray(); + List columns = new ArrayList<>(); + for (Object column : columnArray) { + columns.add(String.valueOf(column)); + } + result.add(new IndexInfo( + rs.getString("index_name"), + columns, + rs.getBoolean("is_unique"), + rs.getBoolean("is_primary"), + null, + rs.getString("index_type"), + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT co.conname AS constraint_name, " + + "a.attname AS column_name, " + + "rc.relname AS ref_table, " + + "ra.attname AS ref_column " + + "FROM pg_catalog.pg_constraint co " + + "JOIN pg_catalog.pg_class c ON c.oid = co.conrelid " + + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + + "JOIN pg_catalog.pg_class rc ON rc.oid = co.confrelid " + + "JOIN LATERAL (SELECT unnest(co.conkey) AS attnum, generate_series(1, array_length(co.conkey, 1)) AS ord) AS fk ON true " + + "JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum = fk.attnum " + + "JOIN LATERAL (SELECT unnest(co.confkey) AS attnum, generate_series(1, array_length(co.confkey, 1)) AS ord) AS pk ON pk.ord = fk.ord " + + "JOIN pg_catalog.pg_attribute ra ON ra.attrelid = rc.oid AND ra.attnum = pk.attnum " + + "WHERE co.contype = 'f' " + + "AND n.nspname = ? " + + "AND c.relname = ? " + + "ORDER BY co.conname, fk.ord"; + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("constraint_name"), + rs.getString("column_name"), + rs.getString("ref_table"), + rs.getString("ref_column") + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT tg.tgname AS trigger_name, " + + "trim(trailing ',' FROM (" + + "CASE WHEN (tg.tgtype & 4) <> 0 THEN 'INSERT,' ELSE '' END || " + + "CASE WHEN (tg.tgtype & 8) <> 0 THEN 'DELETE,' ELSE '' END || " + + "CASE WHEN (tg.tgtype & 16) <> 0 THEN 'UPDATE,' ELSE '' END || " + + "CASE WHEN (tg.tgtype & 32) <> 0 THEN 'TRUNCATE,' ELSE '' END" + + ")) AS event_manipulation, " + + "CASE WHEN (tg.tgtype & 2) <> 0 THEN 'BEFORE' ELSE 'AFTER' END AS action_timing " + + "FROM pg_catalog.pg_trigger tg " + + "JOIN pg_catalog.pg_class c ON c.oid = tg.tgrelid " + + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + + "WHERE n.nspname = ? AND c.relname = ? AND NOT tg.tgisinternal " + + "ORDER BY tg.tgname"; + try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo( + rs.getString("trigger_name"), + rs.getString("event_manipulation"), + rs.getString("action_timing") + )); + } + } + } + return result; + }); + } + + @Override + public String setSchemaSQL(String schema) { + return "SET search_path TO " + quoteIdentifier(schema); + } + + private java.sql.Connection requireConnection() { + return requireConnected(); + } + + private static String normalizeTableType(String type) { + if (type == null || type.trim().isEmpty()) return "TABLE"; + if ("BASE TABLE".equals(type)) return "TABLE"; + return type; + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + private String quoteQualifiedIdentifier(String schema, String name) { + return quoteIdentifier(schema) + "." + quoteIdentifier(name); + } + + private static Integer intObject(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgentProfile.java b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgentProfile.java new file mode 100644 index 0000000000..f4ba2a8e97 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgentProfile.java @@ -0,0 +1,30 @@ +package com.dbx.agent; + +public final class PostgresLikeAgentProfile { + private final String driverClass; + private final String urlTemplate; + + public PostgresLikeAgentProfile(String driverClass, String urlTemplate) { + this.driverClass = driverClass; + this.urlTemplate = urlTemplate; + } + + public String getDriverClass() { + return driverClass; + } + + public String getUrlTemplate() { + return urlTemplate; + } + + public String buildUrl(ConnectParams params) { + return new JdbcAgentProfile( + driverClass, + urlTemplate, + 0, + false, + java.util.Collections.emptySet(), + java.util.Arrays.asList("TABLE", "VIEW", "MATERIALIZED VIEW", "SYSTEM TABLE", "SYSTEM VIEW") + ).buildUrl(params); + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/QueryPageOptions.java b/agents/common/src/main/java/com/dbx/agent/QueryPageOptions.java new file mode 100644 index 0000000000..293c7aada6 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/QueryPageOptions.java @@ -0,0 +1,82 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class QueryPageOptions { + private int pageSize; + private Integer fetchSize; + private int maxRows; + private int timeoutSecs; + + public QueryPageOptions() { + this(100, null, JdbcExecutor.DEFAULT_MAX_ROWS, 0); + } + + public QueryPageOptions(int pageSize, Integer fetchSize, int maxRows) { + this(pageSize, fetchSize, maxRows, 0); + } + + public QueryPageOptions(int pageSize, Integer fetchSize, int maxRows, int timeoutSecs) { + this.pageSize = pageSize; + this.fetchSize = fetchSize; + this.maxRows = maxRows; + this.timeoutSecs = timeoutSecs; + } + + public int getPageSize() { + return pageSize; + } + + public Integer getFetchSize() { + return fetchSize; + } + + public int getMaxRows() { + return maxRows; + } + + public int getTimeoutSecs() { + return timeoutSecs; + } + + public void setPageSize(int pageSize) { + this.pageSize = pageSize; + } + + public void setFetchSize(Integer fetchSize) { + this.fetchSize = fetchSize; + } + + public void setMaxRows(int maxRows) { + this.maxRows = maxRows; + } + + public void setTimeoutSecs(int timeoutSecs) { + this.timeoutSecs = timeoutSecs; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof QueryPageOptions)) return false; + QueryPageOptions that = (QueryPageOptions) other; + return pageSize == that.pageSize + && maxRows == that.maxRows + && timeoutSecs == that.timeoutSecs + && Objects.equals(fetchSize, that.fetchSize); + } + + @Override + public int hashCode() { + return Objects.hash(pageSize, fetchSize, maxRows, timeoutSecs); + } + + @Override + public String toString() { + return "QueryPageOptions(pageSize=" + pageSize + + ", fetchSize=" + fetchSize + + ", maxRows=" + maxRows + + ", timeoutSecs=" + timeoutSecs + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/QueryPageParams.java b/agents/common/src/main/java/com/dbx/agent/QueryPageParams.java new file mode 100644 index 0000000000..a0ab76b5b8 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/QueryPageParams.java @@ -0,0 +1,90 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class QueryPageParams { + private String sql; + private String schema; + private Integer pageSize; + private Integer fetchSize; + private Integer maxRows; + + public QueryPageParams() { + this("", null, null, null, null); + } + + public QueryPageParams(String sql, String schema, Integer pageSize, Integer fetchSize, Integer maxRows) { + this.sql = sql; + this.schema = schema; + this.pageSize = pageSize; + this.fetchSize = fetchSize; + this.maxRows = maxRows; + } + + public String getSql() { + return sql; + } + + public String getSchema() { + return schema; + } + + public Integer getPageSize() { + return pageSize; + } + + public Integer getFetchSize() { + return fetchSize; + } + + public Integer getMaxRows() { + return maxRows; + } + + public void setSql(String sql) { + this.sql = sql; + } + + public void setSchema(String schema) { + this.schema = schema; + } + + public void setPageSize(Integer pageSize) { + this.pageSize = pageSize; + } + + public void setFetchSize(Integer fetchSize) { + this.fetchSize = fetchSize; + } + + public void setMaxRows(Integer maxRows) { + this.maxRows = maxRows; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof QueryPageParams)) return false; + QueryPageParams that = (QueryPageParams) other; + return Objects.equals(sql, that.sql) + && Objects.equals(schema, that.schema) + && Objects.equals(pageSize, that.pageSize) + && Objects.equals(fetchSize, that.fetchSize) + && Objects.equals(maxRows, that.maxRows); + } + + @Override + public int hashCode() { + return Objects.hash(sql, schema, pageSize, fetchSize, maxRows); + } + + @Override + public String toString() { + return "QueryPageParams(sql=" + sql + + ", schema=" + schema + + ", pageSize=" + pageSize + + ", fetchSize=" + fetchSize + + ", maxRows=" + maxRows + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/QueryPageResult.java b/agents/common/src/main/java/com/dbx/agent/QueryPageResult.java new file mode 100644 index 0000000000..051cb0a88c --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/QueryPageResult.java @@ -0,0 +1,163 @@ +package com.dbx.agent; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +public final class QueryPageResult { + private List columns; + private List column_types; + private List> rows; + private long affected_rows; + private long execution_time_ms; + private boolean truncated; + private String session_id; + private boolean has_more; + + public QueryPageResult() { + this(Collections.emptyList(), Collections.emptyList(), 0L, 0L, false, null, false); + } + + public QueryPageResult(List columns, List> rows, long affected_rows, long execution_time_ms) { + this(columns, rows, affected_rows, execution_time_ms, false, null, false); + } + + public QueryPageResult( + List columns, + List> rows, + long affected_rows, + long execution_time_ms, + boolean truncated + ) { + this(columns, rows, affected_rows, execution_time_ms, truncated, null, false); + } + + public QueryPageResult( + List columns, + List> rows, + long affected_rows, + long execution_time_ms, + boolean truncated, + String session_id, + boolean has_more + ) { + this(columns, null, rows, affected_rows, execution_time_ms, truncated, session_id, has_more); + } + + public QueryPageResult( + List columns, + List column_types, + List> rows, + long affected_rows, + long execution_time_ms, + boolean truncated, + String session_id, + boolean has_more + ) { + this.columns = columns == null ? Collections.emptyList() : columns; + this.column_types = column_types == null ? Collections.emptyList() : column_types; + this.rows = QueryResult.normalizeRows(rows); + this.affected_rows = affected_rows; + this.execution_time_ms = execution_time_ms; + this.truncated = truncated; + this.session_id = session_id; + this.has_more = has_more; + } + + public List getColumns() { + return columns; + } + + public List getColumn_types() { + return column_types; + } + + public List> getRows() { + return rows; + } + + public long getAffected_rows() { + return affected_rows; + } + + public long getExecution_time_ms() { + return execution_time_ms; + } + + public boolean getTruncated() { + return truncated; + } + + public String getSession_id() { + return session_id; + } + + public boolean getHas_more() { + return has_more; + } + + public void setColumns(List columns) { + this.columns = columns; + } + + public void setColumn_types(List column_types) { + this.column_types = column_types; + } + + public void setRows(List> rows) { + this.rows = rows; + } + + public void setAffected_rows(long affected_rows) { + this.affected_rows = affected_rows; + } + + public void setExecution_time_ms(long execution_time_ms) { + this.execution_time_ms = execution_time_ms; + } + + public void setTruncated(boolean truncated) { + this.truncated = truncated; + } + + public void setSession_id(String session_id) { + this.session_id = session_id; + } + + public void setHas_more(boolean has_more) { + this.has_more = has_more; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof QueryPageResult)) return false; + QueryPageResult that = (QueryPageResult) other; + return affected_rows == that.affected_rows + && execution_time_ms == that.execution_time_ms + && truncated == that.truncated + && has_more == that.has_more + && Objects.equals(columns, that.columns) + && Objects.equals(column_types, that.column_types) + && Objects.equals(rows, that.rows) + && Objects.equals(session_id, that.session_id); + } + + @Override + public int hashCode() { + return Objects.hash(columns, column_types, rows, affected_rows, execution_time_ms, truncated, session_id, has_more); + } + + @Override + public String toString() { + return "QueryPageResult(columns=" + columns + + ", column_types=" + column_types + + ", rows=" + rows + + ", affected_rows=" + affected_rows + + ", execution_time_ms=" + execution_time_ms + + ", truncated=" + truncated + + ", session_id=" + session_id + + ", has_more=" + has_more + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/QueryResult.java b/agents/common/src/main/java/com/dbx/agent/QueryResult.java new file mode 100644 index 0000000000..172966e5a6 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/QueryResult.java @@ -0,0 +1,137 @@ +package com.dbx.agent; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +public final class QueryResult { + private List columns; + private List column_types; + private List> rows; + private long affected_rows; + private long execution_time_ms; + private boolean truncated; + + public QueryResult() { + this(Collections.emptyList(), Collections.emptyList(), 0L, 0L, false); + } + + public QueryResult(List columns, List> rows, long affected_rows, long execution_time_ms) { + this(columns, rows, affected_rows, execution_time_ms, false); + } + + public QueryResult( + List columns, + List> rows, + long affected_rows, + long execution_time_ms, + boolean truncated + ) { + this(columns, null, rows, affected_rows, execution_time_ms, truncated); + } + + public QueryResult( + List columns, + List column_types, + List> rows, + long affected_rows, + long execution_time_ms, + boolean truncated + ) { + this.columns = columns == null ? Collections.emptyList() : columns; + this.column_types = column_types == null ? Collections.emptyList() : column_types; + this.rows = normalizeRows(rows); + this.affected_rows = affected_rows; + this.execution_time_ms = execution_time_ms; + this.truncated = truncated; + } + + public List getColumns() { + return columns; + } + + public List getColumn_types() { + return column_types; + } + + public List> getRows() { + return rows; + } + + public long getAffected_rows() { + return affected_rows; + } + + public long getExecution_time_ms() { + return execution_time_ms; + } + + public boolean getTruncated() { + return truncated; + } + + public void setColumns(List columns) { + this.columns = columns; + } + + public void setColumn_types(List column_types) { + this.column_types = column_types; + } + + public void setRows(List> rows) { + this.rows = rows; + } + + public void setAffected_rows(long affected_rows) { + this.affected_rows = affected_rows; + } + + public void setExecution_time_ms(long execution_time_ms) { + this.execution_time_ms = execution_time_ms; + } + + public void setTruncated(boolean truncated) { + this.truncated = truncated; + } + + static List> normalizeRows(List> input) { + if (input == null) { + return Collections.emptyList(); + } + List> normalized = new ArrayList<>(); + for (List row : input) { + normalized.add(row == null ? Collections.emptyList() : new ArrayList(row)); + } + return normalized; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof QueryResult)) return false; + QueryResult that = (QueryResult) other; + return affected_rows == that.affected_rows + && execution_time_ms == that.execution_time_ms + && truncated == that.truncated + && Objects.equals(columns, that.columns) + && Objects.equals(column_types, that.column_types) + && Objects.equals(rows, that.rows); + } + + @Override + public int hashCode() { + return Objects.hash(columns, column_types, rows, affected_rows, execution_time_ms, truncated); + } + + @Override + public String toString() { + return "QueryResult(columns=" + columns + + ", column_types=" + column_types + + ", rows=" + rows + + ", affected_rows=" + affected_rows + + ", execution_time_ms=" + execution_time_ms + + ", truncated=" + truncated + + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/SchemaTableParams.java b/agents/common/src/main/java/com/dbx/agent/SchemaTableParams.java new file mode 100644 index 0000000000..fc389c6646 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/SchemaTableParams.java @@ -0,0 +1,51 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class SchemaTableParams { + private String schema; + private String table; + + public SchemaTableParams() { + this("", ""); + } + + public SchemaTableParams(String schema, String table) { + this.schema = schema; + this.table = table; + } + + public String getSchema() { + return schema; + } + + public String getTable() { + return table; + } + + public void setSchema(String schema) { + this.schema = schema; + } + + public void setTable(String table) { + this.table = table; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof SchemaTableParams)) return false; + SchemaTableParams that = (SchemaTableParams) other; + return Objects.equals(schema, that.schema) && Objects.equals(table, that.table); + } + + @Override + public int hashCode() { + return Objects.hash(schema, table); + } + + @Override + public String toString() { + return "SchemaTableParams(schema=" + schema + ", table=" + table + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java b/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java new file mode 100644 index 0000000000..7f671b71e0 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java @@ -0,0 +1,350 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class StandardJdbcMetadata { + public static final StandardJdbcMetadata INSTANCE = new StandardJdbcMetadata(); + + private StandardJdbcMetadata() { + } + + public List listDatabases(Connection conn, String configuredDatabase) { + return unchecked(() -> { + Set names = new LinkedHashSet<>(); + try { + try (ResultSet rs = conn.getMetaData().getCatalogs()) { + while (rs.next()) { + addNonBlank(names, rs.getString("TABLE_CAT")); + } + } + } catch (Exception ignored) { + } + addNonBlank(names, configuredDatabase); + addNonBlank(names, conn.getCatalog()); + + List result = new ArrayList<>(); + for (String name : names) { + result.add(new DatabaseInfo(name)); + } + return result; + }); + } + + public List listSchemas(Connection conn, JdbcAgentProfile profile) { + return unchecked(() -> { + Set names = new LinkedHashSet<>(); + DatabaseMetaData meta = conn.getMetaData(); + try { + appendSchemas(names, meta.getSchemas(null, null)); + } catch (Exception | AbstractMethodError first) { + try { + appendSchemas(names, meta.getSchemas()); + } catch (Exception | AbstractMethodError ignored) { + } + } + try { + addNonBlank(names, conn.getSchema()); + } catch (Exception | AbstractMethodError ignored) { + } + + List result = new ArrayList<>(); + for (String name : names) { + if (!profile.getExcludedSchemas().contains(name.toUpperCase(Locale.ROOT))) { + result.add(name); + } + } + Collections.sort(result); + return result; + }); + } + + public List listTables(Connection conn, JdbcAgentProfile profile, String configuredDatabase, String schema) { + return unchecked(() -> { + DatabaseMetaData meta = conn.getMetaData(); + List result = new ArrayList<>(); + appendTables(result, meta, profile, null, schema); + if (result.isEmpty() && profile.getCatalogFallbackEnabled() && !configuredDatabase.trim().isEmpty()) { + appendTables(result, meta, profile, configuredDatabase, schema); + } + result.sort(Comparator.comparing(TableInfo::getName)); + return result; + }); + } + + public List listObjects(List tables, String schema) { + List result = new ArrayList<>(); + for (TableInfo table : tables) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + return result; + } + + public List getColumns(Connection conn, JdbcAgentProfile profile, String configuredDatabase, String schema, String table) { + return unchecked(() -> { + DatabaseMetaData meta = conn.getMetaData(); + Set primaryKeys = primaryKeys(meta, null, schema, table); + List result = new ArrayList<>(); + appendColumns(result, meta, null, schema, table, primaryKeys); + if (result.isEmpty() && profile.getCatalogFallbackEnabled() && !configuredDatabase.trim().isEmpty()) { + Set fallbackPrimaryKeys = primaryKeys(meta, configuredDatabase, schema, table); + appendColumns(result, meta, configuredDatabase, schema, table, fallbackPrimaryKeys); + } + return result; + }); + } + + public List listIndexes(Connection conn, String schema, String table) { + return unchecked(() -> { + DatabaseMetaData meta = conn.getMetaData(); + Map> indexes = new LinkedHashMap<>(); + Map unique = new LinkedHashMap<>(); + try (ResultSet rs = meta.getIndexInfo(null, blankToNull(schema), table, false, false)) { + while (rs.next()) { + String name = rs.getString("INDEX_NAME"); + String column = rs.getString("COLUMN_NAME"); + if (name == null || column == null) { + continue; + } + indexes.computeIfAbsent(name, ignored -> new ArrayList<>()).add( + new IndexColumn(rs.getShort("ORDINAL_POSITION"), column) + ); + unique.put(name, !rs.getBoolean("NON_UNIQUE")); + } + } + + List result = new ArrayList<>(); + for (Map.Entry> entry : indexes.entrySet()) { + List orderedColumns = entry.getValue(); + orderedColumns.sort(Comparator.comparingInt(IndexColumn::getOrdinal)); + List columns = new ArrayList<>(); + for (IndexColumn column : orderedColumns) { + columns.add(column.getName()); + } + String name = entry.getKey(); + result.add(new IndexInfo( + name, + columns, + Boolean.TRUE.equals(unique.get(name)), + "PRIMARY".equalsIgnoreCase(name), + null, + null, + null, + null + )); + } + return result; + }); + } + + public List listForeignKeys(Connection conn, String schema, String table) { + return unchecked(() -> { + DatabaseMetaData meta = conn.getMetaData(); + List result = new ArrayList<>(); + try (ResultSet rs = meta.getImportedKeys(null, blankToNull(schema), table)) { + while (rs.next()) { + String name = rs.getString("FK_NAME"); + result.add(new ForeignKeyInfo( + name == null ? "" : name, + rs.getString("FKCOLUMN_NAME"), + rs.getString("PKTABLE_NAME"), + rs.getString("PKCOLUMN_NAME") + )); + } + } + return result; + }); + } + + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + private static String[] getDriverTableTypes(DatabaseMetaData meta, JdbcAgentProfile profile) throws Exception { + Set configuredTypes = new LinkedHashSet<>(); + for (String type : profile.getTableTypes()) { + if (type != null) { + configuredTypes.add(type.toUpperCase(Locale.ROOT)); + } + } + try (ResultSet rs = meta.getTableTypes()) { + List types = new ArrayList<>(); + while (rs.next()) { + String type = rs.getString("TABLE_TYPE"); + if (type != null && configuredTypes.contains(type.toUpperCase(Locale.ROOT))) { + types.add(type); + } + } + if (!types.isEmpty()) { + return types.toArray(new String[0]); + } + } catch (Exception ignored) { + } + return profile.getTableTypes().toArray(new String[profile.getTableTypes().size()]); + } + + public static String getIdentifierQuote(Connection conn, JdbcAgentProfile profile) { + try { + String quote = conn.getMetaData().getIdentifierQuoteString(); + if (quote != null && !quote.trim().isEmpty()) { + return quote.trim(); + } + } catch (Exception ignored) { + } + return profile.getIdentifierQuote(); + } + + public static String quoteIdentifier(String identifier, String quote) { + return quote + identifier.replace(quote, quote + quote) + quote; + } + + public static String schemaSwitchSql(Connection conn, JdbcAgentProfile profile, String schema) { + if (profile.getSkipExecutionContext()) { + return profile.schemaSwitchSql(schema); + } + String quote = getIdentifierQuote(conn, profile); + return profile.schemaSwitchSql(schema, quote); + } + + private void appendTables(List result, DatabaseMetaData meta, JdbcAgentProfile profile, String catalog, String schema) throws Exception { + String[] tableTypes = getDriverTableTypes(meta, profile); + try (ResultSet rs = meta.getTables(catalog, blankToNull(schema), "%", tableTypes)) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("TABLE_NAME"), + normalizeTableType(rs.getString("TABLE_TYPE")), + rs.getString("REMARKS") + )); + } + } + } + + private Set primaryKeys(DatabaseMetaData meta, String catalog, String schema, String table) { + Set keys = new LinkedHashSet<>(); + try (ResultSet rs = meta.getPrimaryKeys(catalog, blankToNull(schema), table)) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + if (name != null) { + keys.add(name); + } + } + } catch (Exception ignored) { + } + return keys; + } + + private static void appendSchemas(Set names, ResultSet rs) throws Exception { + try (ResultSet closeable = rs) { + while (closeable.next()) { + addNonBlank(names, closeable.getString("TABLE_SCHEM")); + } + } + } + + private void appendColumns( + List result, + DatabaseMetaData meta, + String catalog, + String schema, + String table, + Set primaryKeys + ) throws Exception { + try (ResultSet rs = meta.getColumns(catalog, blankToNull(schema), table, "%")) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + name, + rs.getString("TYPE_NAME"), + rs.getInt("NULLABLE") != DatabaseMetaData.columnNoNulls, + rs.getString("COLUMN_DEF"), + primaryKeys.contains(name), + null, + rs.getString("REMARKS"), + intOrNull(rs, "COLUMN_SIZE"), + intOrNull(rs, "DECIMAL_DIGITS"), + characterLength(rs) + )); + } + } + } + + private static String normalizeTableType(String type) { + if (type == null || type.trim().isEmpty()) { + return "TABLE"; + } + String normalized = type.toUpperCase(Locale.ROOT); + if ("BASE TABLE".equals(normalized) || "COLUMN TABLE".equals(normalized) || "ROW TABLE".equals(normalized)) { + return "TABLE"; + } + return type; + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static Integer characterLength(ResultSet rs) throws Exception { + String typeName = rs.getString("TYPE_NAME"); + if (typeName == null) { + return null; + } + String normalized = typeName.toLowerCase(Locale.ROOT); + if (!normalized.contains("char") && !normalized.contains("text")) { + return null; + } + return intOrNull(rs, "COLUMN_SIZE"); + } + + private static String blankToNull(String value) { + return value == null || value.trim().isEmpty() ? null : value; + } + + private static void addNonBlank(Set values, String value) { + if (value != null && !value.trim().isEmpty()) { + values.add(value); + } + } + + private static T unchecked(ThrowingSupplier supplier) { + try { + return supplier.get(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static final class IndexColumn { + private final int ordinal; + private final String name; + + private IndexColumn(int ordinal, String name) { + this.ordinal = ordinal; + this.name = name; + } + + private int getOrdinal() { + return ordinal; + } + + private String getName() { + return name; + } + } + + private interface ThrowingSupplier { + T get() throws Exception; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/TableInfo.java b/agents/common/src/main/java/com/dbx/agent/TableInfo.java new file mode 100644 index 0000000000..77518afcf2 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/TableInfo.java @@ -0,0 +1,67 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class TableInfo { + private String name; + private String table_type; + private String comment; + + public TableInfo() { + this("", "", null); + } + + public TableInfo(String name, String table_type) { + this(name, table_type, null); + } + + public TableInfo(String name, String table_type, String comment) { + this.name = name; + this.table_type = table_type; + this.comment = comment; + } + + public String getName() { + return name; + } + + public String getTable_type() { + return table_type; + } + + public String getComment() { + return comment; + } + + public void setName(String name) { + this.name = name; + } + + public void setTable_type(String table_type) { + this.table_type = table_type; + } + + public void setComment(String comment) { + this.comment = comment; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof TableInfo)) return false; + TableInfo that = (TableInfo) other; + return Objects.equals(name, that.name) + && Objects.equals(table_type, that.table_type) + && Objects.equals(comment, that.comment); + } + + @Override + public int hashCode() { + return Objects.hash(name, table_type, comment); + } + + @Override + public String toString() { + return "TableInfo(name=" + name + ", table_type=" + table_type + ", comment=" + comment + ")"; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java b/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java new file mode 100644 index 0000000000..9ba7b877d3 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java @@ -0,0 +1,135 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.sql.Statement; +import java.util.Collections; +import java.util.List; +import java.util.function.Function; + +public final class TransactionExecutor { + private TransactionExecutor() { + } + + public static QueryResult executeUpdateStatements( + Connection conn, + List statements, + String schema, + Function setSchemaSql + ) { + return executeStatements(conn, statements, schema, setSchemaSql, new StatementRunner() { + @Override + public long run(Statement stmt, String sql) throws Exception { + return stmt.executeUpdate(sql); + } + }); + } + + public static QueryResult executeStatements( + Connection conn, + List statements, + String schema, + Function setSchemaSql, + StatementRunner runner + ) { + return unchecked(() -> { + long start = System.currentTimeMillis(); + if (!supportsTransactions(conn)) { + long totalAffected = executeAll(conn, statements, schema, setSchemaSql, runner); + return result(totalAffected, start); + } + + boolean savedAutoCommit = conn.getAutoCommit(); + conn.setAutoCommit(false); + try { + long totalAffected = executeAll(conn, statements, schema, setSchemaSql, runner); + conn.commit(); + return result(totalAffected, start); + } catch (Exception e) { + conn.rollback(); + throw e; + } finally { + conn.setAutoCommit(savedAutoCommit); + } + }); + } + + private static long executeAll( + Connection conn, + List statements, + String schema, + Function setSchemaSql, + StatementRunner runner + ) throws Exception { + applySchema(conn, schema, setSchemaSql); + long totalAffected = 0; + for (String statement : statements) { + try (Statement stmt = conn.createStatement()) { + totalAffected += runner.run(stmt, JdbcExecutor.trimSql(statement)); + } + } + return totalAffected; + } + + private static void applySchema(Connection conn, String schema, Function setSchemaSql) throws Exception { + if (schema == null || schema.trim().isEmpty()) { + return; + } + // Prefer JDBC standard APIs over database-specific SQL. + try { + conn.setSchema(schema); + return; + } catch (Exception | AbstractMethodError ignored) { + // setSchema not supported by this driver + } + try { + conn.setCatalog(schema); + return; + } catch (Exception | AbstractMethodError ignored) { + // setCatalog not supported either + } + // Fallback: execute database-specific SQL (e.g. USE, SET SCHEMA, etc.) + String schemaSql = setSchemaSql.apply(schema); + if (schemaSql == null || schemaSql.trim().isEmpty()) { + return; + } + try (Statement stmt = conn.createStatement()) { + stmt.execute(schemaSql); + } + } + + private static boolean supportsTransactions(Connection conn) { + try { + return conn.getMetaData().supportsTransactions(); + } catch (Exception ignored) { + return true; + } + } + + private static QueryResult result(long totalAffected, long start) { + return new QueryResult( + Collections.emptyList(), + Collections.emptyList(), + totalAffected, + System.currentTimeMillis() - start, + false + ); + } + + private static T unchecked(ThrowingSupplier supplier) { + try { + return supplier.get(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + public interface StatementRunner { + long run(Statement stmt, String sql) throws Exception; + } + + private interface ThrowingSupplier { + T get() throws Exception; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/TriggerInfo.java b/agents/common/src/main/java/com/dbx/agent/TriggerInfo.java new file mode 100644 index 0000000000..490b55a187 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/TriggerInfo.java @@ -0,0 +1,63 @@ +package com.dbx.agent; + +import java.util.Objects; + +public final class TriggerInfo { + private String name; + private String event; + private String timing; + + public TriggerInfo() { + this("", "", ""); + } + + public TriggerInfo(String name, String event, String timing) { + this.name = name; + this.event = event; + this.timing = timing; + } + + public String getName() { + return name; + } + + public String getEvent() { + return event; + } + + public String getTiming() { + return timing; + } + + public void setName(String name) { + this.name = name; + } + + public void setEvent(String event) { + this.event = event; + } + + public void setTiming(String timing) { + this.timing = timing; + } + + @Override + public boolean equals(Object other) { + if (this == other) return true; + if (!(other instanceof TriggerInfo)) return false; + TriggerInfo that = (TriggerInfo) other; + return Objects.equals(name, that.name) + && Objects.equals(event, that.event) + && Objects.equals(timing, that.timing); + } + + @Override + public int hashCode() { + return Objects.hash(name, event, timing); + } + + @Override + public String toString() { + return "TriggerInfo(name=" + name + ", event=" + event + ", timing=" + timing + ")"; + } +} diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json new file mode 100644 index 0000000000..12d0356e90 --- /dev/null +++ b/agents/common/src/main/resources/agent-protocol-v1.json @@ -0,0 +1,75 @@ +{ + "protocolVersion": 1, + "handshakeMethod": "handshake", + "handshakeResponseFields": [ + "protocolVersion", + "agentProtocolVersion", + "capabilities" + ], + "allCapabilities": [ + "connect", + "test_connection", + "metadata", + "query", + "paged_query", + "transaction", + "ddl", + "kv" + ], + "capabilities": [ + "connect", + "test_connection", + "metadata", + "query", + "paged_query", + "transaction", + "ddl" + ], + "defaultSqlCapabilities": [ + "connect", + "test_connection", + "metadata", + "query", + "paged_query", + "transaction", + "ddl" + ], + "commonMethods": [ + "handshake", + "connect", + "test_connection", + "validate_connection", + "list_databases", + "list_schemas", + "list_tables", + "list_objects", + "get_object_source", + "get_table_ddl", + "get_columns", + "list_indexes", + "list_foreign_keys", + "list_triggers", + "execute_query", + "execute_query_page", + "fetch_query_page", + "close_query_session", + "get_explain_info", + "execute_transaction", + "disconnect", + "shutdown" + ], + "mongoLegacyMethods": [ + "list_databases", + "list_collections", + "find_documents", + "insert_document", + "update_document", + "delete_document" + ], + "kvMethods": [ + "kv_list_prefix", + "kv_get", + "kv_put", + "kv_delete" + ] +} diff --git a/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java b/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java new file mode 100644 index 0000000000..c6fe1b5710 --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java @@ -0,0 +1,443 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class AbstractJdbcAgentTest { + @Test + void ownsConnectionLifecycleAndConnectedState() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + + assertThrows(IllegalStateException.class, () -> agent.executeQuery("SELECT 1", null, new ExecuteQueryOptions())); + + agent.connect(new ConnectParams("localhost", 0, "demo", "user", "secret", "", "", false)); + + assertNotNull(agent.getConnection()); + assertEquals("demo", agent.database()); + assertEquals(1, tracking.openCount); + assertEquals(1, agent.afterConnectCount); + + agent.disconnect(); + + assertNull(agent.getConnection()); + assertEquals(1, tracking.closeCount); + } + + @Test + void testsConnectionsThroughSharedLifecycle() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + + assertTrue(agent.testConnection(new ConnectParams())); + + assertEquals(1, tracking.openCount); + assertEquals(1, tracking.isValidCount); + assertEquals(1, tracking.closeCount); + } + + @Test + void delegatesQueryExecutionWithSchemaAndValueReader() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + agent.connect(new ConnectParams()); + + QueryResult result = agent.executeQuery("SELECT VALUE", "APP", new ExecuteQueryOptions(25, 7)); + + assertEquals(Collections.singletonList("VALUE"), result.getColumns()); + assertEquals(Collections.singletonList(Collections.singletonList("row-value")), result.getRows()); + assertFalse(result.getTruncated()); + assertEquals(Arrays.asList("setSchema:APP", "setMaxRows:26", "setFetchSize:7", "execute:SELECT VALUE"), tracking.calls); + } + + @Test + void delegatesQueryTimeoutToStatement() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + agent.connect(new ConnectParams()); + + agent.executeQuery("SELECT VALUE", null, new ExecuteQueryOptions(25, 7, 12)); + + assertEquals(Arrays.asList("setMaxRows:26", "setQueryTimeout:12", "setFetchSize:7", "execute:SELECT VALUE"), tracking.calls); + } + + @Test + void delegatesPagedQueryTimeoutToStatement() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + agent.connect(new ConnectParams()); + + agent.executeQueryPage("SELECT VALUE", null, new QueryPageOptions(10, 7, 25, 12)); + + assertEquals(Arrays.asList("setQueryTimeout:12", "setFetchSize:7", "execute:SELECT VALUE"), tracking.calls); + } + + @Test + void preservesPlSqlBlockTerminatorDuringQueryExecution() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + agent.connect(new ConnectParams()); + + agent.executeQuery("BEGIN\n NULL;\nEND;\n/", null, new ExecuteQueryOptions()); + + assertEquals(Arrays.asList("setMaxRows:10001", "execute:BEGIN\n NULL;\nEND;"), tracking.calls); + } + + @Test + void delegatesTransactionsThroughSharedFoundation() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + agent.connect(new ConnectParams()); + + QueryResult result = agent.executeTransaction(Arrays.asList("UPDATE A;", "UPDATE B"), "APP"); + + assertEquals(2L, result.getAffected_rows()); + assertEquals( + Arrays.asList( + "setAutoCommit:false", + "setSchema:APP", + "executeUpdate:UPDATE A", + "executeUpdate:UPDATE B", + "commit", + "setAutoCommit:true" + ), + tracking.calls + ); + } + + @Test + void preservesPlSqlBlockTerminatorDuringTransactionExecution() { + TrackingConnection tracking = new TrackingConnection(); + TestAgent agent = new TestAgent(tracking); + agent.connect(new ConnectParams()); + + agent.executeTransaction(Collections.singletonList("DECLARE\n v NUMBER;\nBEGIN\n NULL;\nEND;"), null); + + assertEquals( + Arrays.asList( + "setAutoCommit:false", + "executeUpdate:DECLARE\n v NUMBER;\nBEGIN\n NULL;\nEND;", + "commit", + "setAutoCommit:true" + ), + tracking.calls + ); + } + + @Test + void inheritedDdlFallbackToleratesMissingOptionalMetadata() { + TestAgent agent = new TestAgent(new TrackingConnection()) { + @Override + public List getColumns(String schema, String table) { + return Collections.singletonList(new ColumnInfo("ID", "INTEGER", false, null, true)); + } + + @Override + public List listIndexes(String schema, String table) { + throw new RuntimeException("indexes unavailable"); + } + + @Override + public List listForeignKeys(String schema, String table) { + throw new RuntimeException("foreign keys unavailable"); + } + }; + + assertEquals( + "CREATE TABLE \"APP\".\"ORDERS\" (\n" + + " \"ID\" INTEGER NOT NULL,\n" + + " PRIMARY KEY (\"ID\")\n" + + ");\n", + agent.getTableDdl("APP", "ORDERS") + ); + } + + private static class TestAgent extends AbstractJdbcAgent { + private final TrackingConnection tracking; + private int afterConnectCount; + + private TestAgent(TrackingConnection tracking) { + this.tracking = tracking; + } + + @Override + protected String driverClass() { + return TestAgent.class.getName(); + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return "jdbc:test"; + } + + @Override + protected Connection openConnection(ConnectParams params) { + tracking.openCount += 1; + return tracking.connection(); + } + + @Override + protected void afterConnect(ConnectParams params, Connection connection) { + afterConnectCount += 1; + } + + @Override + public List listDatabases() { + return Collections.emptyList(); + } + + @Override + public List listSchemas() { + return Collections.emptyList(); + } + + @Override + public List listTables(String schema) { + return Collections.emptyList(); + } + + @Override + public List getColumns(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + schema; + } + + private String database() { + return getConfiguredDatabase(); + } + } + + private static final class TrackingConnection { + private final List calls = new ArrayList(); + private int openCount; + private int closeCount; + private int isValidCount; + private boolean autoCommit = true; + + private Connection connection() { + return proxy(Connection.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("createStatement".equals(name)) { + return statement(); + } + if ("getMetaData".equals(name)) { + return databaseMetaData(); + } + if ("isValid".equals(name)) { + isValidCount += 1; + return true; + } + if ("close".equals(name)) { + closeCount += 1; + return null; + } + if ("getAutoCommit".equals(name)) { + return autoCommit; + } + if ("setSchema".equals(name)) { + calls.add("setSchema:" + args[0]); + return null; + } + if ("setCatalog".equals(name)) { + calls.add("setCatalog:" + args[0]); + return null; + } + if ("setAutoCommit".equals(name)) { + autoCommit = (Boolean) args[0]; + calls.add("setAutoCommit:" + args[0]); + return null; + } + if ("commit".equals(name)) { + calls.add("commit"); + return null; + } + if ("rollback".equals(name)) { + calls.add("rollback"); + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private DatabaseMetaData databaseMetaData() { + return proxy(DatabaseMetaData.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + if ("supportsTransactions".equals(method.getName())) { + return true; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private Statement statement() { + return proxy(Statement.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("execute".equals(name)) { + calls.add("execute:" + args[0]); + return true; + } + if ("getResultSet".equals(name)) { + return resultSet(); + } + if ("getUpdateCount".equals(name)) { + return 1; + } + if ("executeUpdate".equals(name)) { + calls.add("executeUpdate:" + args[0]); + return 1; + } + if ("setMaxRows".equals(name)) { + calls.add("setMaxRows:" + args[0]); + return null; + } + if ("setFetchSize".equals(name)) { + calls.add("setFetchSize:" + args[0]); + return null; + } + if ("setQueryTimeout".equals(name)) { + calls.add("setQueryTimeout:" + args[0]); + return null; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private ResultSet resultSet() { + final int[] index = {-1}; + return proxy(ResultSet.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + index[0] += 1; + return index[0] == 0; + } + if ("getMetaData".equals(name)) { + return resultSetMetaData(); + } + if ("getString".equals(name)) { + return "row-value"; + } + if ("wasNull".equals(name)) { + return false; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private ResultSetMetaData resultSetMetaData() { + return proxy(ResultSetMetaData.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("getColumnCount".equals(name)) { + return 1; + } + if ("getColumnLabel".equals(name)) { + return "VALUE"; + } + if ("getColumnType".equals(name)) { + return Types.VARCHAR; + } + return defaultValue(method.getReturnType()); + } + }); + } + } + + private static T proxy(Class type, final MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java new file mode 100644 index 0000000000..3ec366cd9e --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java @@ -0,0 +1,491 @@ +package com.dbx.agent; + +import com.google.gson.JsonArray; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class CommonJavaCompatibilityTest { + @Test + void definesSharedAgentProtocolContract() { + assertEquals("handshake", AgentProtocol.METHOD_HANDSHAKE); + assertEquals(1, AgentProtocol.PROTOCOL_VERSION); + assertTrue(AgentProtocol.CAPABILITIES.contains(AgentProtocol.CAPABILITY_CONNECT)); + assertTrue(AgentProtocol.CAPABILITIES.contains(AgentProtocol.CAPABILITY_QUERY)); + assertTrue(AgentProtocol.CAPABILITIES.contains(AgentProtocol.CAPABILITY_METADATA)); + } + + @Test + void agentProtocolMatchesContractResource() { + JsonObject contract = protocolContract(); + + assertEquals(AgentProtocol.PROTOCOL_VERSION, contract.get("protocolVersion").getAsInt()); + assertEquals(AgentProtocol.METHOD_HANDSHAKE, contract.get("handshakeMethod").getAsString()); + assertEquals( + Arrays.asList("protocolVersion", "agentProtocolVersion", "capabilities"), + strings(contract.getAsJsonArray("handshakeResponseFields")) + ); + assertEquals(AgentProtocol.ALL_CAPABILITIES, strings(contract.getAsJsonArray("allCapabilities"))); + assertEquals(AgentProtocol.CAPABILITIES, strings(contract.getAsJsonArray("capabilities"))); + assertEquals(AgentProtocol.CAPABILITIES, strings(contract.getAsJsonArray("defaultSqlCapabilities"))); + assertEquals(AgentProtocol.COMMON_METHODS, strings(contract.getAsJsonArray("commonMethods"))); + assertEquals(AgentProtocol.MONGO_LEGACY_METHODS, strings(contract.getAsJsonArray("mongoLegacyMethods"))); + assertEquals(AgentProtocol.KV_METHODS, strings(contract.getAsJsonArray("kvMethods"))); + } + + @Test + void jsonRpcServerExposesProtocolHandshake() { + JsonRpcServer server = new JsonRpcServer(new MinimalAgent()); + + String response = server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"" + AgentProtocol.METHOD_HANDSHAKE + "\",\"params\":{\"appVersion\":\"0.5.13\",\"supportedProtocolVersions\":[1]}}" + ); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + JsonObject result = json.getAsJsonObject("result"); + assertEquals("2.0", json.get("jsonrpc").getAsString()); + assertEquals(7, json.get("id").getAsInt()); + assertEquals(1, result.get("protocolVersion").getAsInt()); + assertEquals(1, result.get("agentProtocolVersion").getAsInt()); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), "connect")); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), "query")); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), "metadata")); + } + + @Test + void jsonRpcServerReconnectsWhenStoredJdbcConnectionIsStale() { + ReconnectingAgent agent = new ReconnectingAgent(); + JsonRpcServer server = new JsonRpcServer(agent); + + server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"connect\",\"params\":{\"host\":\"db.example.com\",\"port\":1521,\"database\":\"ORCL\",\"username\":\"u\",\"password\":\"p\"}}" + ); + String response = server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"list_databases\",\"params\":{}}" + ); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertTrue(json.has("result")); + assertEquals(2, agent.connectCount); + assertEquals(1, agent.disconnectCount); + assertEquals(1, agent.firstConnectionValidChecks); + } + + @Test + void jsonRpcServerValidatesCurrentJdbcConnection() { + ReconnectingAgent agent = new ReconnectingAgent(); + JsonRpcServer server = new JsonRpcServer(agent); + + server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"connect\",\"params\":{\"host\":\"db.example.com\",\"port\":1521,\"database\":\"ORCL\",\"username\":\"u\",\"password\":\"p\"}}" + ); + String response = server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"" + AgentProtocol.METHOD_VALIDATE_CONNECTION + "\",\"params\":{}}" + ); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertTrue(json.has("error")); + assertEquals(1, agent.connectCount); + assertEquals(1, agent.firstConnectionValidChecks); + } + + @Test + void jsonRpcServerSwitchesCatalogBeforeMetadataCalls() { + CatalogSwitchAgent agent = new CatalogSwitchAgent(); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = server.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"" + AgentProtocol.METHOD_LIST_TABLES + "\",\"params\":{\"database\":\"sales\",\"schema\":\"app\"}}" + ); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertTrue(json.has("result")); + assertEquals("sales", agent.catalogs.get(0)); + assertEquals("app", agent.lastSchema); + } + + @Test + void exposesJavaFriendlyDefaultsAndModels() { + ConnectParams params = new ConnectParams("localhost", 5432, "demo", "user", "secret", "ssl=false", "", false); + assertEquals("localhost", params.getHost()); + assertEquals("ssl=false", params.getUrl_params()); + + TableInfo table = new TableInfo("orders", "TABLE"); + assertEquals("TABLE", table.getTable_type()); + + QueryResult result = new QueryResult( + Collections.singletonList("id"), + Collections.singletonList(Collections.singletonList(1)), + 2L, + 3L + ); + assertEquals(2L, result.getAffected_rows()); + assertEquals(3L, result.getExecution_time_ms()); + assertFalse(result.getTruncated()); + + QueryPageResult page = new QueryPageResult( + Collections.singletonList("id"), + Collections.emptyList(), + 0L, + 1L, + false, + "session-1", + true + ); + assertEquals("session-1", page.getSession_id()); + assertEquals(true, page.getHas_more()); + + assertEquals(JdbcExecutor.DEFAULT_MAX_ROWS, new ExecuteQueryOptions().getMaxRows()); + assertEquals(100, new QueryPageOptions().getPageSize()); + assertNotNull(JdbcExecutor.INSTANCE); + } + + @Test + void exposesDatabaseAgentDefaultMethodsToJavaImplementors() { + DatabaseAgent agent = new MinimalAgent(); + + assertEquals(1, agent.listObjects("public").size()); + assertThrows(UnsupportedOperationException.class, () -> + agent.getObjectSource("public", "orders", "TABLE") + ); + assertEquals("SET SCHEMA \"public\"", agent.setSchemaSQL("public")); + assertThrows(IllegalStateException.class, () -> + agent.executeQueryPage("select 1", "public") + ); + assertEquals(0L, agent.executeQuery("select 1", "public").getAffected_rows()); + + String ddl = DatabaseAgent.buildTableDdl( + "public", + "orders", + Collections.singletonList(new ColumnInfo("id", "integer", false, null, true)), + Collections.singletonList(new IndexInfo("orders_name_idx", Collections.singletonList("name"), false, false)), + Collections.singletonList(new ForeignKeyInfo("orders_customer_fk", "customer_id", "customers", "id")) + ); + assertEquals( + "CREATE TABLE \"public\".\"orders\" (\n" + + " \"id\" integer NOT NULL,\n" + + " PRIMARY KEY (\"id\"),\n" + + " CONSTRAINT \"orders_customer_fk\" FOREIGN KEY (\"customer_id\") REFERENCES \"customers\"(\"id\")\n" + + ");\n\n" + + "CREATE INDEX \"orders_name_idx\" ON \"public\".\"orders\" (\"name\");", + ddl + ); + } + + @Test + void executesTransactionsOneByOneWhenJdbcDriverDoesNotSupportTransactions() { + List calls = new ArrayList<>(); + DatabaseAgent agent = new TransactionAgent(nonTransactionalConnection(calls)); + + QueryResult result = agent.executeTransaction(Arrays.asList("UPDATE A SET ID = 1", "UPDATE B SET ID = 2"), "APP"); + + assertEquals(2L, result.getAffected_rows()); + assertEquals( + Arrays.asList("supportsTransactions", "setSchema:APP", "executeUpdate:UPDATE A SET ID = 1", "executeUpdate:UPDATE B SET ID = 2"), + calls + ); + } + + @Test + void buildsTableDdlWithoutSchemaQualifierWhenSchemaIsBlank() { + String ddl = DatabaseAgent.buildTableDdl( + "", + "orders", + Collections.singletonList(new ColumnInfo("id", "integer", false, null, true)), + Collections.emptyList(), + Collections.emptyList() + ); + + assertEquals( + "CREATE TABLE \"orders\" (\n" + + " \"id\" integer NOT NULL,\n" + + " PRIMARY KEY (\"id\")\n" + + ");\n", + ddl + ); + } + + @Test + void buildsTableDdlWithColumnComments() { + String ddl = DdlBuilder.buildTableDdl( + "public", + "orders", + Collections.singletonList(new ColumnInfo( + "display_name", + "varchar", + true, + null, + false, + null, + "User's display name", + null, + null, + 64 + )), + Collections.emptyList(), + Collections.emptyList(), + false, + true + ); + + assertEquals( + "CREATE TABLE \"public\".\"orders\" (\n" + + " \"display_name\" varchar(64)\n" + + ");\n\n" + + "COMMENT ON COLUMN \"public\".\"orders\".\"display_name\" IS 'User''s display name';", + ddl + ); + } + + private static class MinimalAgent implements DatabaseAgent { + @Override + public void connect(ConnectParams params) { + } + + @Override + public boolean testConnection(ConnectParams params) { + return true; + } + + @Override + public List listDatabases() { + return Collections.emptyList(); + } + + @Override + public List listSchemas() { + return Collections.singletonList("public"); + } + + @Override + public List listTables(String schema) { + return Collections.singletonList(new TableInfo("orders", "TABLE")); + } + + @Override + public List getColumns(String schema, String table) { + return Arrays.asList( + new ColumnInfo("id", "integer", false, null, true), + new ColumnInfo("name", "character varying", true, null, false, null, null, null, null, 255) + ); + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return new QueryResult(Collections.emptyList(), Collections.emptyList(), 0L, 0L); + } + + @Override + public void disconnect() { + } + + @Override + public Connection getConnection() { + return null; + } + } + + private static final class TransactionAgent extends MinimalAgent { + private final Connection connection; + + private TransactionAgent(Connection connection) { + this.connection = connection; + } + + @Override + public Connection getConnection() { + return connection; + } + } + + private static final class ReconnectingAgent extends MinimalAgent { + private int connectCount; + private int disconnectCount; + private int firstConnectionValidChecks; + private Connection connection; + + @Override + public void connect(ConnectParams params) { + connectCount += 1; + final int connectionNumber = connectCount; + connection = proxy(Connection.class, (method, args) -> { + String name = method.getName(); + if ("isClosed".equals(name)) { + return false; + } + if ("isValid".equals(name)) { + if (connectionNumber == 1) { + firstConnectionValidChecks += 1; + return false; + } + return true; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + @Override + public void disconnect() { + disconnectCount += 1; + } + + @Override + public Connection getConnection() { + return connection; + } + } + + private static final class CatalogSwitchAgent extends MinimalAgent { + private final List catalogs = new ArrayList<>(); + private String lastSchema = ""; + + @Override + public Connection getConnection() { + return proxy(Connection.class, (method, args) -> { + if ("setCatalog".equals(method.getName())) { + catalogs.add(String.valueOf(args[0])); + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + @Override + public List listTables(String schema) { + lastSchema = schema; + return super.listTables(schema); + } + } + + private static Connection nonTransactionalConnection(List calls) { + return proxy(Connection.class, (method, args) -> { + String name = method.getName(); + if ("getMetaData".equals(name)) { + return proxy(java.sql.DatabaseMetaData.class, (metaMethod, metaArgs) -> { + if ("supportsTransactions".equals(metaMethod.getName())) { + calls.add("supportsTransactions"); + return false; + } + return defaultValue(metaMethod.getReturnType()); + }); + } + if ("createStatement".equals(name)) { + return proxy(java.sql.Statement.class, (stmtMethod, stmtArgs) -> { + if ("execute".equals(stmtMethod.getName())) { + calls.add("execute:" + stmtArgs[0]); + return false; + } + if ("executeUpdate".equals(stmtMethod.getName())) { + calls.add("executeUpdate:" + stmtArgs[0]); + return 1; + } + return defaultValue(stmtMethod.getReturnType()); + }); + } + if ("setSchema".equals(name)) { + calls.add("setSchema:" + args[0]); + return null; + } + if ("setCatalog".equals(name)) { + calls.add("setCatalog:" + args[0]); + return null; + } + if ("setAutoCommit".equals(name) || "commit".equals(name) || "rollback".equals(name)) { + calls.add(name); + return null; + } + if ("getAutoCommit".equals(name)) { + return true; + } + return defaultValue(method.getReturnType()); + }); + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args == null ? new Object[0] : args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + return null; + } + + private static boolean containsCapability(JsonArray capabilities, String expected) { + for (int i = 0; i < capabilities.size(); i++) { + if (expected.equals(capabilities.get(i).getAsString())) { + return true; + } + } + return false; + } + + private static JsonObject protocolContract() { + InputStream stream = CommonJavaCompatibilityTest.class.getResourceAsStream("/agent-protocol-v1.json"); + if (stream == null) { + throw new AssertionError("agent-protocol-v1.json resource missing"); + } + return JsonParser.parseReader(new InputStreamReader(stream, StandardCharsets.UTF_8)).getAsJsonObject(); + } + + private static List strings(JsonArray array) { + List result = new ArrayList<>(); + for (int i = 0; i < array.size(); i++) { + result.add(array.get(i).getAsString()); + } + return result; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/ConfiguredJdbcAgentTest.java b/agents/common/src/test/java/com/dbx/agent/ConfiguredJdbcAgentTest.java new file mode 100644 index 0000000000..7a243436fd --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/ConfiguredJdbcAgentTest.java @@ -0,0 +1,91 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class ConfiguredJdbcAgentTest { + @Test + void buildsJdbcUrlFromProfileTemplateAndUrlParams() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "com.example.Driver", + "jdbc:example://{host}:{port}/{database}" + ); + + String url = profile.buildUrl( + new ConnectParams( + "127.0.0.1", + 1234, + "demo", + "", + "", + "ssl=false", + "", + false + ) + ); + + assertEquals("jdbc:example://127.0.0.1:1234/demo?ssl=false", url); + } + + @Test + void usesExplicitConnectionStringBeforeTemplate() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "com.example.Driver", + "jdbc:example://{host}:{port}/{database}" + ); + + String url = profile.buildUrl( + new ConnectParams( + "127.0.0.1", + 1234, + "demo", + "", + "", + "", + "jdbc:example://server:4321/prod", + false + ) + ); + + assertEquals("jdbc:example://server:4321/prod", url); + } + + @Test + void usesDefaultPortWhenParamsOmitPort() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "com.example.Driver", + "jdbc:example://{host}:{port}/{database}", + 9999 + ); + + String url = profile.buildUrl( + new ConnectParams( + "127.0.0.1", + 0, + "demo", + "", + "", + "", + "", + false + ) + ); + + assertEquals("jdbc:example://127.0.0.1:9999/demo", url); + } + + @Test + void canDisableSchemaSwitchingForDriversThatDoNotSupportContextChanges() { + ConfiguredJdbcAgent agent = new ConfiguredJdbcAgent( + new JdbcAgentProfile( + "com.example.Driver", + "jdbc:example://{host}:{port}/{database}", + 0, + true + ) + ) {}; + + assertEquals("", agent.setSchemaSQL("APP")); + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/EwkbWktDecoderTest.java b/agents/common/src/test/java/com/dbx/agent/EwkbWktDecoderTest.java new file mode 100644 index 0000000000..3be543161d --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/EwkbWktDecoderTest.java @@ -0,0 +1,163 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class EwkbWktDecoderTest { + + private static byte[] hex(String s) { + s = s.replace(" ", ""); + if (s.startsWith("0x") || s.startsWith("0X")) s = s.substring(2); + byte[] out = new byte[s.length() / 2]; + for (int i = 0; i < out.length; i++) { + int hi = Character.digit(s.charAt(i * 2), 16); + int lo = Character.digit(s.charAt(i * 2 + 1), 16); + out[i] = (byte) ((hi << 4) | lo); + } + return out; + } + + // ────────────────────────────────────────────────────────────────────── + // Cross-checked against the Rust ewkb_to_wkt tests in + // crates/dbx-core/src/db/postgres.rs (same hex inputs, same expected WKT). + // ────────────────────────────────────────────────────────────────────── + + @Test + void decodesPointWithSridFromRustFixture() { + byte[] raw = hex("0101000020E6100000C520B07268195D404E62105839F44340"); + assertEquals("POINT(116.397 39.908)", EwkbWktDecoder.decode(raw)); + } + + @Test + void decodesMultiPolygonFromRustFixture() { + byte[] raw = hex( + "0106000020E610000002000000010300000001000000050000000000000000005D40000000000000444000000000" + + "00405D4000000000000044400000000000405D4000000000008044400000000000005D40000000000080444000" + + "00000000005D400000000000004440010300000001000000050000000000000000805D40000000000080434000" + + "00000000C05D4000000000008043400000000000C05D4000000000000044400000000000805D40000000000000" + + "44400000000000805D400000000000804340" + ); + assertEquals( + "MULTIPOLYGON(((116 40,117 40,117 41,116 41,116 40)),((118 39,119 39,119 40,118 40,118 39)))", + EwkbWktDecoder.decode(raw) + ); + } + + @Test + void decodesGeometryCollectionFromRustFixture() { + byte[] raw = hex( + "0107000020E61000000200000001010000000000000000005D40000000000000444001020000000200000000000000" + + "00405D4000000000008044400000000000805D400000000000004540" + ); + assertEquals( + "GEOMETRYCOLLECTION(POINT(116 40),LINESTRING(117 41,118 42))", + EwkbWktDecoder.decode(raw) + ); + } + + @Test + void decodesHexStringWithoutPrefix() { + String hex = "0101000020E6100000C520B07268195D404E62105839F44340"; + assertEquals("POINT(116.397 39.908)", EwkbWktDecoder.decode(hex)); + } + + @Test + void decodesHexStringWith0xPrefix() { + String hex = "0x0101000020E6100000C520B07268195D404E62105839F44340"; + assertEquals("POINT(116.397 39.908)", EwkbWktDecoder.decode(hex)); + } + + @Test + void preservesAlreadyWktTextFromStAsText() { + String wkt = "POINT(116.397 39.908)"; + assertEquals(wkt, EwkbWktDecoder.decode(wkt)); + } + + @Test + void preservesAlreadyEwktTextWithSrid() { + String ewkt = "SRID=4326;POINT(116.397 39.908)"; + assertEquals(ewkt, EwkbWktDecoder.decode(ewkt)); + } + + @Test + void truncatedBytesFallBackToHex() { + // First few bytes of a valid point but truncated mid-coordinate. + byte[] raw = hex("0101000020E6100000C520B072"); + Object decoded = EwkbWktDecoder.decode(raw); + assertTrue(decoded instanceof String); + String s = (String) decoded; + assertTrue(s.startsWith("0x"), "expected hex fallback, got: " + s); + assertEquals("0x0101000020e6100000c520b072", s); + } + + @Test + void unparseableHexStringPreservedWithHexPrefix() { + // Looks like hex (even, leading 01), but bytes are not valid EWKB. + String mostly = "0123456789ABCDEF"; + Object decoded = EwkbWktDecoder.decode(mostly); + assertTrue(decoded instanceof String); + // Shape A: parseEwkb returns null, so we re-emit the input ensuring 0x prefix. + assertTrue(((String) decoded).toLowerCase().startsWith("0x"), decoded.toString()); + } + + @Test + void nullInputReturnsNull() { + assertNull(EwkbWktDecoder.decode(null)); + } + + @Test + void formatsWholeNumberDoublesWithoutTrailingDot() { + // Ensure POINT(116 40) is "116 40", not "116.0 40.0" — matches Rust f64 Display. + byte[] raw = hex("0101000020E61000000000000000005D400000000000004440"); + assertEquals("POINT(116 40)", EwkbWktDecoder.decode(raw)); + } + + @Test + void emptyPointBecomesPointEmpty() { + // A 2D point with both coordinates = NaN encodes as "POINT EMPTY" in the decoder. + // EWKB: little-endian Point with SRID=0 omitted; type=1, then two NaN doubles. + byte[] raw = hex("0101000000" + nanLe() + nanLe()); + assertEquals("POINT EMPTY", EwkbWktDecoder.decode(raw)); + } + + private static String nanLe() { + // IEEE 754 quiet NaN, little-endian — same value Rust's f64::NAN encodes to. + long bits = Double.doubleToRawLongBits(Double.NaN); + StringBuilder sb = new StringBuilder(16); + for (int i = 0; i < 8; i++) { + int b = (int) ((bits >>> (i * 8)) & 0xFF); + sb.append(String.format("%02x", b)); + } + return sb.toString(); + } + + @Test + void pgGeometryReflectiveDecodeReturnsValueText() { + // Synthetic stand-in for org.postgis.PGgeometry — same shape: getValue() returns EWKT. + Object pgGeom = new FakePGgeometry("SRID=4326;POINT(116.397 39.908)"); + assertEquals("SRID=4326;POINT(116.397 39.908)", EwkbWktDecoder.decode(pgGeom)); + } + + @SuppressWarnings("unused") + public static final class FakePGgeometry { + private final String value; + + public FakePGgeometry(String value) { + this.value = value; + } + + // Class name won't match the reflective list so this falls through to toString(); + // we keep it here purely to assert the toString() fallback also works (no NPE). + public String getValue() { + return value; + } + + @Override + public String toString() { + return value; + } + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcAgentProfileExtendedTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcAgentProfileExtendedTest.java new file mode 100644 index 0000000000..e8848f2fde --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/JdbcAgentProfileExtendedTest.java @@ -0,0 +1,109 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class JdbcAgentProfileExtendedTest { + @Test + void appendsQueryParamsToJdbcUrlTemplates() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:example://{host}:{port}/{database}", + 1234 + ); + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(0); + params.setDatabase("analytics"); + params.setUrl_params("ssl=true"); + + assertEquals("jdbc:example://db.example.com:1234/analytics?ssl=true", profile.buildUrl(params)); + } + + @Test + void appendsSemicolonParamsToSemicolonJdbcUrlTemplates() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:exa:{host}:{port};schema={database}", + 8563 + ); + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(0); + params.setDatabase("analytics"); + params.setUrl_params("encryption=1"); + + assertEquals("jdbc:exa:db.example.com:8563;schema=analytics;encryption=1", profile.buildUrl(params)); + } + + @Test + void appendsCommaParamsToTeradataJdbcUrlTemplates() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:teradata://{host}/DBS_PORT={port},DATABASE={database}", + 1025 + ); + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(0); + params.setDatabase("analytics"); + params.setUrl_params("LOGMECH=LDAP"); + + assertEquals("jdbc:teradata://db.example.com/DBS_PORT=1025,DATABASE=analytics,LOGMECH=LDAP", profile.buildUrl(params)); + } + + @Test + void appendsColonPropertiesToDb2JdbcUrlTemplates() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:db2://{host}:{port}/{database}", + 50000 + ); + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(0); + params.setDatabase("analytics"); + params.setUrl_params("sslConnection=true"); + + assertEquals("jdbc:db2://db.example.com:50000/analytics:sslConnection=true;", profile.buildUrl(params)); + } + + @Test + void preservesExplicitConnectionStringForVendorSpecificUrls() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:example://{host}:{port}/{database}", + 1234 + ); + ConnectParams params = new ConnectParams(); + params.setConnection_string("jdbc:example:special://host/path;token=secret"); + + assertEquals("jdbc:example:special://host/path;token=secret", profile.buildUrl(params)); + } + + @Test + void exposesDialectProfileOptions() { + JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:example://{host}:{port}/{database}", + 1234, + false, + java.util.Collections.singleton("SYS"), + java.util.Arrays.asList("TABLE", "VIEW"), + "`", + "USE", + false, + true, + true, + true + ); + + assertEquals("USE `analytics``prod`", profile.schemaSwitchSql("analytics`prod")); + assertEquals("`orders`", profile.quoteIdentifier("orders")); + assertEquals(false, profile.getCatalogFallbackEnabled()); + assertEquals(true, profile.getNativeTableDdlSupported()); + assertEquals(true, profile.getObjectSourceSupported()); + assertEquals(true, profile.getTriggersSupported()); + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java new file mode 100644 index 0000000000..9acee273bc --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -0,0 +1,98 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.ResultSet; +import java.sql.Types; +import javax.sql.rowset.serial.SerialBlob; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class JdbcExecutorTest { + @Test + void stringResultValueFormatsBlobWithoutUsingStringConversion() throws Exception { + ResultSet rs = resultSet( + new byte[]{0x01, 0x2A, (byte) 0xFF}, + () -> { + throw new AssertionError("BLOB columns should not be read with getString"); + } + ); + + assertEquals("0x012aff", JdbcExecutor.stringResultValue(rs, 1, Types.BLOB)); + } + + @Test + void defaultResultValueReadsClobColumnsAsText() throws Exception { + ResultSet rs = resultSet(null, () -> "hello clob"); + + assertEquals("hello clob", JdbcExecutor.INSTANCE.defaultResultValue(rs, 1, Types.CLOB)); + } + + @Test + void defaultResultValueNormalizesBlobObjectsFromFallbackTypes() throws Exception { + ResultSet rs = resultSet(new SerialBlob(new byte[]{0x0A, 0x0B}), null, false); + + assertEquals("0x0a0b", JdbcExecutor.INSTANCE.defaultResultValue(rs, 1, Types.OTHER)); + } + + private static ResultSet resultSet(byte[] bytes, StringSupplier stringSupplier) { + return resultSet(bytes, stringSupplier, false); + } + + private static ResultSet resultSet(Object objectValue, StringSupplier stringSupplier, boolean wasNull) { + InvocationHandler handler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "getObject": + return objectValue; + case "getBytes": + return objectValue instanceof byte[] ? objectValue : null; + case "getString": + return stringSupplier == null ? null : stringSupplier.get(); + case "wasNull": + return wasNull; + default: + return defaultValue(method.getReturnType()); + } + }; + return (ResultSet) Proxy.newProxyInstance( + ResultSet.class.getClassLoader(), + new Class[]{ResultSet.class}, + handler + ); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) { + return false; + } + if (type == Byte.TYPE) { + return (byte) 0; + } + if (type == Short.TYPE) { + return (short) 0; + } + if (type == Integer.TYPE) { + return 0; + } + if (type == Long.TYPE) { + return 0L; + } + if (type == Float.TYPE) { + return 0f; + } + if (type == Double.TYPE) { + return 0.0d; + } + if (type == Character.TYPE) { + return '\0'; + } + return null; + } + + private interface StringSupplier { + String get() throws Exception; + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java b/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java new file mode 100644 index 0000000000..130c3dd9e1 --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java @@ -0,0 +1,304 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Types; +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class PostgresLikeAgentTest { + @Test + void metadataQueriesUsePgCatalogInsteadOfInformationSchema() { + TestPostgresLikeAgent agent = new TestPostgresLikeAgent(); + agent.connect(new ConnectParams()); + + agent.listSchemas(); + agent.listTables("app"); + agent.listObjects("app"); + agent.getColumns("app", "orders"); + agent.listForeignKeys("app", "orders"); + agent.listTriggers("app", "orders"); + + String sql = String.join("\n", MetadataSqlFake.statements); + + assertFalse(sql.contains("FROM information_schema"), sql); + assertFalse(sql.contains("JOIN information_schema"), sql); + assertTrue(sql.contains("pg_catalog.pg_namespace"), sql); + assertTrue(sql.contains("pg_catalog.pg_class"), sql); + assertTrue(sql.contains("pg_catalog.pg_proc"), sql); + assertTrue(sql.contains("pg_catalog.pg_attribute"), sql); + assertTrue(sql.contains("pg_catalog.pg_constraint"), sql); + assertTrue(sql.contains("pg_catalog.pg_trigger"), sql); + assertFalse(sql.contains(" AS key "), sql); + assertFalse(sql.contains(" key."), sql); + } + + @Test + void postgisGeometryTypeNameDetection() { + assertTrue(PostgresLikeAgent.isPostgisGeometryTypeName("geometry")); + assertTrue(PostgresLikeAgent.isPostgisGeometryTypeName("GEOMETRY")); + assertTrue(PostgresLikeAgent.isPostgisGeometryTypeName(" Geography ")); + assertTrue(PostgresLikeAgent.isPostgisGeometryTypeName("public.geometry")); + assertTrue(PostgresLikeAgent.isPostgisGeometryTypeName("geometry(Point,4326)")); + assertFalse(PostgresLikeAgent.isPostgisGeometryTypeName("")); + assertFalse(PostgresLikeAgent.isPostgisGeometryTypeName(null)); + assertFalse(PostgresLikeAgent.isPostgisGeometryTypeName("text")); + assertFalse(PostgresLikeAgent.isPostgisGeometryTypeName("vector")); + } + + @Test + void executeQueryDecodesGeometryColumnsToWktAndReportsColumnTypes() { + TestPostgresLikeAgent agent = new TestPostgresLikeAgent(); + agent.connect(new ConnectParams()); + + QueryResult result = JdbcExecutor.INSTANCE.readResultSet( + GeometryResultSet.create(), + 5L, + JdbcExecutor.DEFAULT_MAX_ROWS, + agent.geometryAwareResolverForTest() + ); + + assertEquals(2, result.getColumns().size()); + assertEquals("id", result.getColumns().get(0)); + assertEquals("geom", result.getColumns().get(1)); + + // column_types is reported via JDBC getColumnTypeName + assertEquals(2, result.getColumn_types().size()); + assertEquals("int4", result.getColumn_types().get(0)); + assertEquals("geometry", result.getColumn_types().get(1)); + + // The geometry cell is decoded to WKT (matches Rust ewkb_to_wkt fixture). + assertEquals(1, result.getRows().size()); + assertEquals(1, result.getRows().get(0).get(0)); + assertEquals("POINT(116.397 39.908)", result.getRows().get(0).get(1)); + } + + private static final class TestPostgresLikeAgent extends PostgresLikeAgent { + private TestPostgresLikeAgent() { + super(new PostgresLikeAgentProfile( + PostgresLikeAgentTest.class.getName(), + "jdbc:test://{host}:{port}/{database}" + )); + } + + @Override + protected Connection openConnection(ConnectParams params) { + return MetadataSqlFake.connection(); + } + + JdbcExecutor.ColumnAwareResultValueReader geometryAwareResolverForTest() { + return (rs, index, sqlType, columnTypeName) -> { + if (PostgresLikeAgent.isPostgisGeometryTypeName(columnTypeName)) { + Object raw = rs.getObject(index); + if (rs.wasNull() || raw == null) { + return null; + } + return EwkbWktDecoder.decode(raw); + } + if (sqlType == Types.INTEGER) { + return rs.getInt(index); + } + return rs.getObject(index); + }; + } + } + + private static final class GeometryResultSet { + static ResultSet create() { + ResultSetMetaData meta = (ResultSetMetaData) Proxy.newProxyInstance( + GeometryResultSet.class.getClassLoader(), + new Class[]{ResultSetMetaData.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "getColumnCount": + return 2; + case "getColumnLabel": + case "getColumnName": + return ((Integer) args[0]) == 1 ? "id" : "geom"; + case "getColumnType": + return ((Integer) args[0]) == 1 ? Types.INTEGER : Types.OTHER; + case "getColumnTypeName": + return ((Integer) args[0]) == 1 ? "int4" : "geometry"; + default: + return null; + } + } + ); + + // POINT(116.397 39.908) with SRID=4326, little-endian. + String hex = "0101000020E6100000C520B07268195D404E62105839F44340"; + byte[] geomBytes = parseHex(hex); + + int[] cursor = new int[]{0}; // 0 = before first row + return (ResultSet) Proxy.newProxyInstance( + GeometryResultSet.class.getClassLoader(), + new Class[]{ResultSet.class}, + (proxy, method, args) -> { + String name = method.getName(); + switch (name) { + case "getMetaData": + return meta; + case "next": + cursor[0] += 1; + return cursor[0] == 1; + case "getInt": + return cursor[0] == 1 ? 1 : 0; + case "getObject": + return ((Integer) args[0]) == 2 ? geomBytes : 1; + case "wasNull": + return false; + case "close": + return null; + default: + return defaultPrimitive(method.getReturnType()); + } + } + ); + } + + private static byte[] parseHex(String s) { + byte[] out = new byte[s.length() / 2]; + for (int i = 0; i < out.length; i++) { + int hi = Character.digit(s.charAt(i * 2), 16); + int lo = Character.digit(s.charAt(i * 2 + 1), 16); + out[i] = (byte) ((hi << 4) | lo); + } + return out; + } + + private static Object defaultPrimitive(Class t) { + if (Boolean.TYPE.equals(t)) return false; + if (Integer.TYPE.equals(t)) return 0; + if (Long.TYPE.equals(t)) return 0L; + if (Double.TYPE.equals(t)) return 0.0; + return null; + } + } + + private static final class MetadataSqlFake { + private static final List statements = new ArrayList(); + + private static Connection connection() { + statements.clear(); + return proxy(Connection.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("prepareStatement".equals(name)) { + statements.add((String) args[0]); + return preparedStatement(); + } + if ("isClosed".equals(name)) { + return false; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static PreparedStatement preparedStatement() { + final ResultSet resultSet = emptyResultSet(); + return proxy(PreparedStatement.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("executeQuery".equals(name)) { + return resultSet; + } + if ("setString".equals(name)) { + statements.add("param:" + args[0] + "=" + args[1]); + return null; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet emptyResultSet() { + final ResultSetMetaData metadata = proxy(ResultSetMetaData.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + if ("getColumnCount".equals(method.getName())) { + return 0; + } + return defaultValue(method.getReturnType()); + } + }); + return proxy(ResultSet.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + return false; + } + if ("getMetaData".equals(name)) { + return metadata; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static T proxy(Class type, final MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java b/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java new file mode 100644 index 0000000000..55ace77357 --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java @@ -0,0 +1,471 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSet; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class StandardJdbcMetadataTest { + private final JdbcAgentProfile profile = new JdbcAgentProfile( + "example.Driver", + "jdbc:example://{host}:{port}/{database}", + 0, + false, + Collections.singleton("SYS"), + Arrays.asList("TABLE", "VIEW", "BASE TABLE") + ); + + @Test + void listsSchemasWithFilteringAndStableOrdering() { + Connection conn = connection( + rows(row("TABLE_SCHEM", "APP"), row("TABLE_SCHEM", "SYS"), row("TABLE_SCHEM", "PUBLIC")), + rows(), + rows(), + rows(), + rows(), + rows() + ); + + assertEquals(Arrays.asList("APP", "PUBLIC"), StandardJdbcMetadata.INSTANCE.listSchemas(conn, profile)); + } + + @Test + void listsSchemasWhenConnectionGetSchemaIsUnsupported() { + Connection conn = connection( + rows(row("TABLE_SCHEM", "APP"), row("TABLE_SCHEM", "PUBLIC")), + rows(), + rows(), + rows(), + rows(), + rows(), + true + ); + + assertEquals(Arrays.asList("APP", "PUBLIC"), StandardJdbcMetadata.INSTANCE.listSchemas(conn, profile)); + } + + @Test + void listsSchemasWhenConnectionGetSchemaThrowsAbstractMethodError() { + Connection conn = connection( + rows(row("TABLE_SCHEM", "APP"), row("TABLE_SCHEM", "PUBLIC")), + rows(), + rows(), + rows(), + rows(), + rows(), + UnsupportedSchemaCall.ABSTRACT_METHOD_ERROR + ); + + assertEquals(Arrays.asList("APP", "PUBLIC"), StandardJdbcMetadata.INSTANCE.listSchemas(conn, profile)); + } + + @Test + void listsSchemasWithEmptyResultWhenDriverSchemaMetadataIsUnsupported() { + Connection conn = connection( + rows(), + rows(), + rows(), + rows(), + rows(), + rows(), + UnsupportedSchemaCall.METADATA_AND_CONNECTION_ABSTRACT_METHOD_ERROR + ); + + assertEquals(Collections.emptyList(), StandardJdbcMetadata.INSTANCE.listSchemas(conn, profile)); + } + + @Test + void listsTablesWithNormalizedTypesAndStableOrdering() { + Connection conn = connection( + rows(), + rows(row("TABLE_NAME", "ZETA", "TABLE_TYPE", "BASE TABLE", "REMARKS", "last"), + row("TABLE_NAME", "ALPHA", "TABLE_TYPE", "VIEW", "REMARKS", "first")), + rows(), + rows(), + rows(), + rows() + ); + + List tables = StandardJdbcMetadata.INSTANCE.listTables(conn, profile, "", "APP"); + + assertEquals("ALPHA", tables.get(0).getName()); + assertEquals("VIEW", tables.get(0).getTable_type()); + assertEquals("ZETA", tables.get(1).getName()); + assertEquals("TABLE", tables.get(1).getTable_type()); + } + + @Test + void listTablesUsesDriverTableTypesWithinProfileAllowList() { + AtomicReference capturedTypes = new AtomicReference<>(); + Connection conn = connection( + rows(), + rows(row("TABLE_NAME", "ORDERS", "TABLE_TYPE", "TABLE", "REMARKS", null)), + rows(), + rows(), + rows(), + rows(), + rows( + row("TABLE_TYPE", "TABLE"), + row("TABLE_TYPE", "LOCAL TEMPORARY"), + row("TABLE_TYPE", "BASE TABLE") + ), + capturedTypes + ); + + StandardJdbcMetadata.INSTANCE.listTables(conn, profile, "", "APP"); + + assertEquals(Arrays.asList("TABLE", "BASE TABLE"), Arrays.asList(capturedTypes.get())); + } + + @Test + void listTablesAllowsHanaColumnAndRowTables() { + JdbcAgentProfile hanaProfile = new JdbcAgentProfile( + "com.sap.db.jdbc.Driver", + "jdbc:sap://{host}:{port}/?databaseName={database}", + 30015, + false, + Collections.emptySet(), + Arrays.asList("COLUMN TABLE", "ROW TABLE", "TABLE", "VIEW") + ); + AtomicReference capturedTypes = new AtomicReference<>(); + Connection conn = connection( + rows(), + rows( + row("TABLE_NAME", "ORDERS", "TABLE_TYPE", "COLUMN TABLE", "REMARKS", null), + row("TABLE_NAME", "AUDIT_LOG", "TABLE_TYPE", "ROW TABLE", "REMARKS", null), + row("TABLE_NAME", "ACTIVE_ORDERS", "TABLE_TYPE", "VIEW", "REMARKS", null) + ), + rows(), + rows(), + rows(), + rows(), + rows( + row("TABLE_TYPE", "COLUMN TABLE"), + row("TABLE_TYPE", "ROW TABLE"), + row("TABLE_TYPE", "VIEW") + ), + capturedTypes + ); + + List tables = StandardJdbcMetadata.INSTANCE.listTables(conn, hanaProfile, "", "APP"); + + assertEquals(Arrays.asList("COLUMN TABLE", "ROW TABLE", "VIEW"), Arrays.asList(capturedTypes.get())); + assertEquals(Arrays.asList("ACTIVE_ORDERS", "AUDIT_LOG", "ORDERS"), Arrays.asList( + tables.get(0).getName(), + tables.get(1).getName(), + tables.get(2).getName() + )); + assertEquals("VIEW", tables.get(0).getTable_type()); + assertEquals("TABLE", tables.get(1).getTable_type()); + assertEquals("TABLE", tables.get(2).getTable_type()); + } + + @Test + void mapsColumnsWithPrimaryKeysAndLengths() { + Connection conn = connection( + rows(), + rows(), + rows(row("COLUMN_NAME", "ID")), + rows(row( + "COLUMN_NAME", "ID", + "TYPE_NAME", "INTEGER", + "NULLABLE", DatabaseMetaData.columnNoNulls, + "COLUMN_DEF", null, + "REMARKS", "identifier", + "COLUMN_SIZE", 10, + "DECIMAL_DIGITS", 0 + ), row( + "COLUMN_NAME", "NAME", + "TYPE_NAME", "VARCHAR", + "NULLABLE", DatabaseMetaData.columnNullable, + "COLUMN_DEF", null, + "REMARKS", null, + "COLUMN_SIZE", 64, + "DECIMAL_DIGITS", 0 + )), + rows(), + rows() + ); + + List columns = StandardJdbcMetadata.INSTANCE.getColumns(conn, profile, "", "APP", "ORDERS"); + + assertEquals("ID", columns.get(0).getName()); + assertTrue(columns.get(0).getIs_primary_key()); + assertFalse(columns.get(0).getIs_nullable()); + assertEquals("NAME", columns.get(1).getName()); + assertEquals(Integer.valueOf(64), columns.get(1).getCharacter_maximum_length()); + } + + @Test + void groupsIndexColumnsInOrdinalOrder() { + Connection conn = connection( + rows(), + rows(), + rows(), + rows(), + rows(row("INDEX_NAME", "IDX_ORDERS", "COLUMN_NAME", "B", "ORDINAL_POSITION", (short) 2, "NON_UNIQUE", true), + row("INDEX_NAME", "IDX_ORDERS", "COLUMN_NAME", "A", "ORDINAL_POSITION", (short) 1, "NON_UNIQUE", true)), + rows() + ); + + List indexes = StandardJdbcMetadata.INSTANCE.listIndexes(conn, "APP", "ORDERS"); + + assertEquals(1, indexes.size()); + assertEquals(Arrays.asList("A", "B"), indexes.get(0).getColumns()); + assertFalse(indexes.get(0).getIs_unique()); + } + + @Test + void mapsForeignKeysAndEmptyTriggers() { + Connection conn = connection( + rows(), + rows(), + rows(), + rows(), + rows(), + rows(row("FK_NAME", null, "FKCOLUMN_NAME", "CUSTOMER_ID", "PKTABLE_NAME", "CUSTOMERS", "PKCOLUMN_NAME", "ID")) + ); + + List foreignKeys = StandardJdbcMetadata.INSTANCE.listForeignKeys(conn, "APP", "ORDERS"); + + assertEquals("", foreignKeys.get(0).getName()); + assertEquals("CUSTOMER_ID", foreignKeys.get(0).getColumn()); + assertEquals(Collections.emptyList(), StandardJdbcMetadata.INSTANCE.listTriggers("APP", "ORDERS")); + } + + private static Connection connection( + ResultSet schemas, + ResultSet tables, + ResultSet primaryKeys, + ResultSet columns, + ResultSet indexes, + ResultSet foreignKeys + ) { + return connection(schemas, tables, primaryKeys, columns, indexes, foreignKeys, false); + } + + private static Connection connection( + ResultSet schemas, + ResultSet tables, + ResultSet primaryKeys, + ResultSet columns, + ResultSet indexes, + ResultSet foreignKeys, + boolean unsupportedGetSchema + ) { + return connection( + schemas, + tables, + primaryKeys, + columns, + indexes, + foreignKeys, + unsupportedGetSchema ? UnsupportedSchemaCall.RUNTIME_EXCEPTION : UnsupportedSchemaCall.NONE, + null, + null + ); + } + + private static Connection connection( + ResultSet schemas, + ResultSet tables, + ResultSet primaryKeys, + ResultSet columns, + ResultSet indexes, + ResultSet foreignKeys, + UnsupportedSchemaCall unsupportedSchemaCall + ) { + return connection(schemas, tables, primaryKeys, columns, indexes, foreignKeys, unsupportedSchemaCall, null, null); + } + + private static Connection connection( + ResultSet schemas, + ResultSet tables, + ResultSet primaryKeys, + ResultSet columns, + ResultSet indexes, + ResultSet foreignKeys, + ResultSet tableTypes, + AtomicReference capturedTableTypes + ) { + return connection(schemas, tables, primaryKeys, columns, indexes, foreignKeys, UnsupportedSchemaCall.NONE, tableTypes, capturedTableTypes); + } + + private static Connection connection( + ResultSet schemas, + ResultSet tables, + ResultSet primaryKeys, + ResultSet columns, + ResultSet indexes, + ResultSet foreignKeys, + UnsupportedSchemaCall unsupportedSchemaCall, + ResultSet tableTypes, + AtomicReference capturedTableTypes + ) { + DatabaseMetaData meta = proxy(DatabaseMetaData.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("getSchemas".equals(name)) { + if (unsupportedSchemaCall == UnsupportedSchemaCall.METADATA_AND_CONNECTION_ABSTRACT_METHOD_ERROR) { + throw new AbstractMethodError("Unimplemented method: getSchemas()"); + } + return schemas; + } + if ("getTables".equals(name)) { + if (capturedTableTypes != null && args != null && args.length > 3) { + capturedTableTypes.set((String[]) args[3]); + } + return tables; + } + if ("getTableTypes".equals(name) && tableTypes != null) { + return tableTypes; + } + if ("getPrimaryKeys".equals(name)) { + return primaryKeys; + } + if ("getColumns".equals(name)) { + return columns; + } + if ("getIndexInfo".equals(name)) { + return indexes; + } + if ("getImportedKeys".equals(name)) { + return foreignKeys; + } + if ("getCatalogs".equals(name)) { + return rows(); + } + return defaultValue(method.getReturnType()); + } + }); + return proxy(Connection.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + if ("getMetaData".equals(method.getName())) { + return meta; + } + if ("getSchema".equals(method.getName())) { + if (unsupportedSchemaCall == UnsupportedSchemaCall.RUNTIME_EXCEPTION) { + throw new RuntimeException("Unimplemented method: getSchema()"); + } + if (unsupportedSchemaCall == UnsupportedSchemaCall.ABSTRACT_METHOD_ERROR + || unsupportedSchemaCall == UnsupportedSchemaCall.METADATA_AND_CONNECTION_ABSTRACT_METHOD_ERROR) { + throw new AbstractMethodError("Unimplemented method: getSchema()"); + } + return null; + } + if ("getCatalog".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet rows(Map... rows) { + return proxy(ResultSet.class, new MethodHandler() { + private int index = -1; + + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + index += 1; + return index < rows.length; + } + if ("close".equals(name)) { + return null; + } + Object value = rows[index].get(args[0]); + if ("getString".equals(name)) { + return value == null ? null : String.valueOf(value); + } + if ("getObject".equals(name)) { + return value; + } + if ("getInt".equals(name)) { + return value instanceof Number ? ((Number) value).intValue() : 0; + } + if ("getShort".equals(name)) { + return value instanceof Number ? ((Number) value).shortValue() : 0; + } + if ("getBoolean".equals(name)) { + return value instanceof Boolean && (Boolean) value; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static Map row(Object... values) { + Map row = new LinkedHashMap<>(); + for (int i = 0; i < values.length; i += 2) { + row.put(String.valueOf(values[i]), values[i + 1]); + } + return row; + } + + private static T proxy(Class type, final MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } + + private enum UnsupportedSchemaCall { + NONE, + RUNTIME_EXCEPTION, + ABSTRACT_METHOD_ERROR, + METADATA_AND_CONNECTION_ABSTRACT_METHOD_ERROR + } +} diff --git a/agents/docs/agent-authoring.md b/agents/docs/agent-authoring.md new file mode 100644 index 0000000000..ec6b6ede78 --- /dev/null +++ b/agents/docs/agent-authoring.md @@ -0,0 +1,318 @@ +# Agent Authoring Guide + +This guide defines the expected shape of a DBX agent. Treat it as the checklist for adding or reviewing an agent. + +## Agent Contract + +Every agent is a standalone JVM process that: + +- Implements `com.dbx.agent.DatabaseAgent`. +- Prefer extending `com.dbx.agent.ConfiguredJdbcAgent` for standard JDBC agents. +- Extend `com.dbx.agent.AbstractJdbcAgent` when the database needs custom metadata SQL but can still share lifecycle and execution behavior. +- Starts with `new JsonRpcServer(new ()).run()` in its `main` method. +- Talks to DBX over stdin/stdout JSON-RPC 2.0. +- Uses JDBC for database access unless the module is explicitly designed around a non-JDBC protocol. +- Produces one shadow JAR named `dbx-agent-.jar`. + +The public behavior should be consistent across agents even when each database has different SQL dialects. + +## Required Methods + +Each agent must implement these capabilities: + +- `connect(params)`: handled by the shared JDBC foundation for JDBC agents. +- `testConnection(params)`: handled by the shared JDBC foundation for JDBC agents. +- `listDatabases()`: return visible catalogs/databases when the database supports them; otherwise return one sensible default database. +- `listSchemas()`: return schemas in stable order. +- `listTables(schema)`: return table-like objects for the selected schema, with normalized `type` values where possible. +- `getColumns(schema, table)`: return column metadata, nullability, defaults, key flags, numeric precision/scale, character length, and comments when available. +- `listIndexes(schema, table)`: return one `IndexInfo` per index, preserving column order. +- `listForeignKeys(schema, table)`: return outbound foreign keys when available. +- `listTriggers(schema, table)`: return triggers when available; return an empty list if the database has no trigger metadata. +- `executeQuery(sql, schema, options)`: inherited from the shared JDBC foundation unless the agent has a documented custom execution behavior. +- `disconnect()`: inherited from the shared JDBC foundation. +- `getConnection()`: inherited from the shared JDBC foundation. + +Throw `IllegalStateException("Not connected")` when a method requires a connection and none exists. + +## SQL Execution Rules + +Do not classify statements by SQL prefix inside individual agents. + +For most JDBC agents, inherit this from `AbstractJdbcAgent` or `ConfiguredJdbcAgent`. Do not copy the execution method into the agent class. If a custom value reader is required, override `resultValue(...)`: + +```java +@Override +protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> rs.getString(index)); +} +``` + +For simple standard JDBC metadata agents, start with a profile: + +```java +public final class ExampleAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile EXAMPLE_PROFILE = new JdbcAgentProfile( + "com.example.Driver", + "jdbc:example://{host}:{port}/{database}", + 1234 + ); + + public ExampleAgent() { + super(EXAMPLE_PROFILE); + } +} +``` + +`JdbcExecutor` owns these behaviors: + +- Trims a trailing semicolon before execution. +- Handles `BEGIN`, `COMMIT`, and `ROLLBACK`. +- Runs schema switching SQL before the user statement. +- Executes statements with `Statement.execute(...)`. +- Reads `ResultSet` output for any statement type, not only `SELECT`. +- Returns update counts for update statements. +- Caps result rows at `options.maxRows`, defaulting to `JdbcExecutor.DEFAULT_MAX_ROWS`. +- Applies `options.fetchSize` to the JDBC statement when provided. +- Marks `truncated = true` only when more rows exist beyond the cap. + +An agent must not reintroduce local copies of: + +- `QUERY_PREFIXES` +- `MAX_ROWS` +- `executeUpdate(trimmedSql)` in `executeQuery` +- result truncation based on `rows.size >= MAX_ROWS` +- `Class.forName(...)` / `DriverManager.getConnection(...)` lifecycle boilerplate outside shared foundation extension points +- local `disconnect()` copies for JDBC agents +- local `JdbcExecutor.INSTANCE.execute(...)` copies for standard query execution + +## Schema And Identifier Rules + +Override `setSchemaSQL(schema)` for the database dialect, or configure the profile dialect options when standard quoting is enough. + +Use `JdbcIdentifiers` helpers when quoting identifiers: + +```java +@Override +public String setSchemaSQL(String schema) { + return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); +} +``` + +If the database does not support a schema switching statement, return an empty string: + +```java +@Override +public String setSchemaSQL(String schema) { + return ""; +} +``` + +Never concatenate unquoted user-provided schema names into schema-switching SQL unless the target database requires unquoted identifiers and the value has already been validated. + +For metadata queries, prefer prepared statements for `schema`, `table`, and other user-controlled values. + +## Shared JDBC Foundation + +Use the smallest shared base that matches the database: + +- `ConfiguredJdbcAgent`: standard JDBC lifecycle, execution, and standard `DatabaseMetaData`-based metadata. +- `PostgresLikeAgent`: PostgreSQL-family metadata with shared lifecycle and execution. +- `AbstractJdbcAgent`: custom metadata SQL with shared lifecycle, execution, paging, transactions, result conversion, and DDL fallback. + +If a database cannot yet use the shared foundation, add it to the validation allowlist with a specific reason and tests covering the custom behavior. Remove the allowlist entry when the module migrates. + +## Driver Packaging + +Each module chooses one of two driver modes. + +### Bundled Driver + +Use this when the driver is redistributable from Maven Central or another permitted repository: + +```groovy +dependencies { + implementation 'com.example:example-jdbc:1.2.3' +} +``` + +The root Gradle convention supplies `project(':common')`, `project(':test-support')`, JUnit, Java toolchains, the Shadow plugin, and the `dbx-agent-` archive name for included agent modules. + +No manifest flag is needed. + +### External Driver + +Use this when the driver cannot be redistributed or must be supplied by the user: + +```groovy +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) +} + +tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'Example DB', + 'Agent-External-Driver': 'true', + 'Main-Class': 'com.dbx.agent.example.ExampleAgent' + ) + } +} +``` + +The release workflow reads `Agent-External-Driver: true` and emits `external_driver_required: true` in `agent-registry.json`. + +## Module Registration + +When adding an agent named `exampledb`: + +- Create the module under `drivers/exampledb`. +- Add `exampledb` to `driverModules` in `settings.gradle`. +- Add `"exampledb": "0.1.0"` to `versions.json`. +- Set `Agent-Label` to the user-facing database name. +- Set `Main-Class` to the Java agent class, usually `com.dbx.agent.exampledb.ExampledbAgent`. +- Add the database to the README support table. + +The root `build.gradle` convention derives the archive name from the module name, so `exampledb` builds `dbx-agent-exampledb.jar` without per-module archive configuration. + +`versions.json` must contain only modules included in `settings.gradle`, excluding infrastructure modules such as `common` and `test-support`. + +## Runtime Selection + +Default to: + +```groovy +def java8Projects = ['common', 'test-support', 'oracle-10g'] as Set + +subprojects { + java { + toolchain { + languageVersion = JavaLanguageVersion.of(java8Projects.contains(name) ? 8 : 21) + } + } +} +``` + +Most agents use JRE 21. Only use JRE 8 for drivers that require it, such as legacy Oracle 10g support. If an agent needs a special runtime, update the root Gradle convention, the release workflow JRE detection logic, and document the reason in the module. + +## Tests + +Every JDBC agent should have at least one execution-path regression test. + +For agents that can run with an embedded or in-memory database, prefer both shared behavior contracts: + +```java +import java.util.List; + +class H2ExecutionBehaviorTest extends JdbcExecutionBehaviorTest { + @Override + protected DatabaseAgent createConnectedAgent(String databaseName) { + return H2AgentFixtures.createConnectedAgent(databaseName); + } + + @Override + protected String resultSetSql() { + return "CALL 42"; + } + + @Override + protected List expectedResultSetColumns() { + return List.of("42"); + } + + @Override + protected List> expectedResultSetRows() { + return List.of(List.of(42)); + } + + @Override + protected String rowsSql(int rowCount) { + return "SELECT X FROM SYSTEM_RANGE(1, " + rowCount + ")"; + } +} + +class H2MetadataBehaviorTest extends JdbcMetadataBehaviorTest { + @Override + protected DatabaseAgent createConnectedAgent(String databaseName) { + return H2AgentFixtures.createConnectedAgent(databaseName); + } + + @Override + protected List metadataFixtureSql() { + return List.of( + "CREATE TABLE BETA_TABLE (ID INT PRIMARY KEY)", + "CREATE TABLE ALPHA_TABLE (ID INT PRIMARY KEY)", + "CREATE TABLE COLUMN_ORDER_SAMPLE (ID INT PRIMARY KEY, NAME VARCHAR(64), CREATED_AT TIMESTAMP)" + ); + } + + @Override + protected String metadataSchema() { + return "PUBLIC"; + } + + @Override + protected List expectedTablesInOrder() { + return List.of("ALPHA_TABLE", "BETA_TABLE", "COLUMN_ORDER_SAMPLE"); + } + + @Override + protected String metadataColumnsTable() { + return "COLUMN_ORDER_SAMPLE"; + } + + @Override + protected List expectedColumnsInOrder() { + return List.of("ID", "NAME", "CREATED_AT"); + } +} + +final class H2AgentFixtures { + private H2AgentFixtures() { + } + + static DatabaseAgent createConnectedAgent(String databaseName) { + H2Agent agent = new H2Agent(); + agent.connect(new ConnectParams("mem:" + databaseName + ";DB_CLOSE_DELAY=-1")); + return agent; + } +} +``` + +`JdbcExecutionBehaviorTest` verifies non-`SELECT` result set execution, max-row truncation boundaries, and transaction control statements. `JdbcMetadataBehaviorTest` verifies stable metadata ordering. Agents with limited local test infrastructure can adopt the execution contract first and add the metadata contract later. + +For agents that need unavailable commercial or external drivers, use the fake execution contract: + +```java +class ExampleAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new ExampleAgent(); + } + + @Override + protected String resultSetSql() { + return "SHOW TABLES"; + } +} +``` + +`JdbcFakeExecutionBehaviorTest` injects a fake JDBC connection and verifies that `executeQuery` uses `Statement.execute`, reads the returned `ResultSet`, and does not fall back to `executeQuery` or `executeUpdate`. Use `testImplementation project(':test-support')` for this contract. + +## Review Checklist + +Before opening a PR or release: + +- `executeQuery` delegates to `JdbcExecutor.execute`. +- No local SQL-prefix classifier exists. +- No local result row cap exists unless there is a database-specific reason documented in code. +- Metadata methods use prepared statements for user-controlled schema/table inputs. +- `setSchemaSQL` quotes identifiers or returns an empty string. +- `disconnect` closes and clears the connection. +- `testConnection` closes its temporary connection. +- Root `build.gradle` conventions cover Java plugin, toolchain, JUnit, common/test-support dependencies, Shadow plugin, and archive name. +- Module `build.gradle` has correct driver dependencies, `Agent-Label`, `Main-Class`, and external driver flag. +- `settings.gradle`, `versions.json`, and README are updated together. +- At least one execution-path regression test exists. +- `python3 scripts/validate_agents.py` passes. +- `./gradlew test shadowJar --continue` passes. diff --git a/agents/docs/examples/jdbc-agent-template/README.md b/agents/docs/examples/jdbc-agent-template/README.md new file mode 100644 index 0000000000..18c8152a8b --- /dev/null +++ b/agents/docs/examples/jdbc-agent-template/README.md @@ -0,0 +1,20 @@ +# JDBC Agent Template + +Copy this directory when adding a new JDBC-based DBX agent. + +Recommended flow: + +1. Rename `template` paths, package names, class names, and artifact names to the new agent name. +2. Replace the driver class and JDBC URL builder. +3. Implement metadata SQL using the target database system catalogs or `DatabaseMetaData`. +4. Keep `executeQuery` delegated to `JdbcExecutor`. +5. Place the module under `drivers/`, then add it to `settings.gradle`, `versions.json`, and the root README support table. +6. Run `./gradlew ::test ::shadowJar`. + +To verify this template before copying it, run: + +```bash +./gradlew -p docs/examples/jdbc-agent-template clean test shadowJar +``` + +See `docs/agent-authoring.md` for the full checklist. diff --git a/agents/docs/examples/jdbc-agent-template/build.gradle b/agents/docs/examples/jdbc-agent-template/build.gradle new file mode 100644 index 0000000000..45bec3ac2d --- /dev/null +++ b/agents/docs/examples/jdbc-agent-template/build.gradle @@ -0,0 +1,64 @@ +plugins { + id 'java' + id 'com.gradleup.shadow' +} + +configure([project(':common'), project(':test-support')]) { + if (name == 'test-support') { + apply plugin: 'java-library' + } else { + apply plugin: 'java' + } + + java { + toolchain { + languageVersion = JavaLanguageVersion.of(8) + } + } + + dependencies { + testImplementation 'org.junit.jupiter:junit-jupiter:5.11.4' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher:1.11.4' + } + + tasks.withType(Test).configureEach { + useJUnitPlatform() + } +} + +dependencies { + implementation project(':common') + + // Bundled driver example: + // implementation 'com.example:example-jdbc:1.2.3' + + // External driver example: + // implementation fileTree(dir: 'libs', include: ['*.jar']) + + testImplementation 'org.junit.jupiter:junit-jupiter:5.11.4' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher:1.11.4' + testImplementation project(':test-support') +} + +java { + toolchain { + languageVersion = JavaLanguageVersion.of(21) + } +} + +tasks.named('shadowJar') { + archiveBaseName = 'dbx-agent-template' + archiveClassifier = '' + manifest { + attributes( + 'Agent-Label': 'Template DB', + // Add this only when users must provide the JDBC driver jar: + // 'Agent-External-Driver': 'true', + 'Main-Class': 'com.dbx.agent.template.TemplateAgent' + ) + } +} + +tasks.named('test') { + useJUnitPlatform() +} diff --git a/agents/docs/examples/jdbc-agent-template/settings.gradle b/agents/docs/examples/jdbc-agent-template/settings.gradle new file mode 100644 index 0000000000..0a2ac70843 --- /dev/null +++ b/agents/docs/examples/jdbc-agent-template/settings.gradle @@ -0,0 +1,21 @@ +pluginManagement { + repositories { + gradlePluginPortal() + mavenCentral() + } + plugins { + id 'com.gradleup.shadow' version '9.0.0-beta12' + } +} + +dependencyResolutionManagement { + repositoriesMode.set(RepositoriesMode.FAIL_ON_PROJECT_REPOS) + repositories { + mavenCentral() + } +} + +rootProject.name = 'jdbc-agent-template' +include 'common', 'test-support' +project(':common').projectDir = file('../../../common') +project(':test-support').projectDir = file('../../../test-support') diff --git a/agents/docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java b/agents/docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java new file mode 100644 index 0000000000..bca4d45a08 --- /dev/null +++ b/agents/docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java @@ -0,0 +1,27 @@ +package com.dbx.agent.template; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; + +public final class TemplateAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile TEMPLATE_PROFILE = new JdbcAgentProfile( + "com.example.jdbc.TemplateDriver", + "jdbc:template://{host}:{port}/{database}", + 1234 + ); + + public TemplateAgent() { + super(TEMPLATE_PROFILE); + } + + @Override + public String setSchemaSQL(String schema) { + return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + public static void main(String[] args) { + new JsonRpcServer(new TemplateAgent()).run(); + } +} diff --git a/agents/docs/examples/jdbc-agent-template/src/test/java/com/dbx/agent/template/TemplateAgentTest.java b/agents/docs/examples/jdbc-agent-template/src/test/java/com/dbx/agent/template/TemplateAgentTest.java new file mode 100644 index 0000000000..d074e67519 --- /dev/null +++ b/agents/docs/examples/jdbc-agent-template/src/test/java/com/dbx/agent/template/TemplateAgentTest.java @@ -0,0 +1,23 @@ +package com.dbx.agent.template; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.test.JdbcAgentFake; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TemplateAgentTest { + @Test + void executeQueryUsesStatementExecute() { + TemplateAgent agent = new TemplateAgent(); + TestSupport.setPrivateConnection(agent, JdbcAgentFake.connection()); + + agent.executeQuery("SHOW TABLES", null, new ExecuteQueryOptions()); + + Assertions.assertEquals( + List.of("setMaxRows:10001", "execute"), + JdbcAgentFake.calls + ); + } +} diff --git a/agents/docs/release-checklist.md b/agents/docs/release-checklist.md new file mode 100644 index 0000000000..dfd898cfd5 --- /dev/null +++ b/agents/docs/release-checklist.md @@ -0,0 +1,200 @@ +# Release Checklist + +Use this checklist before pushing release tags or publishing DBX agent jars. + +## 1. Confirm Scope + +- Review the recent commits: + +```bash +git log --oneline -10 +``` + +- Confirm the working tree is clean before tagging: + +```bash +git status --short +``` + +- Check that each intended agent change touches the right module directory, `common`, `test-support`, docs, or workflows. +- Avoid mixing unrelated driver updates, behavior changes, and release-only edits in the same commit. + +## 2. Local Verification + +On this workspace, use the local JDKs under `/private/tmp`: + +```bash +env JAVA_HOME=/private/tmp/dbx-jdk/jdk-21.0.11+10/Contents/Home \ + ./gradlew \ + -Dorg.gradle.java.installations.paths=/private/tmp/dbx-jdk/jdk-21.0.11+10/Contents/Home,/private/tmp/dbx-jdk8/jdk8/Contents/Home \ + test shadowJar --continue +``` + +Also run the lightweight validation gates: + +```bash +python3 -m unittest discover -s scripts -p '*_test.py' +python3 scripts/validate_agents.py +python3 scripts/validate_agent_jars.py +git diff --check +``` + +Expected result: + +- Python script tests pass. +- `Agent validation passed`. +- `Agent jar validation passed` after `shadowJar` has produced jars. +- `git diff --check` prints nothing. +- Gradle finishes with `BUILD SUCCESSFUL`. + +## 3. Agent Contract Checks + +For changed agent modules: + +- `executeQuery` delegates to `JdbcExecutor.execute`. +- No module-local SQL prefix classifier is present. +- No module-local result row cap is present. +- `Statement.execute(...)` is used for arbitrary user SQL. +- Metadata methods return stable ordering for schemas, tables, columns, indexes, foreign keys, and triggers. +- User-controlled schema and table names are passed through prepared statements or quoted with `JdbcIdentifiers`. +- `connect` stores one connection, `disconnect` closes and clears it, and `testConnection` closes its temporary connection. +- A behavior test exists: + - Use `JdbcExecutionBehaviorTest` and `JdbcMetadataBehaviorTest` for agents with a local/embedded test database. + - Use `JdbcFakeExecutionBehaviorTest` for agents that require unavailable external drivers. + - Use targeted `JdbcMetadataSqlFake` tests for metadata SQL that must interpolate quoted identifiers. + +## 4. Registry And Module Checks + +Run: + +```bash +python3 scripts/validate_agents.py +``` + +This checks: + +- `versions.json` keys match included agent modules in `settings.gradle`. +- Root Gradle conventions define agent archive names. +- Agent manifests define `Agent-Label`. +- Agent manifests define `Main-Class`. +- Agent `Main-Class` source exists and the built jar contains the matching `.class`. +- Forbidden legacy execution patterns are absent. +- Disallowed JVM source/build DSL residue is absent outside Gradle output directories. + +When adding or removing an agent, update these files together: + +- `settings.gradle` +- `versions.json` +- `README.md` +- Agent module `build.gradle` for driver dependencies and manifest attributes +- Root `build.gradle` only when the module needs non-standard shared build behavior + +## 5. Driver Packaging + +For bundled drivers: + +- Verify the dependency is redistributable. +- Prefer Maven dependencies over checked-in jars. + +For external drivers: + +- Use `implementation fileTree(dir: 'libs', include: ['*.jar'])`. +- Set the manifest attribute: + +```groovy +attributes( + 'Agent-External-Driver': 'true' +) +``` + +- Confirm release registry generation emits `external_driver_required: true`. + +Current external-driver agents include BigQuery and SunDB. + +## 6. JRE Selection + +Most agents are built for the default JRE key `21`, backed by JDK 21 in the release workflow. + +Special case: + +- `oracle-10g` uses JRE key `8`. + +If another agent needs a different runtime: + +- Update the release workflow JRE detection logic. +- Document why in the module or release notes. +- Verify DBX can download the matching runtime artifact. + +## 7. CI Expectations + +The CI workflow runs on `main` and pull requests: + +```bash +python3 -m unittest discover -s scripts -p '*_test.py' +python3 scripts/validate_agents.py +./gradlew test shadowJar --continue +python3 scripts/validate_agent_jars.py +``` + +Do not tag a release while CI is failing on `main`. + +## 8. Release Tag Flow + +Release workflow runs on tags matching `agents-v*`. + +Before tagging: + +```bash +git status --short +git log --oneline -5 +git tag --list 'agents-v*' --sort=-creatordate | head +``` + +Choose a new tag that does not already exist locally or on GitHub. For example: + +```bash +RELEASE_TAG=agents-v0.3.0 +git tag --list "$RELEASE_TAG" +git ls-remote --tags origin "$RELEASE_TAG" +``` + +Both commands should print nothing before you create the tag. + +Create and push the tag: + +```bash +git tag "$RELEASE_TAG" +git push origin main +git push origin "$RELEASE_TAG" +``` + +The release workflow will: + +- Bump changed module versions in `versions.json`. +- Build all agent shadow jars. +- Build/download JRE artifacts. +- Generate `agent-registry.json`. +- Create a GitHub release with jars, JRE archives, and registry. + +## 9. Post-Release Verification + +After the GitHub release finishes: + +- Download or inspect `agent-registry.json`. +- Confirm every expected agent appears under `drivers`. +- Confirm labels preserve spaces, for example `Google BigQuery`. +- Confirm `oracle-10g` uses JRE key `8`. +- Confirm other agents use JRE key `21`. +- Confirm `external_driver_required` is correct. +- Confirm every jar URL, sha256, and size is present. +- Spot-check at least one agent jar manifest: + +```bash +unzip -p dbx-agent-h2.jar META-INF/MANIFEST.MF +``` + +## 10. Known Follow-Ups + +- Gradle currently reports deprecated features that will be incompatible with Gradle 10. Use `--warning-mode all` in a separate cleanup task. +- Commercial or external-driver agents still need manual smoke tests with real driver jars and databases. +- Add real containerized smoke tests only one database family at a time. diff --git a/agents/drivers/access/build.gradle b/agents/drivers/access/build.gradle new file mode 100644 index 0000000000..c022da63b6 --- /dev/null +++ b/agents/drivers/access/build.gradle @@ -0,0 +1,9 @@ +dependencies { + implementation 'io.github.spannm:ucanaccess:5.1.5' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Microsoft Access', 'Main-Class': 'com.dbx.agent.access.AccessAgent') + } +} diff --git a/agents/drivers/access/src/main/java/com/dbx/agent/access/AccessAgent.java b/agents/drivers/access/src/main/java/com/dbx/agent/access/AccessAgent.java new file mode 100644 index 0000000000..2f0d829e54 --- /dev/null +++ b/agents/drivers/access/src/main/java/com/dbx/agent/access/AccessAgent.java @@ -0,0 +1,366 @@ +package com.dbx.agent.access; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class AccessAgent extends BaseDatabaseAgent { + private Connection connection; + private String databaseFile = ""; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("net.ucanaccess.jdbc.UcanaccessDriver"); + String file = databasePath(params); + String url = jdbcUrl(params, true); + connection = DriverManager.getConnection(url, params.getUsername(), params.getPassword()); + databaseFile = file; + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("net.ucanaccess.jdbc.UcanaccessDriver"); + String file = databasePath(params); + if (file.isBlank() || !Files.exists(Path.of(file))) { + return false; + } + try (Connection conn = DriverManager.getConnection(jdbcUrl(params, false), params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + String name = "default"; + if (!databaseFile.isBlank()) { + Path fileName = Path.of(databaseFile).getFileName(); + if (fileName != null && !fileName.toString().isBlank()) { + name = fileName.toString(); + } + } + return Collections.singletonList(new DatabaseInfo(name)); + } + + @Override + public List listSchemas() { + return Collections.emptyList(); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + DatabaseMetaData meta = requireConnected().getMetaData(); + List result = new ArrayList<>(); + try (ResultSet rs = meta.getTables(null, null, "%", new String[]{"TABLE", "VIEW"})) { + while (rs.next()) { + String name = rs.getString("TABLE_NAME"); + if (name == null || startsWithIgnoreCase(name, "MSys")) { + continue; + } + result.add(new TableInfo(name, normalizeTableType(rs.getString("TABLE_TYPE")), null)); + } + } + result.sort(Comparator.comparing(table -> table.getName().toLowerCase(Locale.ROOT))); + return result; + }); + } + + @Override + public List listObjects(String schema) { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), null, table.getComment())); + } + return result; + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + DatabaseMetaData meta = requireConnected().getMetaData(); + Set primaryKeys = primaryKeyColumns(meta, table); + List result = new ArrayList<>(); + try (ResultSet rs = meta.getColumns(null, null, table, "%")) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + name, + typeName(rs), + rs.getInt("NULLABLE") != DatabaseMetaData.columnNoNulls, + rs.getString("COLUMN_DEF"), + primaryKeys.contains(name), + null, + rs.getString("REMARKS"), + intObject(rs, "COLUMN_SIZE"), + intObject(rs, "DECIMAL_DIGITS"), + intObject(rs, "COLUMN_SIZE") + )); + } + } + result.sort(Comparator.comparingInt(column -> columnOrder(meta, table, column.getName()))); + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + DatabaseMetaData meta = requireConnected().getMetaData(); + Set primaryKeys = primaryKeyColumns(meta, table); + Map> columnsByIndex = new LinkedHashMap<>(); + Map uniqueByIndex = new LinkedHashMap<>(); + Map primaryByIndex = new LinkedHashMap<>(); + Map typeByIndex = new LinkedHashMap<>(); + + try (ResultSet rs = meta.getIndexInfo(null, null, table, false, false)) { + while (rs.next()) { + if (rs.getShort("TYPE") == DatabaseMetaData.tableIndexStatistic) { + continue; + } + String indexName = rs.getString("INDEX_NAME"); + String columnName = rs.getString("COLUMN_NAME"); + if (indexName == null || columnName == null) { + continue; + } + short ordinal = rs.getShort("ORDINAL_POSITION"); + + columnsByIndex.computeIfAbsent(indexName, ignored -> new ArrayList<>()).add(new OrdinalColumn(ordinal, columnName)); + uniqueByIndex.put(indexName, !rs.getBoolean("NON_UNIQUE")); + primaryByIndex.put(indexName, primaryKeys.contains(columnName) || indexName.equalsIgnoreCase("PrimaryKey")); + typeByIndex.put(indexName, indexTypeName(rs.getShort("TYPE"))); + } + } + + List result = new ArrayList<>(); + for (Map.Entry> entry : columnsByIndex.entrySet()) { + String name = entry.getKey(); + List pairs = entry.getValue(); + pairs.sort(Comparator.comparingInt(pair -> pair.ordinal)); + List columns = new ArrayList<>(); + for (OrdinalColumn pair : pairs) { + columns.add(pair.column); + } + result.add(new IndexInfo( + name, + columns, + uniqueByIndex.getOrDefault(name, false), + primaryByIndex.getOrDefault(name, false), + null, + typeByIndex.get(name), + null, + null + )); + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + DatabaseMetaData meta = requireConnected().getMetaData(); + List result = new ArrayList<>(); + try (ResultSet rs = meta.getImportedKeys(null, null, table)) { + while (rs.next()) { + String name = rs.getString("FK_NAME"); + result.add(new ForeignKeyInfo( + name == null ? "" : name, + rs.getString("FKCOLUMN_NAME"), + rs.getString("PKTABLE_NAME"), + rs.getString("PKCOLUMN_NAME") + )); + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + JdbcExecutor.INSTANCE::defaultResultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return ""; + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Set primaryKeyColumns(DatabaseMetaData meta, String table) throws Exception { + Set result = new LinkedHashSet<>(); + try (ResultSet rs = meta.getPrimaryKeys(null, null, table)) { + while (rs.next()) { + result.add(rs.getString("COLUMN_NAME")); + } + } + return result; + } + + private int columnOrder(DatabaseMetaData meta, String table, String column) { + return unchecked(() -> { + try (ResultSet rs = meta.getColumns(null, null, table, column)) { + return rs.next() ? rs.getInt("ORDINAL_POSITION") : Integer.MAX_VALUE; + } + }); + } + + private static String databasePath(ConnectParams params) { + if (startsWithIgnoreCase(params.getConnection_string(), "jdbc:ucanaccess:")) { + return stripJdbcPrefix(params.getConnection_string()); + } + if (startsWithIgnoreCase(params.getDatabase(), "jdbc:ucanaccess:")) { + return stripJdbcPrefix(params.getDatabase()); + } + if (!params.getHost().isBlank()) { + return params.getHost(); + } + return params.getDatabase(); + } + + static String jdbcUrl(ConnectParams params, boolean createIfMissing) { + String rawUrl; + if (startsWithIgnoreCase(params.getConnection_string(), "jdbc:ucanaccess:")) { + rawUrl = params.getConnection_string(); + } else if (startsWithIgnoreCase(params.getDatabase(), "jdbc:ucanaccess:")) { + rawUrl = params.getDatabase(); + } else { + rawUrl = "jdbc:ucanaccess://" + databasePath(params); + } + rawUrl = appendUCanAccessParams(rawUrl, params.getUrl_params()); + + if (!createIfMissing || Files.exists(Path.of(databasePath(params)))) { + return rawUrl; + } + if (containsIgnoreCase(rawUrl, "newDatabaseVersion=")) { + return rawUrl; + } + return rawUrl + ";newDatabaseVersion=V2010"; + } + + private static String appendUCanAccessParams(String url, String urlParams) { + String params = trimUCanAccessParams(urlParams); + if (params.isEmpty()) { + return url; + } + return url + (url.endsWith(";") ? "" : ";") + params; + } + + private static String trimUCanAccessParams(String urlParams) { + String value = urlParams == null ? "" : urlParams.trim(); + while (value.startsWith("?") || value.startsWith("&") || value.startsWith(";")) { + value = value.substring(1); + } + return value; + } + + private static String normalizeTableType(String value) { + if (value != null && value.toUpperCase(Locale.ROOT).equals("BASE TABLE")) { + return "TABLE"; + } + return value == null ? "TABLE" : value; + } + + private static String indexTypeName(short value) { + if (value == DatabaseMetaData.tableIndexClustered) { + return "CLUSTERED"; + } + if (value == DatabaseMetaData.tableIndexHashed) { + return "HASHED"; + } + if (value == DatabaseMetaData.tableIndexOther) { + return "OTHER"; + } + return ""; + } + + private static String typeName(ResultSet rs) throws Exception { + String value = rs.getString("TYPE_NAME"); + return value == null ? Integer.toString(rs.getInt("DATA_TYPE")) : value; + } + + private static Integer intObject(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number number ? number.intValue() : null; + } + + private static String stripJdbcPrefix(String value) { + return substringBefore(value.replaceFirst("(?i)^jdbc:ucanaccess://", ""), ';'); + } + + private static String substringBefore(String value, char delimiter) { + int index = value.indexOf(delimiter); + return index >= 0 ? value.substring(0, index) : value; + } + + private static boolean startsWithIgnoreCase(String value, String prefix) { + return value.regionMatches(true, 0, prefix, 0, prefix.length()); + } + + private static boolean containsIgnoreCase(String value, String needle) { + return value.toLowerCase(Locale.ROOT).contains(needle.toLowerCase(Locale.ROOT)); + } + + private record OrdinalColumn(short ordinal, String column) { + } + + public static void main(String[] args) { + new JsonRpcServer(new AccessAgent()).run(); + } +} diff --git a/agents/drivers/access/src/test/java/com/dbx/agent/access/AccessAgentTest.java b/agents/drivers/access/src/test/java/com/dbx/agent/access/AccessAgentTest.java new file mode 100644 index 0000000000..b49ab56c19 --- /dev/null +++ b/agents/drivers/access/src/test/java/com/dbx/agent/access/AccessAgentTest.java @@ -0,0 +1,113 @@ +package com.dbx.agent.access; + +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.QueryPageOptions; +import com.dbx.agent.test.JdbcConnectedAgentTest; +import java.nio.file.Path; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class AccessAgentTest extends JdbcConnectedAgentTest { + @TempDir + Path tempDir; + + @Override + protected DatabaseAgent createConnectedAgent(String databaseName) { + Path file = tempDir.resolve(databaseName + ".accdb"); + AccessAgent agent = new AccessAgent(); + agent.connect(new ConnectParams("", 0, file.toString(), "", "", "", "", false)); + return agent; + } + + @Test + void connectsToAnAccessFileAndReturnsTableMetadata() { + withAgent("metadata", agent -> { + agent.executeQuery( + "CREATE TABLE People (ID COUNTER PRIMARY KEY, Name TEXT(64), Active YESNO)", + null, + new ExecuteQueryOptions() + ); + agent.executeQuery("INSERT INTO People (Name, Active) VALUES ('Ada', TRUE)", null, new ExecuteQueryOptions()); + + var tables = agent.listTables(""); + var columns = agent.getColumns("", "People"); + var result = agent.executeQuery("SELECT Name, Active FROM People", null, new ExecuteQueryOptions()); + + Assertions.assertTrue(tables.stream().anyMatch(table -> table.getName().equals("People"))); + Assertions.assertEquals(List.of("ID", "Name", "Active"), columns.stream().map(column -> column.getName()).toList()); + Assertions.assertTrue(columns.stream() + .filter(column -> column.getName().equals("ID")) + .findFirst() + .orElseThrow() + .getIs_primary_key()); + Assertions.assertEquals(List.of("Name", "Active"), result.getColumns()); + Assertions.assertEquals("Ada", result.getRows().get(0).get(0)); + Assertions.assertEquals(true, result.getRows().get(0).get(1)); + }); + } + + @Test + void returnsEmptySchemaListForAccessFiles() { + withAgent("schemas", agent -> { + Assertions.assertEquals(List.of(), agent.listSchemas()); + Assertions.assertFalse(agent.testConnection( + new ConnectParams("", 0, tempDir.resolve("missing.accdb").toString(), "", "", "", "", false) + )); + }); + } + + @Test + void appendsUcanaccessUrlParams() { + Path file = tempDir.resolve("params.accdb"); + ConnectParams params = new ConnectParams("", 0, file.toString(), "", "", "memory=false;ignoreCase=true", "", false); + + Assertions.assertEquals( + "jdbc:ucanaccess://" + file + ";memory=false;ignoreCase=true;newDatabaseVersion=V2010", + AccessAgent.jdbcUrl(params, true) + ); + } + + @Test + void trimsUcanaccessUrlParamSeparators() { + Path file = tempDir.resolve("params-existing.accdb"); + ConnectParams params = new ConnectParams("", 0, "jdbc:ucanaccess://" + file + ";", "", "", ";memory=false", "", false); + + Assertions.assertEquals( + "jdbc:ucanaccess://" + file + ";memory=false;newDatabaseVersion=V2010", + AccessAgent.jdbcUrl(params, true) + ); + } + + @Test + void supportsLimitOffsetSqlAndCursorResultPages() { + withAgent("pages", agent -> { + agent.executeQuery("CREATE TABLE Items (ID COUNTER PRIMARY KEY, Value INTEGER)", null, new ExecuteQueryOptions()); + for (int i = 1; i <= 5; i++) { + agent.executeQuery("INSERT INTO Items (Value) VALUES (" + i + ")", null, new ExecuteQueryOptions()); + } + + var limited = agent.executeQuery( + "SELECT [Value] FROM [Items] ORDER BY [ID] LIMIT 2 OFFSET 1", + null, + new ExecuteQueryOptions() + ); + var firstPage = agent.executeQueryPage( + "SELECT Value FROM Items ORDER BY ID", + null, + new QueryPageOptions(2, null, 5) + ); + String sessionId = firstPage.getSession_id(); + Assertions.assertNotNull(sessionId); + var secondPage = agent.fetchQueryPage(sessionId, 2); + + Assertions.assertEquals(List.of(2, 3), limited.getRows().stream().map(row -> row.get(0)).toList()); + Assertions.assertEquals(List.of(1, 2), firstPage.getRows().stream().map(row -> row.get(0)).toList()); + Assertions.assertEquals(true, firstPage.getHas_more()); + Assertions.assertEquals(List.of(3, 4), secondPage.getRows().stream().map(row -> row.get(0)).toList()); + }); + } +} diff --git a/agents/drivers/bigquery/build.gradle b/agents/drivers/bigquery/build.gradle new file mode 100644 index 0000000000..ae0da1b389 --- /dev/null +++ b/agents/drivers/bigquery/build.gradle @@ -0,0 +1,13 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) +} + +tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'Google BigQuery', + 'Agent-External-Driver': 'true', + 'Main-Class': 'com.dbx.agent.bigquery.BigQueryAgent' + ) + } +} diff --git a/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java b/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java new file mode 100644 index 0000000000..11feaefd29 --- /dev/null +++ b/agents/drivers/bigquery/src/main/java/com/dbx/agent/bigquery/BigQueryAgent.java @@ -0,0 +1,145 @@ +package com.dbx.agent.bigquery; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +public final class BigQueryAgent extends AbstractJdbcAgent { + + @Override + protected String driverClass() { + return "com.simba.googlebigquery.jdbc.Driver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SELECT schema_name FROM INFORMATION_SCHEMA.SCHEMATA ORDER BY schema_name")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String quotedSchema = JdbcIdentifiers.INSTANCE.backtick(schema); + String sql = "SELECT table_name, table_type FROM " + quotedSchema + ".INFORMATION_SCHEMA.TABLES ORDER BY table_name"; + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql); + ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("table_name"), + normalizeTableType(rs.getString("table_type")), + null + )); + } + } + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String quotedSchema = JdbcIdentifiers.INSTANCE.backtick(schema); + String sql = "SELECT column_name, data_type, is_nullable " + + "FROM " + quotedSchema + ".INFORMATION_SCHEMA.COLUMNS " + + "WHERE table_name = ? " + + "ORDER BY ordinal_position"; + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ColumnInfo( + rs.getString("column_name"), + rs.getString("data_type"), + "YES".equals(rs.getString("is_nullable")), + null, + false, + null, + null, + null, + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return ""; + } + + static String buildUrl(ConnectParams params) { + String url = "jdbc:bigquery://" + params.getHost() + ":" + params.getPort() + ";ProjectId=" + params.getDatabase(); + String extraParams = trimSemicolons(params.getUrl_params()); + return extraParams.isBlank() ? url : url + ";" + extraParams; + } + + private static String trimSemicolons(String value) { + if (value == null) { + return ""; + } + return value.trim().replaceAll("^;+", "").replaceAll(";+\\s*$", ""); + } + + private static String normalizeTableType(String type) { + return "BASE TABLE".equals(type) ? "TABLE" : type; + } + + public static void main(String[] args) { + new JsonRpcServer(new BigQueryAgent()).run(); + } +} diff --git a/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java b/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java new file mode 100644 index 0000000000..d0340b5571 --- /dev/null +++ b/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentMetadataTest.java @@ -0,0 +1,49 @@ +package com.dbx.agent.bigquery; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.test.JdbcAgentFake; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import java.sql.Connection; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class BigQueryAgentMetadataTest { + @Test + void quotesSchemaIdentifiersInMetadataSql() { + BigQueryAgent agent = new BigQueryAgent(); + Connection fake = JdbcMetadataSqlFake.connection(); + TestSupport.setPrivateConnection(agent, fake); + + agent.listTables("bad`schema"); + agent.getColumns("bad`schema", "sample"); + + Assertions.assertTrue( + JdbcMetadataSqlFake.statements.stream() + .anyMatch(statement -> statement.contains("FROM `bad``schema`.INFORMATION_SCHEMA.TABLES")) + ); + Assertions.assertTrue( + JdbcMetadataSqlFake.statements.stream() + .anyMatch(statement -> statement.contains("FROM `bad``schema`.INFORMATION_SCHEMA.COLUMNS")) + ); + } + + @Test + void buildUrlAppendsAuthenticationUrlParams() { + ConnectParams params = new ConnectParams( + "https://www.googleapis.com/bigquery/v2", + 443, + "demo-project", + "", + "", + "OAuthType=0;OAuthServiceAcctEmail=svc@demo.iam.gserviceaccount.com;OAuthPvtKeyPath=C:\\keys\\demo.json", + "", + false + ); + + Assertions.assertEquals( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId=demo-project;OAuthType=0;OAuthServiceAcctEmail=svc@demo.iam.gserviceaccount.com;OAuthPvtKeyPath=C:\\keys\\demo.json", + BigQueryAgent.buildUrl(params) + ); + } +} diff --git a/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentTest.java b/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentTest.java new file mode 100644 index 0000000000..147d1f6c56 --- /dev/null +++ b/agents/drivers/bigquery/src/test/java/com/dbx/agent/bigquery/BigQueryAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.bigquery; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class BigQueryAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new BigQueryAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL dataset.proc()"; + } +} diff --git a/agents/drivers/cassandra/build.gradle b/agents/drivers/cassandra/build.gradle new file mode 100644 index 0000000000..38a8d6ffd4 --- /dev/null +++ b/agents/drivers/cassandra/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.ing.data:cassandra-jdbc-wrapper:4.12.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Apache Cassandra', 'Main-Class': 'com.dbx.agent.cassandra.CassandraAgent') + } +} diff --git a/agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java b/agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java new file mode 100644 index 0000000000..5c6f630501 --- /dev/null +++ b/agents/drivers/cassandra/src/main/java/com/dbx/agent/cassandra/CassandraAgent.java @@ -0,0 +1,185 @@ +package com.dbx.agent.cassandra; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +public final class CassandraAgent extends AbstractJdbcAgent { + private static final Pattern TARGET_PATTERN = Pattern.compile("target[\"']?\\s*[:=]\\s*[\"']?([\\w]+)"); + + @Override + protected String driverClass() { + return "com.ing.data.cassandra.jdbc.CassandraDriver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT keyspace_name FROM system_schema.keyspaces"; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + result.sort(Comparator.comparing(DatabaseInfo::getName)); + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT table_name FROM system_schema.tables WHERE keyspace_name = ?"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), "TABLE", null)); + } + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT column_name, type, kind FROM system_schema.columns WHERE keyspace_name = ? AND table_name = ?"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String kind = coalesce(rs.getString("kind")); + boolean isPrimaryKey = "partition_key".equals(kind) || "clustering".equals(kind); + result.add(new ColumnInfo( + rs.getString("column_name"), + coalesce(rs.getString("type"), "unknown"), + !isPrimaryKey, + null, + isPrimaryKey, + kind.trim().isEmpty() ? null : kind, + null, + null, + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT index_name, options FROM system_schema.indexes WHERE keyspace_name = ? AND table_name = ?"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String indexName = coalesce(rs.getString("index_name")); + String options = coalesce(rs.getString("options")); + result.add(new IndexInfo( + indexName, + targetColumns(options), + false, + false, + null, + null, + null, + null + )); + } + } + } + result.sort(Comparator.comparing(IndexInfo::getName)); + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:cassandra://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private static List targetColumns(String options) { + Matcher matcher = TARGET_PATTERN.matcher(options); + if (!matcher.find()) { + return Collections.emptyList(); + } + return Collections.singletonList(matcher.group(1)); + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + private static String coalesce(String value, String fallback) { + return value == null ? fallback : value; + } + + public static void main(String[] args) { + new JsonRpcServer(new CassandraAgent()).run(); + } +} diff --git a/agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java b/agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java new file mode 100644 index 0000000000..484ed1e024 --- /dev/null +++ b/agents/drivers/cassandra/src/test/java/com/dbx/agent/cassandra/CassandraAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.cassandra; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class CassandraAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new CassandraAgent(); + } + + @Override + protected String resultSetSql() { + return "LIST ROLES"; + } +} diff --git a/agents/drivers/dameng/build.gradle b/agents/drivers/dameng/build.gradle new file mode 100644 index 0000000000..6a80eb99a5 --- /dev/null +++ b/agents/drivers/dameng/build.gradle @@ -0,0 +1,13 @@ +dependencies { + // Local override: place a newer DmJdbcDriver*.jar in libs/ for versions + // not yet published to Maven Central (the fileTree takes priority over Maven). + implementation fileTree(dir: 'libs', include: ['*.jar']) + // DmJdbcDriver8 supports JDK 8+ (compatible with JRE 21 used by DAMENG agent). + implementation 'com.dameng:DmJdbcDriver8:8.1.5.45' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': '达梦 DM8', 'Main-Class': 'com.dbx.agent.dameng.DamengAgent') + } +} diff --git a/agents/drivers/dameng/libs/.gitkeep b/agents/drivers/dameng/libs/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java b/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java new file mode 100644 index 0000000000..3ff1919edd --- /dev/null +++ b/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java @@ -0,0 +1,758 @@ +package com.dbx.agent.dameng; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryPageOptions; +import com.dbx.agent.QueryPageResult; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.io.PrintStream; +import java.sql.Blob; +import java.sql.Clob; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLXML; +import java.sql.Statement; +import java.sql.Types; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class DamengAgent extends BaseDatabaseAgent { + private static final String AGENT_VERSION = "9999.06.04.1-fix-default"; + private static final Set SYSTEM_USERS = Set.of( + "SYS", "SYSAUDITOR", "SYSSSO", "CTISYS", + "SYS_DBA", "_SYS_STATISTICS", "SYS_PHM" + ); + + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + withSuppressedStdout(() -> { + Class.forName("dm.jdbc.driver.DmDriver"); + connection = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword()); + }); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + return withSuppressedStdout(() -> { + Class.forName("dm.jdbc.driver.DmDriver"); + try (Connection conn = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + }); + } + + /** + * The DM JDBC driver writes a banner to {@code System.out} during + * {@code Class.forName} / driver initialization. This corrupts the + * JSON-RPC stdout protocol. Temporarily redirect {@code System.out} + * to {@code System.err} so driver output lands on stderr instead. + */ + private static T withSuppressedStdout(ThrowingSupplier action) throws Exception { + PrintStream originalOut = System.out; + try { + System.setOut(System.err); + return action.get(); + } finally { + System.setOut(originalOut); + } + } + + private static void withSuppressedStdout(ThrowingRunnable action) throws Exception { + PrintStream originalOut = System.out; + try { + System.setOut(System.err); + action.run(); + } finally { + System.setOut(originalOut); + } + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + String placeholders = String.join(",", SYSTEM_USERS.stream().map(user -> "?").toList()); + String sql = "SELECT USERNAME FROM ALL_USERS WHERE USERNAME NOT IN (" + placeholders + ") ORDER BY USERNAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + int index = 1; + for (String user : SYSTEM_USERS) { + stmt.setString(index++, user); + } + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT DISTINCT OWNER FROM ALL_OBJECTS ORDER BY OWNER"; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(rs.getString(1)); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = """ + SELECT TABLE_NAME, 'TABLE' AS TABLE_TYPE, COMMENTS + FROM ALL_TAB_COMMENTS + WHERE OWNER = ? AND TABLE_TYPE = 'TABLE' + UNION ALL + SELECT TABLE_NAME, 'VIEW' AS TABLE_TYPE, COMMENTS + FROM ALL_TAB_COMMENTS + WHERE OWNER = ? AND TABLE_TYPE = 'VIEW' + ORDER BY 1 + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + String sql = """ + SELECT OBJECT_NAME, OBJECT_TYPE FROM ALL_OBJECTS + WHERE OWNER = ? AND OBJECT_TYPE IN ('PROCEDURE', 'FUNCTION') + ORDER BY CASE OBJECT_TYPE WHEN 'PROCEDURE' THEN 0 ELSE 1 END, OBJECT_NAME + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String dbmsType = switch (objectType.toUpperCase(Locale.ROOT)) { + case "VIEW" -> "VIEW"; + case "PROCEDURE" -> "PROCEDURE"; + case "FUNCTION" -> "FUNCTION"; + default -> throw new IllegalArgumentException("Unsupported object type: " + objectType); + }; + String source; + String sql = "SELECT DBMS_METADATA.GET_DDL(?, ?, ?) FROM DUAL"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, dbmsType); + stmt.setString(2, name); + stmt.setString(3, schema); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public String getTableDdl(String schema, String table) { + return unchecked(() -> { + String sql = "SELECT DBMS_METADATA.GET_DDL(?, ?, ?) FROM DUAL"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, "TABLE"); + stmt.setString(2, table); + stmt.setString(3, schema); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + return appendTableAndColumnComments(coalesce(rs.getString(1)), schema, table); + } + } + } + throw new IllegalArgumentException("Table not found: " + schema + "." + table); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set pkColumns = new java.util.HashSet<>(); + String pkSql = """ + SELECT cols.COLUMN_NAME FROM ALL_CONS_COLUMNS cols + JOIN ALL_CONSTRAINTS cons ON cols.CONSTRAINT_NAME = cons.CONSTRAINT_NAME AND cols.OWNER = cons.OWNER + WHERE cons.CONSTRAINT_TYPE = 'P' AND cons.OWNER = ? AND cons.TABLE_NAME = ? + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(pkSql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + pkColumns.add(rs.getString(1)); + } + } + } + + List result = new ArrayList<>(); + // DATA_DEFAULT is a LONG column — it must be selected first and read first + // in JDBC, otherwise the data is truncated. + String colSql = """ + SELECT c.DATA_DEFAULT, + c.COLUMN_NAME, + c.DATA_TYPE, + c.NULLABLE, + c.DATA_PRECISION, + c.DATA_SCALE, + c.DATA_LENGTH, + c.CHAR_LENGTH, + cc.COMMENTS + FROM ALL_TAB_COLUMNS c + LEFT JOIN ALL_COL_COMMENTS cc + ON cc.OWNER = c.OWNER + AND cc.TABLE_NAME = c.TABLE_NAME + AND cc.COLUMN_NAME = c.COLUMN_NAME + WHERE c.OWNER = ? AND c.TABLE_NAME = ? + ORDER BY c.COLUMN_ID + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(colSql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + // DATA_DEFAULT is a LONG — must be read first, before all other columns. + String dataDefault = readLongColumn(rs, "DATA_DEFAULT"); + String name = rs.getString("COLUMN_NAME"); + String baseType = rs.getString("DATA_TYPE"); + Integer numPrec = intObject(rs, "DATA_PRECISION"); + Integer numScale = intObject(rs, "DATA_SCALE"); + Integer dataLen = intObject(rs, "DATA_LENGTH"); + Integer charLen = intObject(rs, "CHAR_LENGTH"); + String dataType = formatDataType(baseType, numPrec, numScale, dataLen, charLen); + + result.add(new ColumnInfo( + name, + dataType, + "Y".equals(rs.getString("NULLABLE")), + dataDefault, + pkColumns.contains(name), + null, + rs.getString("COMMENTS"), + numPrec, + numScale, + charLen + )); + } + } + } + fillMissingColumnComments(schema, table, result); + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = """ + SELECT i.INDEX_NAME, + LISTAGG(ic.COLUMN_NAME, ',') WITHIN GROUP (ORDER BY ic.COLUMN_POSITION) AS COLUMNS, + i.UNIQUENESS, + CASE WHEN c.CONSTRAINT_TYPE = 'P' THEN 1 ELSE 0 END AS IS_PK, + i.INDEX_TYPE + FROM ALL_INDEXES i + JOIN ALL_IND_COLUMNS ic ON i.INDEX_NAME = ic.INDEX_NAME AND i.OWNER = ic.INDEX_OWNER AND i.TABLE_OWNER = ic.TABLE_OWNER + LEFT JOIN ALL_CONSTRAINTS c ON i.INDEX_NAME = c.INDEX_NAME AND i.TABLE_OWNER = c.OWNER + AND c.CONSTRAINT_TYPE = 'P' + WHERE i.TABLE_OWNER = ? AND i.TABLE_NAME = ? + GROUP BY i.INDEX_NAME, i.UNIQUENESS, c.CONSTRAINT_TYPE, i.INDEX_TYPE + ORDER BY i.INDEX_NAME + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new IndexInfo( + rs.getString(1), + splitNonEmpty(coalesce(rs.getString(2)), ","), + "UNIQUE".equals(rs.getString(3)), + "1".equals(rs.getString(4)), + null, + rs.getString(5), + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = """ + SELECT c.CONSTRAINT_NAME, cc.COLUMN_NAME, rc.TABLE_NAME, rcc.COLUMN_NAME + FROM ALL_CONSTRAINTS c + JOIN ALL_CONS_COLUMNS cc ON c.CONSTRAINT_NAME = cc.CONSTRAINT_NAME AND c.OWNER = cc.OWNER + JOIN ALL_CONSTRAINTS rc ON c.R_CONSTRAINT_NAME = rc.CONSTRAINT_NAME AND c.R_OWNER = rc.OWNER + JOIN ALL_CONS_COLUMNS rcc ON rc.CONSTRAINT_NAME = rcc.CONSTRAINT_NAME AND rc.OWNER = rcc.OWNER + WHERE c.CONSTRAINT_TYPE = 'R' AND c.OWNER = ? AND c.TABLE_NAME = ? + ORDER BY c.CONSTRAINT_NAME + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString(1), + rs.getString(2), + rs.getString(3), + rs.getString(4) + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = """ + SELECT TRIGGER_NAME, TRIGGERING_EVENT, '' AS TRIGGER_TYPE + FROM ALL_TRIGGERS + WHERE OWNER = ? AND TABLE_NAME = ? + ORDER BY TRIGGER_NAME + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::stringResultValue + ); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + return JdbcExecutor.INSTANCE.executePage( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options, + this::stringResultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Object stringResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = switch (sqlType) { + case Types.BIGINT -> rs.getLong(index); + case Types.INTEGER, Types.SMALLINT, Types.TINYINT -> rs.getInt(index); + case Types.FLOAT, Types.REAL -> rs.getFloat(index); + case Types.DOUBLE -> rs.getDouble(index); + case Types.DECIMAL, Types.NUMERIC -> rs.getBigDecimal(index); + case Types.BOOLEAN, Types.BIT -> rs.getBoolean(index); + case Types.CHAR, Types.VARCHAR, Types.LONGVARCHAR, + Types.NCHAR, Types.NVARCHAR, Types.LONGNVARCHAR, + Types.CLOB, Types.NCLOB -> rs.getString(index); + case Types.BINARY, Types.VARBINARY, Types.LONGVARBINARY, + Types.BLOB -> bytesToHex(rs.getBytes(index)); + case Types.SQLXML -> sqlXmlToString(rs.getSQLXML(index)); + default -> normalizeResultValue(rs.getObject(index)); + }; + return rs.wasNull() ? null : value; + }); + } + + private static Object normalizeResultValue(Object value) { + if (value == null) return null; + if (value instanceof Clob) { + Clob clob = (Clob) value; + return unchecked(() -> clob.getSubString(1, Math.toIntExact(clob.length()))); + } + if (value instanceof Blob) { + Blob blob = (Blob) value; + return unchecked(() -> bytesToHex(blob.getBytes(1, Math.toIntExact(blob.length())))); + } + if (value instanceof SQLXML) { + SQLXML sqlxml = (SQLXML) value; + return unchecked(sqlxml::getString); + } + if (value instanceof byte[]) { + return bytesToHex((byte[]) value); + } + return value instanceof Number || value instanceof Boolean ? value : value.toString(); + } + + private static String sqlXmlToString(SQLXML value) { + return value == null ? null : unchecked(value::getString); + } + + private static String bytesToHex(byte[] bytes) { + if (bytes == null) { + return null; + } + StringBuilder result = new StringBuilder(bytes.length * 2 + 2); + result.append("0x"); + for (byte b : bytes) { + result.append(Character.forDigit((b >> 4) & 0xF, 16)); + result.append(Character.forDigit(b & 0xF, 16)); + } + return result.toString(); + } + + private static String buildUrl(ConnectParams params) { + String database = params.getDatabase() == null ? "" : params.getDatabase().trim(); + String suffix = database.isEmpty() ? "" : "/" + database; + return "jdbc:dm://" + params.getHost() + ":" + params.getPort() + suffix; + } + + private static String formatDataType(String base, Integer numPrec, Integer numScale, Integer dataLen, Integer charLen) { + return switch (base.toUpperCase(Locale.ROOT)) { + case "VARCHAR2", "NVARCHAR2", "VARCHAR", "CHAR", "NCHAR" -> { + Integer length = charLen != null ? charLen : dataLen; + yield length != null ? base + "(" + length + ")" : base; + } + case "NUMBER", "NUMERIC", "DECIMAL" -> { + if (numPrec != null && numScale != null && numScale > 0) { + yield base + "(" + numPrec + "," + numScale + ")"; + } + yield numPrec != null && numPrec > 0 ? base + "(" + numPrec + ")" : base; + } + case "RAW" -> dataLen != null ? "RAW(" + dataLen + ")" : "RAW"; + default -> base; + }; + } + + private static Integer intObject(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value == null ? null : ((Number) value).intValue(); + } + + // DATA_DEFAULT is stored as a LONG column in Oracle/Dameng. JDBC requires LONG + // columns to be read before other columns. We also try getCharacterStream as a + // fallback because some drivers don't support getString on LONG columns. + private static String readLongColumn(ResultSet rs, String column) throws Exception { + try { + String value = rs.getString(column); + if (value != null) { + return value; + } + } catch (Exception ignored) { + } + try (java.io.Reader reader = rs.getCharacterStream(column)) { + if (reader == null) { + return null; + } + StringBuilder sb = new StringBuilder(); + char[] buf = new char[4096]; + int n; + while ((n = reader.read(buf)) != -1) { + sb.append(buf, 0, n); + } + return sb.toString(); + } + } + + private static List splitNonEmpty(String value, String delimiter) { + List result = new ArrayList<>(); + Arrays.stream(value.split(delimiter)) + .filter(part -> !part.isEmpty()) + .forEach(result::add); + return result; + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + private void fillMissingColumnComments(String schema, String table, List columns) { + if (columns.stream().noneMatch(column -> !notBlank(column.getComment()))) { + return; + } + Map comments = new HashMap<>(); + queryColumnComments( + comments, + "SELECT COLUMN_NAME, COMMENTS FROM USER_COL_COMMENTS WHERE TABLE_NAME = ?", + table + ); + queryColumnComments( + comments, + "SELECT COLNAME, COMMENT$ FROM SYS.SYSCOLUMNCOMMENTS WHERE SCHNAME = ? AND TVNAME = ?", + schema, + table + ); + queryColumnComments( + comments, + "SELECT COLUMN_NAME, COMMENTS FROM ALL_COL_COMMENTS WHERE UPPER(OWNER) = UPPER(?) AND UPPER(TABLE_NAME) = UPPER(?)", + schema, + table + ); + for (ColumnInfo column : columns) { + if (notBlank(column.getComment())) { + continue; + } + String comment = comments.get(column.getName().toUpperCase(Locale.ROOT)); + if (notBlank(comment)) { + column.setComment(comment); + } + } + } + + private void queryColumnComments(Map comments, String sql, String... params) { + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + for (int i = 0; i < params.length; i++) { + stmt.setString(i + 1, params[i]); + } + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String column = rs.getString(1); + String comment = rs.getString(2); + if (notBlank(column) && notBlank(comment)) { + comments.putIfAbsent(column.toUpperCase(Locale.ROOT), comment); + } + } + } + } catch (Exception ignored) { + // Some Dameng versions or users do not expose every comment view. + } + } + + private String appendTableAndColumnComments(String ddl, String schema, String table) throws Exception { + StringBuilder result = new StringBuilder(ddl == null ? "" : ddl.trim()); + String tableRef = qualifiedName(schema, table); + String tableComment = tableComment(schema, table); + if (notBlank(tableComment) && !containsCommentOnTable(result.toString(), schema, table)) { + appendCommentStatement(result, "COMMENT ON TABLE " + tableRef + " IS '" + sqlStringBody(tableComment) + "';"); + } + for (ColumnInfo column : getColumns(schema, table)) { + if (!notBlank(column.getComment()) || containsCommentOnColumn(result.toString(), schema, table, column.getName())) { + continue; + } + appendCommentStatement(result, "COMMENT ON COLUMN " + tableRef + "." + JdbcIdentifiers.INSTANCE.doubleQuote(column.getName()) + " IS '" + sqlStringBody(column.getComment()) + "';"); + } + return result.toString(); + } + + private String tableComment(String schema, String table) throws Exception { + String sql = """ + SELECT COMMENTS + FROM ALL_TAB_COMMENTS + WHERE OWNER = ? AND TABLE_NAME = ? + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + return rs.next() ? rs.getString(1) : null; + } + } + } + + private static void appendCommentStatement(StringBuilder ddl, String statement) { + if (!ddl.isEmpty()) { + if (ddl.charAt(ddl.length() - 1) != '\n') { + ddl.append("\n"); + } + ddl.append("\n"); + } + ddl.append(statement); + } + + private static boolean containsCommentOnTable(String ddl, String schema, String table) { + return normalizedDdl(ddl).contains("COMMENT ON TABLE " + normalizedQualifiedName(schema, table)); + } + + private static boolean containsCommentOnColumn(String ddl, String schema, String table, String column) { + return normalizedDdl(ddl).contains("COMMENT ON COLUMN " + normalizedQualifiedName(schema, table) + "." + normalizedIdentifier(column)); + } + + private static String qualifiedName(String schema, String name) { + if (!notBlank(schema)) { + return JdbcIdentifiers.INSTANCE.doubleQuote(name); + } + return JdbcIdentifiers.INSTANCE.doubleQuote(schema) + "." + JdbcIdentifiers.INSTANCE.doubleQuote(name); + } + + private static String normalizedQualifiedName(String schema, String name) { + if (!notBlank(schema)) { + return normalizedIdentifier(name); + } + return normalizedIdentifier(schema) + "." + normalizedIdentifier(name); + } + + private static String normalizedIdentifier(String value) { + return JdbcIdentifiers.INSTANCE.doubleQuote(value).toUpperCase(Locale.ROOT); + } + + private static String normalizedDdl(String ddl) { + return coalesce(ddl).toUpperCase(Locale.ROOT).replaceAll("\\s+", " "); + } + + private static boolean notBlank(String value) { + return value != null && !value.trim().isEmpty(); + } + + private static String sqlStringBody(String value) { + return value.replace("'", "''"); + } + + @Override + public String getExplainInfo(String sql, String database, String schema, int timeoutSecs, String mode) { + return unchecked(() -> { + Connection conn = requireConnected(); + boolean autotrace = "autotrace".equalsIgnoreCase(mode); + String planText = null; + + if (autotrace) { + boolean monitorEnabled = false; + try (Statement s = conn.createStatement()) { + s.execute("SF_SET_SESSION_PARA_VALUE('MONITOR_SQL_EXEC', 1)"); + monitorEnabled = true; + } catch (Exception ignored) {} + + try (Statement stmt = conn.createStatement()) { + if (timeoutSecs >= 0) { + try { stmt.setQueryTimeout(timeoutSecs); } catch (Exception ignored) {} + } + boolean hasResultSet = stmt.execute(sql); + if (hasResultSet) { + try (ResultSet rs = stmt.getResultSet()) { + while (rs.next()) { /* consume all rows */ } + } + } + try { + Class dmConnClass = Class.forName("dm.jdbc.driver.DmdbConnection"); + if (dmConnClass.isInstance(conn)) { + Method m = dmConnClass.getMethod("getExplainInfo", Statement.class); + planText = (String) m.invoke(dmConnClass.cast(conn), stmt); + } + } catch (Exception ignored) {} + } finally { + if (monitorEnabled) { + try (Statement s = conn.createStatement()) { + s.execute("SF_SET_SESSION_PARA_VALUE('MONITOR_SQL_EXEC', 0)"); + } catch (Exception ignored) {} + } + } + } else { + try { + Class dmConnClass = Class.forName("dm.jdbc.driver.DmdbConnection"); + if (dmConnClass.isInstance(conn)) { + Method m = dmConnClass.getMethod("getExplainInfo", String.class); + planText = (String) m.invoke(dmConnClass.cast(conn), sql); + } + } catch (Exception ignored) {} + } + + if (planText == null || planText.trim().isEmpty()) { + try (Statement explainStmt = conn.createStatement(); + ResultSet rs = explainStmt.executeQuery("EXPLAIN " + sql)) { + StringBuilder sb = new StringBuilder(); + while (rs.next()) { + sb.append(rs.getString(1)).append("\n"); + } + planText = sb.toString().trim(); + } catch (Exception ignored) {} + } + return planText != null ? planText : ""; + }); + } + + public static void main(String[] args) { + new JsonRpcServer(new DamengAgent()).run(); + } +} diff --git a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentMetadataTest.java b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentMetadataTest.java new file mode 100644 index 0000000000..fc79e2ba25 --- /dev/null +++ b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentMetadataTest.java @@ -0,0 +1,279 @@ +package com.dbx.agent.dameng; + +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.TableInfo; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import com.dbx.agent.test.TestSupport; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.Arrays; +import java.util.List; + +class DamengAgentMetadataTest { + @Test + void usesColumnCommentsMetadataQuery() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.getColumns("APP", "USERS"); + + String columnsSql = JdbcMetadataSqlFake.statements.stream() + .filter(sql -> sql.contains("ALL_TAB_COLUMNS")) + .findFirst() + .orElseThrow(); + Assertions.assertTrue(columnsSql.contains("LEFT JOIN ALL_COL_COMMENTS"), columnsSql); + } + + @Test + void usesTableCommentsMetadataQuery() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTables("APP"); + + String tablesSql = JdbcMetadataSqlFake.statements.stream() + .filter(sql -> sql.contains("ALL_TAB_COMMENTS")) + .findFirst() + .orElseThrow(); + Assertions.assertTrue(tablesSql.contains("COMMENTS"), tablesSql); + Assertions.assertFalse(tablesSql.contains("ALL_TABLES"), tablesSql); + Assertions.assertFalse(tablesSql.contains("ALL_VIEWS"), tablesSql); + } + + @Test + void mapsTableCommentFromMetadata() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, metadataConnection()); + + List tables = agent.listTables("APP"); + + Assertions.assertEquals(1, tables.size()); + Assertions.assertEquals("USERS", tables.get(0).getName()); + Assertions.assertEquals("用户示例表", tables.get(0).getComment()); + } + + @Test + void mapsTableCommentToObjectMetadata() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, metadataConnection()); + + List objects = agent.listObjects("APP"); + + Assertions.assertEquals(1, objects.size()); + Assertions.assertEquals("USERS", objects.get(0).getName()); + Assertions.assertEquals("用户示例表", objects.get(0).getComment()); + } + + @Test + void triggerMetadataDoesNotRequireOracleTriggerTypeColumn() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listTriggers("APP", "USERS"); + + String triggersSql = JdbcMetadataSqlFake.statements.stream() + .filter(sql -> sql.contains("ALL_TRIGGERS")) + .findFirst() + .orElseThrow(); + Assertions.assertFalse(triggersSql.contains("TRIGGERING_EVENT, TRIGGER_TYPE"), triggersSql); + Assertions.assertTrue(triggersSql.contains("'' AS TRIGGER_TYPE"), triggersSql); + } + + @Test + void mapsColumnCommentFromMetadata() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, metadataConnection()); + + List columns = agent.getColumns("APP", "USERS"); + + Assertions.assertEquals(1, columns.size()); + Assertions.assertEquals("id comment", columns.get(0).getComment()); + } + + @Test + void mapsColumnCommentFromFallbackMetadataViews() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, metadataConnection(null, "fallback id comment")); + + List columns = agent.getColumns("APP", "USERS"); + + Assertions.assertEquals(1, columns.size()); + Assertions.assertEquals("fallback id comment", columns.get(0).getComment()); + } + + @Test + void appendsTableAndColumnCommentsToTableDdl() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, metadataConnection()); + + String ddl = agent.getTableDdl("APP", "USERS"); + + Assertions.assertTrue(ddl.contains("CREATE TABLE \"APP\".\"USERS\""), ddl); + Assertions.assertTrue(ddl.contains("COMMENT ON TABLE \"APP\".\"USERS\" IS '用户示例表';"), ddl); + Assertions.assertTrue(ddl.contains("COMMENT ON COLUMN \"APP\".\"USERS\".\"ID\" IS 'id comment';"), ddl); + } + + private static Connection metadataConnection() { + return metadataConnection("id comment", null); + } + + private static Connection metadataConnection(String allColumnComment, String fallbackColumnComment) { + return proxy(Connection.class, (method, args) -> { + String name = method.getName(); + if ("prepareStatement".equals(name)) { + String sql = (String) args[0]; + if (sql.contains("DBMS_METADATA.GET_DDL")) { + return metadataStatement(List.of(List.of("CREATE TABLE \"APP\".\"USERS\" (\n \"ID\" NUMBER\n);"))); + } + if (sql.contains("ALL_CONS_COLUMNS")) { + return metadataStatement(List.of(List.of("ID"))); + } + if (sql.startsWith("SELECT COMMENTS")) { + return metadataStatement(List.of(List.of("用户示例表"))); + } + if (sql.contains("USER_COL_COMMENTS")) { + return metadataStatement(fallbackColumnComment == null ? List.of() : List.of(List.of("ID", fallbackColumnComment))); + } + if (sql.contains("SYSCOLUMNCOMMENTS")) { + return metadataStatement(List.of()); + } + if (sql.contains("ALL_TAB_COMMENTS")) { + return metadataStatement(List.of(List.of("USERS", "TABLE", "用户示例表"))); + } + if (sql.contains("ALL_COL_COMMENTS") && !sql.contains("ALL_TAB_COLUMNS")) { + return metadataStatement(List.of()); + } + if (sql.contains("ALL_OBJECTS")) { + return metadataStatement(List.of()); + } + if (sql.contains("ALL_TAB_COLUMNS")) { + return metadataStatement(List.of(Arrays.asList( + "ID", + "NUMBER", + "N", + Integer.valueOf(10), + Integer.valueOf(0), + Integer.valueOf(22), + Integer.valueOf(10), + allColumnComment + ))); + } + } + if ("close".equals(name)) { + return null; + } + if ("isClosed".equals(name)) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static PreparedStatement metadataStatement(List> rows) { + return proxy(PreparedStatement.class, (method, args) -> { + String name = method.getName(); + if ("executeQuery".equals(name)) { + return metadataResultSet(rows); + } + if ("setString".equals(name) || "close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet metadataResultSet(List> rows) { + int[] index = {-1}; + return proxy(ResultSet.class, (method, args) -> { + String name = method.getName(); + if ("next".equals(name)) { + index[0] += 1; + return index[0] < rows.size(); + } + if ("getString".equals(name)) { + if (args[0] instanceof Integer columnIndex) { + Object value = rows.get(index[0]).get(columnIndex - 1); + return value == null ? null : value.toString(); + } + return switch (((String) args[0]).toUpperCase()) { + case "COLUMN_NAME" -> string(rows, index[0], 0); + case "DATA_TYPE" -> string(rows, index[0], 1); + case "NULLABLE" -> string(rows, index[0], 2); + case "DATA_DEFAULT" -> null; + case "COMMENTS", "COMMENT$" -> string(rows, index[0], rows.get(index[0]).size() - 1); + case "COLNAME" -> string(rows, index[0], 0); + default -> null; + }; + } + if ("getObject".equals(name)) { + return switch (((String) args[0]).toUpperCase()) { + case "DATA_PRECISION" -> rows.get(index[0]).get(3); + case "DATA_SCALE" -> rows.get(index[0]).get(4); + case "DATA_LENGTH" -> rows.get(index[0]).get(5); + case "CHAR_LENGTH" -> rows.get(index[0]).get(6); + default -> null; + }; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static String string(List> rows, int rowIndex, int columnIndex) { + Object value = rows.get(rowIndex).get(columnIndex); + return value == null ? null : value.toString(); + } + + @SuppressWarnings("unchecked") + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentPagingTest.java b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentPagingTest.java new file mode 100644 index 0000000000..bca5f6b01e --- /dev/null +++ b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentPagingTest.java @@ -0,0 +1,277 @@ +package com.dbx.agent.dameng; + +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.test.TestSupport; +import dm.jdbc.driver.DmdbTimestamp; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.sql.Types; +import javax.sql.rowset.serial.SerialBlob; +import javax.sql.rowset.serial.SerialClob; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class DamengAgentPagingTest { + @Test + void executeQueryPageStringifiesTimestampValuesBeforeJsonRpcSerialization() { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection( + agent, + singleValueConnection( + "CURRENT_TIMESTAMP", + Types.TIMESTAMP, + DmdbTimestamp.valueOf("2026-05-20 12:34:56.123456789"), + null, + null + ) + ); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = handleRequest(server, """ + { + "jsonrpc": "2.0", + "id": 1, + "method": "execute_query_page", + "params": { + "sql": "SELECT CURRENT_TIMESTAMP", + "schema": "SYS", + "pageSize": 100 + } + } + """); + + assertFalse(response.contains("\"error\""), response); + assertTrue(response.contains("2026-05-20 12:34:56.123456"), response); + } + + @Test + void executeQueryPageReadsClobValuesAsText() throws Exception { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection( + agent, + singleValueConnection( + "content", + Types.CLOB, + new SerialClob("真实文本".toCharArray()), + "真实文本", + null + ) + ); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = handleRequest(server, """ + { + "jsonrpc": "2.0", + "id": 1, + "method": "execute_query_page", + "params": { + "sql": "SELECT content FROM sample", + "schema": "SYS", + "pageSize": 100 + } + } + """); + + assertFalse(response.contains("\"error\""), response); + assertTrue(response.contains("真实文本"), response); + assertFalse(response.contains("SerialClob"), response); + } + + @Test + void executeQueryPageFormatsBlobValuesAsHexText() throws Exception { + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection( + agent, + singleValueConnection( + "payload", + Types.BLOB, + new SerialBlob(new byte[]{0x01, 0x2A, (byte) 0xFF}), + null, + new byte[]{0x01, 0x2A, (byte) 0xFF} + ) + ); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = handleRequest(server, """ + { + "jsonrpc": "2.0", + "id": 1, + "method": "execute_query_page", + "params": { + "sql": "SELECT payload FROM sample", + "schema": "SYS", + "pageSize": 100 + } + } + """); + + assertFalse(response.contains("\"error\""), response); + assertTrue(response.contains("0x012aff"), response); + assertFalse(response.contains("SerialBlob"), response); + } + + private static Connection singleValueConnection( + String columnLabel, + int sqlType, + Object objectValue, + String stringValue, + byte[] bytesValue + ) { + Statement statement = singleValueStatement(columnLabel, sqlType, objectValue, stringValue, bytesValue); + return proxy(Connection.class, (method, args) -> { + String name = method.getName(); + if ("createStatement".equals(name)) { + return statement; + } + if ("getAutoCommit".equals(name)) { + return true; + } + if ("setAutoCommit".equals(name) || "commit".equals(name) || "rollback".equals(name) || "close".equals(name)) { + return null; + } + if ("isClosed".equals(name)) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static Statement singleValueStatement( + String columnLabel, + int sqlType, + Object objectValue, + String stringValue, + byte[] bytesValue + ) { + ResultSet resultSet = singleValueResultSet(columnLabel, sqlType, objectValue, stringValue, bytesValue); + return proxy(Statement.class, (method, args) -> { + String name = method.getName(); + if ("execute".equals(name)) { + return true; + } + if ("getResultSet".equals(name)) { + return resultSet; + } + if ("getUpdateCount".equals(name)) { + return 0; + } + if ("setMaxRows".equals(name) || "setFetchSize".equals(name) || "close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet singleValueResultSet( + String columnLabel, + int sqlType, + Object objectValue, + String stringValue, + byte[] bytesValue + ) { + final int[] index = {-1}; + ResultSetMetaData metadata = singleValueMetadata(columnLabel, sqlType); + return proxy(ResultSet.class, (method, args) -> { + String name = method.getName(); + if ("next".equals(name)) { + index[0] += 1; + return index[0] == 0; + } + if ("getMetaData".equals(name)) { + return metadata; + } + if ("getObject".equals(name)) { + return objectValue; + } + if ("getString".equals(name)) { + return stringValue; + } + if ("getBytes".equals(name)) { + return bytesValue; + } + if ("wasNull".equals(name)) { + return false; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSetMetaData singleValueMetadata(String columnLabel, int sqlType) { + return proxy(ResultSetMetaData.class, (method, args) -> { + String name = method.getName(); + if ("getColumnCount".equals(name)) { + return 1; + } + if ("getColumnLabel".equals(name)) { + return columnLabel; + } + if ("getColumnType".equals(name)) { + return sqlType; + } + return defaultValue(method.getReturnType()); + }); + } + + private static String handleRequest(JsonRpcServer server, String request) { + try { + Method method = JsonRpcServer.class.getDeclaredMethod("handleRequest", String.class); + method.setAccessible(true); + return (String) method.invoke(server, request); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java new file mode 100644 index 0000000000..9670eb5464 --- /dev/null +++ b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.dameng; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class DamengAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new DamengAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL SP_SAMPLE()"; + } +} diff --git a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentUrlTest.java b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentUrlTest.java new file mode 100644 index 0000000000..c1d72a87b3 --- /dev/null +++ b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentUrlTest.java @@ -0,0 +1,29 @@ +package com.dbx.agent.dameng; + +import com.dbx.agent.ConnectParams; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Method; + +class DamengAgentUrlTest { + @Test + void omitsDatabasePathWhenDatabaseIsBlank() throws Exception { + String url = invokeBuildUrl(new ConnectParams("127.0.0.1", 5236, "", "SYSDBA", "pwd", "", "", false)); + + Assertions.assertEquals("jdbc:dm://127.0.0.1:5236", url); + } + + @Test + void appendsDatabasePathWhenDatabaseIsProvided() throws Exception { + String url = invokeBuildUrl(new ConnectParams("127.0.0.1", 5236, "MAIN", "SYSDBA", "pwd", "", "", false)); + + Assertions.assertEquals("jdbc:dm://127.0.0.1:5236/MAIN", url); + } + + private static String invokeBuildUrl(ConnectParams params) throws Exception { + Method method = DamengAgent.class.getDeclaredMethod("buildUrl", ConnectParams.class); + method.setAccessible(true); + return (String) method.invoke(null, params); + } +} diff --git a/agents/drivers/databend/build.gradle b/agents/drivers/databend/build.gradle new file mode 100644 index 0000000000..a1db92d0a1 --- /dev/null +++ b/agents/drivers/databend/build.gradle @@ -0,0 +1,9 @@ +dependencies { + implementation 'com.databend:databend-jdbc:0.4.6' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Databend', 'Main-Class': 'com.dbx.agent.databend.DatabendAgent') + } +} diff --git a/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java b/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java new file mode 100644 index 0000000000..e70e195d29 --- /dev/null +++ b/agents/drivers/databend/src/main/java/com/dbx/agent/databend/DatabendAgent.java @@ -0,0 +1,233 @@ +package com.dbx.agent.databend; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +public final class DatabendAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile DATABEND_PROFILE = new JdbcAgentProfile( + "com.databend.jdbc.DatabendDriver", + "jdbc:databend://{host}:{port}/{database}", + 8000, + false, + Collections.singleton("INFORMATION_SCHEMA"), + Arrays.asList("TABLE", "VIEW", "BASE TABLE", "MATERIALIZED VIEW", "SYSTEM TABLE", "SYSTEM VIEW"), + "\"", + "USE", + true, + false, + false, + false + ); + + public DatabendAgent() { + super(DATABEND_PROFILE); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(super.listObjects(schema)); + try { + useSchema(schema); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW PROCEDURES")) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString("name"), "PROCEDURE", schema, rs.getString("comment"))); + } + } + } catch (Exception ignored) { + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + if (!"PROCEDURE".equalsIgnoreCase(objectType)) { + return super.getObjectSource(schema, name, objectType); + } + return unchecked(() -> { + useSchema(schema); + ProcedureMetadata procedure = findProcedure(name); + if (procedure == null) { + throw new IllegalArgumentException("Procedure not found: " + name); + } + Map properties = describeProcedure(name, procedure.inputTypes); + String source = buildProcedureSource(name, procedure, properties); + return new ObjectSource(name, objectType, schema, source); + }); + } + + private void useSchema(String schema) throws SQLException { + if (schema == null || schema.trim().isEmpty()) { + return; + } + try (java.sql.Statement stmt = requireConnected().createStatement()) { + stmt.execute(DATABEND_PROFILE.schemaSwitchSql(schema.trim())); + } + } + + private ProcedureMetadata findProcedure(String name) throws SQLException { + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + "SELECT arguments, comment FROM system.procedures WHERE name = ? ORDER BY procedure_id LIMIT 1" + )) { + stmt.setString(1, name); + try (ResultSet rs = stmt.executeQuery()) { + if (!rs.next()) { + return null; + } + return new ProcedureMetadata( + rs.getString("comment"), + inputTypesFromArguments(rs.getString("arguments")) + ); + } + } + } + + private Map describeProcedure(String name, List inputTypes) throws SQLException { + Map result = new LinkedHashMap<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("DESC PROCEDURE " + simpleProcedureName(name) + "(" + String.join(", ", inputTypes) + ")")) { + while (rs.next()) { + result.put(rs.getString("Property").toLowerCase(Locale.ROOT), rs.getString("Value")); + } + } + return result; + } + + private static String buildProcedureSource(String name, ProcedureMetadata procedure, Map properties) { + List names = signatureNames(properties.get("signature")); + List arguments = new ArrayList<>(); + for (int index = 0; index < procedure.inputTypes.size(); index += 1) { + String argumentName = index < names.size() && !names.get(index).isEmpty() ? names.get(index) : "arg" + (index + 1); + arguments.add(argumentName + " " + procedure.inputTypes.get(index)); + } + + String returns = trimOuterParens(properties.getOrDefault("returns", "")); + String language = properties.getOrDefault("language", "SQL"); + String body = properties.getOrDefault("body", "").trim(); + StringBuilder source = new StringBuilder(); + source.append("CREATE PROCEDURE ").append(simpleProcedureName(name)).append("(").append(String.join(", ", arguments)).append(")"); + if (!returns.isEmpty()) { + source.append("\nRETURNS ").append(returns); + } + source.append("\nLANGUAGE ").append(language); + if (procedure.comment != null && !procedure.comment.isEmpty()) { + source.append("\nCOMMENT = '").append(procedure.comment.replace("'", "''")).append("'"); + } + source.append("\nAS $$\n").append(body).append("\n$$;"); + return source.toString(); + } + + private static String simpleProcedureName(String name) { + if (name != null && name.matches("[A-Za-z_][A-Za-z0-9_]*")) { + return name; + } + throw new IllegalArgumentException("Databend procedure names with special characters are not supported: " + name); + } + + private static List inputTypesFromArguments(String arguments) { + if (arguments == null) { + return Collections.emptyList(); + } + int open = arguments.indexOf('('); + if (open < 0) { + return Collections.emptyList(); + } + int close = matchingParen(arguments, open); + if (close < 0) { + return Collections.emptyList(); + } + return splitTopLevel(arguments.substring(open + 1, close)); + } + + private static List signatureNames(String signature) { + if (signature == null) { + return Collections.emptyList(); + } + String trimmed = trimOuterParens(signature); + return splitTopLevel(trimmed); + } + + private static List splitTopLevel(String value) { + if (value == null || value.trim().isEmpty()) { + return Collections.emptyList(); + } + List result = new ArrayList<>(); + StringBuilder current = new StringBuilder(); + int depth = 0; + for (int index = 0; index < value.length(); index += 1) { + char ch = value.charAt(index); + if (ch == '(') { + depth += 1; + } else if (ch == ')') { + depth = Math.max(0, depth - 1); + } else if (ch == ',' && depth == 0) { + addTrimmed(result, current.toString()); + current.setLength(0); + continue; + } + current.append(ch); + } + addTrimmed(result, current.toString()); + return result; + } + + private static void addTrimmed(List result, String value) { + String trimmed = value.trim(); + if (!trimmed.isEmpty()) { + result.add(trimmed); + } + } + + private static int matchingParen(String value, int open) { + int depth = 0; + for (int index = open; index < value.length(); index += 1) { + char ch = value.charAt(index); + if (ch == '(') { + depth += 1; + } else if (ch == ')') { + depth -= 1; + if (depth == 0) { + return index; + } + } + } + return -1; + } + + private static String trimOuterParens(String value) { + String trimmed = value == null ? "" : value.trim(); + if (trimmed.startsWith("(") && trimmed.endsWith(")") && matchingParen(trimmed, 0) == trimmed.length() - 1) { + return trimmed.substring(1, trimmed.length() - 1).trim(); + } + return trimmed; + } + + private static final class ProcedureMetadata { + private final String comment; + private final List inputTypes; + + ProcedureMetadata(String comment, List inputTypes) { + this.comment = comment; + this.inputTypes = inputTypes; + } + } + + public static void main(String[] args) { + new JsonRpcServer(new DatabendAgent()).run(); + } +} diff --git a/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java b/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java new file mode 100644 index 0000000000..327aa848c2 --- /dev/null +++ b/agents/drivers/databend/src/test/java/com/dbx/agent/databend/DatabendAgentTest.java @@ -0,0 +1,54 @@ +package com.dbx.agent.databend; + +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 DatabendAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new DatabendAgent(); + } + + @Override + protected String resultSetSql() { + return "SELECT 1"; + } + + @Test + void buildsDefaultJdbcUrl() { + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(0); + params.setDatabase("default"); + + Assertions.assertEquals( + "jdbc:databend://db.example.com:8000/default", + DatabendAgent.DATABEND_PROFILE.buildUrl(params) + ); + } + + @Test + void appendsQueryParametersToJdbcUrl() { + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(20080); + params.setDatabase("default"); + params.setUrl_params("sslmode=disable"); + + Assertions.assertEquals( + "jdbc:databend://db.example.com:20080/default?sslmode=disable", + DatabendAgent.DATABEND_PROFILE.buildUrl(params) + ); + } + + @Test + void usesDoubleQuotedDatabaseAsExecutionContext() { + Assertions.assertEquals( + "USE \"analytics\"\"prod\"", + DatabendAgent.DATABEND_PROFILE.schemaSwitchSql("analytics\"prod") + ); + } +} diff --git a/agents/drivers/databricks/build.gradle b/agents/drivers/databricks/build.gradle new file mode 100644 index 0000000000..f6675d8cb2 --- /dev/null +++ b/agents/drivers/databricks/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.databricks:databricks-jdbc:3.3.3' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Databricks SQL', 'Main-Class': 'com.dbx.agent.databricks.DatabricksAgent') + } +} diff --git a/agents/drivers/databricks/src/main/java/com/dbx/agent/databricks/DatabricksAgent.java b/agents/drivers/databricks/src/main/java/com/dbx/agent/databricks/DatabricksAgent.java new file mode 100644 index 0000000000..5b5ef6b99d --- /dev/null +++ b/agents/drivers/databricks/src/main/java/com/dbx/agent/databricks/DatabricksAgent.java @@ -0,0 +1,41 @@ +package com.dbx.agent.databricks; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +public final class DatabricksAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile DATABRICKS_PROFILE = new DatabricksProfile(); + + public DatabricksAgent() { + super(DATABRICKS_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new DatabricksAgent()).run(); + } + + private static final class DatabricksProfile extends JdbcAgentProfile { + private DatabricksProfile() { + super("com.databricks.client.jdbc.Driver", "jdbc:databricks://{host}:{port}/{database}", 443, true); + } + + @Override + public String buildUrl(ConnectParams params) { + if (!params.getConnection_string().trim().isEmpty()) { + return params.getConnection_string(); + } + int port = params.getPort() > 0 ? params.getPort() : getDefaultPort(); + String base = getUrlTemplate() + .replace("{host}", params.getHost()) + .replace("{port}", Integer.toString(port)) + .replace("{database}", params.getDatabase()); + String urlParams = params.getUrl_params() == null ? "" : params.getUrl_params().trim(); + while (urlParams.startsWith(";") || urlParams.startsWith("?") || urlParams.startsWith("&")) { + urlParams = urlParams.substring(1); + } + return urlParams.isEmpty() ? base : base + ";" + urlParams; + } + } +} diff --git a/agents/drivers/databricks/src/test/java/com/dbx/agent/databricks/DatabricksAgentTest.java b/agents/drivers/databricks/src/test/java/com/dbx/agent/databricks/DatabricksAgentTest.java new file mode 100644 index 0000000000..80f5bcba5f --- /dev/null +++ b/agents/drivers/databricks/src/test/java/com/dbx/agent/databricks/DatabricksAgentTest.java @@ -0,0 +1,22 @@ +package com.dbx.agent.databricks; + +import com.dbx.agent.ConnectParams; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class DatabricksAgentTest { + @Test + void buildsHttpPathUrl() { + ConnectParams params = new ConnectParams(); + params.setHost("adb-123.azuredatabricks.net"); + params.setPort(443); + params.setDatabase("default"); + params.setUrl_params("httpPath=/sql/1.0/warehouses/abc;AuthMech=3"); + + assertEquals( + "jdbc:databricks://adb-123.azuredatabricks.net:443/default;httpPath=/sql/1.0/warehouses/abc;AuthMech=3", + DatabricksAgent.DATABRICKS_PROFILE.buildUrl(params) + ); + } +} diff --git a/agents/drivers/db2/build.gradle b/agents/drivers/db2/build.gradle new file mode 100644 index 0000000000..d9af1aac4b --- /dev/null +++ b/agents/drivers/db2/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.ibm.db2:jcc:11.5.9.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'IBM DB2', 'Main-Class': 'com.dbx.agent.db2.Db2Agent') + } +} diff --git a/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java b/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java new file mode 100644 index 0000000000..613050547f --- /dev/null +++ b/agents/drivers/db2/src/main/java/com/dbx/agent/db2/Db2Agent.java @@ -0,0 +1,385 @@ +package com.dbx.agent.db2; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public final class Db2Agent extends BaseDatabaseAgent { + private static final Set NUMERIC_PRECISION_TYPES = Set.of( + "DECIMAL", "NUMERIC", "INTEGER", "SMALLINT", "BIGINT", "REAL", "DOUBLE", "FLOAT" + ); + private static final Set NUMERIC_SCALE_TYPES = Set.of("DECIMAL", "NUMERIC"); + private static final Set CHARACTER_LENGTH_TYPES = Set.of("VARCHAR", "CHAR", "CLOB", "GRAPHIC", "VARGRAPHIC"); + + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public String setSchemaSQL(String schema) { + return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("com.ibm.db2.jcc.DB2Driver"); + connection = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword()); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("com.ibm.db2.jcc.DB2Driver"); + try (Connection conn = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT CURRENT_SERVER FROM SYSIBM.SYSDUMMY1"; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1).trim())); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT SCHEMANAME FROM SYSCAT.SCHEMATA ORDER BY SCHEMANAME"; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(rs.getString(1).trim()); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT TABNAME, TYPE FROM SYSCAT.TABLES WHERE TABSCHEMA = ? AND TYPE IN ('T','V') ORDER BY TABNAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String db2Type = rs.getString(2).trim(); + String type = switch (db2Type) { + case "T" -> "TABLE"; + case "V" -> "VIEW"; + default -> db2Type; + }; + result.add(new TableInfo(rs.getString(1).trim(), type, null)); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + + String sql = "SELECT PROCNAME, 'PROCEDURE' FROM SYSCAT.PROCEDURES WHERE PROCSCHEMA = ? ORDER BY PROCNAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1).trim(), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String sql = "SELECT TEXT FROM SYSCAT.ROUTINES WHERE ROUTINESCHEMA = ? AND ROUTINENAME = ?"; + String source; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, name); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set pkColumns = new java.util.HashSet<>(); + String pkSql = """ + SELECT kc.COLNAME FROM SYSCAT.KEYCOLUSE kc + JOIN SYSCAT.TABCONST tc ON kc.CONSTNAME = tc.CONSTNAME AND kc.TABSCHEMA = tc.TABSCHEMA AND kc.TABNAME = tc.TABNAME + WHERE tc.TYPE = 'P' AND tc.TABSCHEMA = ? AND tc.TABNAME = ? + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(pkSql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + pkColumns.add(rs.getString(1).trim()); + } + } + } + + List result = new ArrayList<>(); + String colSql = """ + SELECT COLNAME, TYPENAME, NULLS, DEFAULT, LENGTH, SCALE + FROM SYSCAT.COLUMNS + WHERE TABSCHEMA = ? AND TABNAME = ? + ORDER BY COLNO + """.stripIndent().trim(); + try (PreparedStatement stmt = requireConnected().prepareStatement(colSql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String name = rs.getString("COLNAME").trim(); + String typeName = rs.getString("TYPENAME").trim(); + Integer length = intObject(rs, "LENGTH"); + Integer scale = intObject(rs, "SCALE"); + String dataType = formatDataType(typeName, length, scale); + + result.add(new ColumnInfo( + name, + dataType, + "Y".equals(rs.getString("NULLS").trim()), + trimNullable(rs.getString("DEFAULT")), + pkColumns.contains(name), + null, + null, + NUMERIC_PRECISION_TYPES.contains(typeName) ? length : null, + NUMERIC_SCALE_TYPES.contains(typeName) ? scale : null, + CHARACTER_LENGTH_TYPES.contains(typeName) ? length : null + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT INDNAME, COLNAMES, UNIQUERULE FROM SYSCAT.INDEXES WHERE TABSCHEMA = ? AND TABNAME = ? ORDER BY INDNAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String colNames = trimToEmpty(rs.getString(2)); + List columns = splitColumns(colNames, "[+-]"); + String uniqueRule = trimToEmpty(rs.getString(3)); + result.add(new IndexInfo( + rs.getString(1).trim(), + columns, + "U".equals(uniqueRule) || "P".equals(uniqueRule), + "P".equals(uniqueRule), + null, + null, + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT CONSTNAME, FK_COLNAMES, REFTABNAME, PK_COLNAMES FROM SYSCAT.REFERENCES WHERE TABSCHEMA = ? AND TABNAME = ? ORDER BY CONSTNAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + List fkColList = splitColumns(trimToEmpty(rs.getString(2)), "\\s+"); + List pkColList = splitColumns(trimToEmpty(rs.getString(4)), "\\s+"); + String refTable = rs.getString(3).trim(); + String constName = rs.getString(1).trim(); + for (int i = 0; i < fkColList.size(); i++) { + result.add(new ForeignKeyInfo( + constName, + fkColList.get(i), + refTable, + i < pkColList.size() ? pkColList.get(i) : "" + )); + } + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT TRIGNAME, TRIGEVENT, TRIGTIME FROM SYSCAT.TRIGGERS WHERE TABSCHEMA = ? AND TABNAME = ? ORDER BY TRIGNAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo( + rs.getString(1).trim(), + rs.getString(2).trim(), + rs.getString(3).trim() + )); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::stringResultValue + ); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Object stringResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + static String buildUrl(ConnectParams params) { + if (!params.getConnection_string().trim().isEmpty()) { + return params.getConnection_string(); + } + String url = "jdbc:db2://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + String extraParams = trimDb2UrlParams(params.getUrl_params()); + if (extraParams.isEmpty()) { + return url; + } + return url + ":" + extraParams + (extraParams.endsWith(";") ? "" : ";"); + } + + private static String trimDb2UrlParams(String urlParams) { + String value = urlParams == null ? "" : urlParams.trim(); + while (value.startsWith("?") || value.startsWith("&") || value.startsWith(":") || value.startsWith(";")) { + value = value.substring(1); + } + return value; + } + + private static String formatDataType(String typeName, Integer length, Integer scale) { + return switch (typeName.toUpperCase(Locale.ROOT)) { + case "VARCHAR", "CHAR", "CLOB", "GRAPHIC", "VARGRAPHIC" -> length != null ? typeName + "(" + length + ")" : typeName; + case "DECIMAL", "NUMERIC" -> { + if (length != null && scale != null && scale > 0) { + yield typeName + "(" + length + "," + scale + ")"; + } + yield length != null ? typeName + "(" + length + ")" : typeName; + } + default -> typeName; + }; + } + + private static Integer intObject(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value == null ? null : ((Number) value).intValue(); + } + + private static List splitColumns(String value, String regex) { + List result = new ArrayList<>(); + for (String part : value.split(regex)) { + if (!part.isBlank()) { + result.add(part); + } + } + return result; + } + + private static String trimNullable(String value) { + return value == null ? null : value.trim(); + } + + private static String trimToEmpty(String value) { + return value == null ? "" : value.trim(); + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + public static void main(String[] args) { + new JsonRpcServer(new Db2Agent()).run(); + } +} diff --git a/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java b/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java new file mode 100644 index 0000000000..e36e528caa --- /dev/null +++ b/agents/drivers/db2/src/test/java/com/dbx/agent/db2/Db2AgentTest.java @@ -0,0 +1,163 @@ +package com.dbx.agent.db2; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.TestSupport; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; + +class Db2AgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new Db2Agent(); + } + + @Override + protected String resultSetSql() { + return "CALL ADMIN_CMD('list applications')"; + } + + @Test + void buildsJdbcUrlWithDb2PropertySuffix() { + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(50000); + params.setDatabase("SAMPLE"); + params.setUrl_params("sslConnection=true"); + + assertEquals("jdbc:db2://db.example.com:50000/SAMPLE:sslConnection=true;", Db2Agent.buildUrl(params)); + } + + @Test + void listsAllCatalogSchemasWithoutOwnerTypeFiltering() { + Db2Agent agent = new Db2Agent(); + AtomicReference executedSql = new AtomicReference<>(); + TestSupport.setPrivateConnection(agent, connection( + executedSql, + rows(row("SCHEMANAME", "APP"), row("SCHEMANAME", "SZ"), row("SCHEMANAME", "TOOLS")) + )); + + List schemas = agent.listSchemas(); + + assertEquals(Arrays.asList("APP", "SZ", "TOOLS"), schemas); + assertEquals("SELECT SCHEMANAME FROM SYSCAT.SCHEMATA ORDER BY SCHEMANAME", executedSql.get()); + assertFalse(executedSql.get().contains("OWNERTYPE")); + } + + private static Connection connection(AtomicReference executedSql, ResultSet resultSet) { + Statement statement = proxy(Statement.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + if ("executeQuery".equals(method.getName())) { + executedSql.set(String.valueOf(args[0])); + return resultSet; + } + if ("close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + return proxy(Connection.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + if ("createStatement".equals(method.getName())) { + return statement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet rows(Map... rows) { + return proxy(ResultSet.class, new MethodHandler() { + private int index = -1; + + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + index += 1; + return index < rows.length; + } + if ("getString".equals(name)) { + Object key = args[0] instanceof Number ? rows[index].keySet().iterator().next() : args[0]; + Object value = rows[index].get(key); + return value == null ? null : String.valueOf(value); + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static Map row(Object... values) { + Map row = new LinkedHashMap<>(); + for (int i = 0; i < values.length; i += 2) { + row.put(String.valueOf(values[i]), values[i + 1]); + } + return row; + } + + private static T proxy(Class type, final MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/drivers/etcd/build.gradle b/agents/drivers/etcd/build.gradle new file mode 100644 index 0000000000..d37a950619 --- /dev/null +++ b/agents/drivers/etcd/build.gradle @@ -0,0 +1,11 @@ +dependencies { + implementation 'com.google.code.gson:gson:2.12.1' + implementation 'io.etcd:jetcd-core:0.8.6' +} + +tasks.named('shadowJar') { + mergeServiceFiles() + manifest { + attributes('Agent-Label': 'etcd', 'Main-Class': 'com.dbx.agent.etcd.EtcdAgent') + } +} diff --git a/agents/drivers/etcd/src/main/java/com/dbx/agent/etcd/EtcdAgent.java b/agents/drivers/etcd/src/main/java/com/dbx/agent/etcd/EtcdAgent.java new file mode 100644 index 0000000000..c29286f047 --- /dev/null +++ b/agents/drivers/etcd/src/main/java/com/dbx/agent/etcd/EtcdAgent.java @@ -0,0 +1,420 @@ +package com.dbx.agent.etcd; + +import com.dbx.agent.AgentProtocol; +import com.google.gson.Gson; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import io.grpc.netty.GrpcSslContexts; +import io.etcd.jetcd.ByteSequence; +import io.etcd.jetcd.Client; +import io.etcd.jetcd.ClientBuilder; +import io.etcd.jetcd.KV; +import io.etcd.jetcd.KeyValue; +import io.etcd.jetcd.kv.DeleteResponse; +import io.etcd.jetcd.kv.GetResponse; +import io.etcd.jetcd.kv.PutResponse; +import io.etcd.jetcd.options.GetOption; +import io.etcd.jetcd.options.PutOption; +import io.netty.handler.ssl.SslContext; +import io.netty.handler.ssl.SslContextBuilder; +import java.io.BufferedReader; +import java.io.File; +import java.io.InputStreamReader; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CharsetDecoder; +import java.nio.charset.CodingErrorAction; +import java.util.Arrays; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +public final class EtcdAgent { + private static final Gson GSON = new Gson(); + private static final int DEFAULT_LIMIT = 100; + private static final int RPC_TIMEOUT_SECONDS = 30; + private static final List CAPABILITIES = Collections.unmodifiableList(Arrays.asList( + AgentProtocol.CAPABILITY_CONNECT, + AgentProtocol.CAPABILITY_TEST_CONNECTION, + AgentProtocol.CAPABILITY_KV + )); + private static Client client; + private static KV kv; + + private EtcdAgent() { + } + + private static Object handshakeResult() { + return new HandshakeResult(AgentProtocol.PROTOCOL_VERSION, AgentProtocol.PROTOCOL_VERSION, CAPABILITIES); + } + + private static Object connect(JsonObject params) throws Exception { + JsonObject connection = connectionObject(params); + Client nextClient = buildClient(connection); + nextClient.getKVClient().get(byteSequence("\0")).get(RPC_TIMEOUT_SECONDS, TimeUnit.SECONDS); + closeClient(); + client = nextClient; + kv = client.getKVClient(); + return Collections.singletonMap("ok", true); + } + + static Client buildClient(JsonObject connection) throws Exception { + List endpoints = endpoints(connection); + ClientBuilder builder = Client.builder().endpoints(endpoints.toArray(String[]::new)); + String username = stringOrEmpty(connection, "username"); + String password = stringOrEmpty(connection, "password"); + if (!username.isBlank()) { + builder.user(byteSequence(username)); + builder.password(byteSequence(password)); + } + if (boolOrDefault(connection, "ssl", false)) { + builder.sslContext(sslContext(connection)); + } + return builder.build(); + } + + static List endpoints(JsonObject connection) { + String configured = firstNonBlank( + stringOrNull(connection, "etcd_endpoints"), + stringOrNull(connection, "endpoints"), + stringOrNull(connection, "connection_string") + ); + List result = new ArrayList<>(); + if (configured != null) { + for (String endpoint : configured.split("[,\\n]")) { + String normalized = normalizeEndpoint(endpoint.trim(), boolOrDefault(connection, "ssl", false)); + if (!normalized.isBlank()) { + result.add(normalized); + } + } + } + if (result.isEmpty()) { + String host = stringOrDefault(connection, "host", "127.0.0.1"); + int port = intOrDefault(connection, "port", 2379); + result.add(normalizeEndpoint(host + ":" + port, boolOrDefault(connection, "ssl", false))); + } + return result; + } + + private static String normalizeEndpoint(String endpoint, boolean tls) { + if (endpoint.isBlank()) { + return ""; + } + if (endpoint.startsWith("http://") || endpoint.startsWith("https://")) { + return endpoint; + } + return (tls ? "https://" : "http://") + endpoint; + } + + private static SslContext sslContext(JsonObject connection) throws Exception { + SslContextBuilder builder = GrpcSslContexts.forClient(); + String ca = stringOrEmpty(connection, "ca_cert_path"); + if (!ca.isBlank()) { + builder.trustManager(new File(ca)); + } + String cert = firstNonBlank(stringOrNull(connection, "client_cert_path"), stringOrNull(connection, "cert_path")); + String key = firstNonBlank(stringOrNull(connection, "client_key_path"), stringOrNull(connection, "key_path")); + if ((cert == null) != (key == null)) { + throw new IllegalArgumentException("Client certificate and key must be provided together"); + } + if (cert != null) { + builder.keyManager(new File(cert), new File(key)); + } + return builder.build(); + } + + private static Object listPrefix(JsonObject params) throws Exception { + KV active = requireKv(); + String prefix = stringOrDefault(params, "prefix", ""); + int limit = intOrDefault(params, "limit", DEFAULT_LIMIT); + String continuation = stringOrNull(params, "continuation"); + ByteSequence start = continuation == null || continuation.isBlank() + ? prefixStart(prefix) + : ByteSequence.from(Base64.getDecoder().decode(continuation)); + GetOption option = GetOption.newBuilder() + .withRange(prefixEnd(byteSequence(prefix))) + .withLimit(Math.max(1, limit)) + .withSortField(GetOption.SortTarget.KEY) + .withSortOrder(GetOption.SortOrder.ASCEND) + .build(); + GetResponse response = active.get(start, option).get(RPC_TIMEOUT_SECONDS, TimeUnit.SECONDS); + + List> keys = new ArrayList<>(); + List kvs = response.getKvs(); + for (KeyValue item : kvs) { + Map row = metadata(item); + row.put("key", displayBytes(item.getKey().getBytes())); + keys.add(row); + } + + Map result = new LinkedHashMap<>(); + result.put("keys", keys); + result.put("continuation", response.isMore() && !kvs.isEmpty() ? nextContinuation(kvs.get(kvs.size() - 1)) : null); + result.put("revision", response.getHeader().getRevision()); + return result; + } + + private static Object get(JsonObject params) throws Exception { + KV active = requireKv(); + String key = params.get("key").getAsString(); + GetResponse response = active.get(byteSequence(key)).get(RPC_TIMEOUT_SECONDS, TimeUnit.SECONDS); + Map result = new LinkedHashMap<>(); + if (response.getKvs().isEmpty()) { + result.put("found", false); + result.put("key", null); + result.put("value", null); + result.put("metadata", null); + return result; + } + KeyValue item = response.getKvs().get(0); + result.put("found", true); + result.put("key", displayBytes(item.getKey().getBytes())); + result.put("value", valueObject(item.getValue().getBytes())); + result.put("metadata", metadata(item)); + return result; + } + + private static Object put(JsonObject params) throws Exception { + KV active = requireKv(); + String key = params.get("key").getAsString(); + byte[] value = parseValue(params.getAsJsonObject("value")); + JsonElement leaseElement = params.get("lease"); + PutResponse response; + if (leaseElement != null && !leaseElement.isJsonNull()) { + PutOption option = PutOption.newBuilder().withLeaseId(leaseElement.getAsLong()).build(); + response = active.put(byteSequence(key), ByteSequence.from(value), option).get(RPC_TIMEOUT_SECONDS, TimeUnit.SECONDS); + } else { + response = active.put(byteSequence(key), ByteSequence.from(value)).get(RPC_TIMEOUT_SECONDS, TimeUnit.SECONDS); + } + return Collections.singletonMap("revision", response.getHeader().getRevision()); + } + + private static Object delete(JsonObject params) throws Exception { + KV active = requireKv(); + DeleteResponse response = + active.delete(byteSequence(params.get("key").getAsString())).get(RPC_TIMEOUT_SECONDS, TimeUnit.SECONDS); + Map result = new LinkedHashMap<>(); + result.put("deleted", response.getDeleted()); + result.put("revision", response.getHeader().getRevision()); + return result; + } + + private static Object dispatch(String method, JsonObject params) throws Exception { + return switch (method) { + case AgentProtocol.METHOD_HANDSHAKE -> handshakeResult(); + case AgentProtocol.METHOD_CONNECT, AgentProtocol.METHOD_TEST_CONNECTION -> connect(params); + case AgentProtocol.KV_METHOD_LIST_PREFIX -> listPrefix(params); + case AgentProtocol.KV_METHOD_GET -> get(params); + case AgentProtocol.KV_METHOD_PUT -> put(params); + case AgentProtocol.KV_METHOD_DELETE -> delete(params); + case AgentProtocol.METHOD_DISCONNECT -> { + closeClient(); + yield Collections.singletonMap("ok", true); + } + case AgentProtocol.METHOD_SHUTDOWN -> { + closeClient(); + System.exit(0); + yield Collections.singletonMap("ok", true); + } + default -> throw new IllegalArgumentException("Unknown method: " + method); + }; + } + + static String handleRequest(String line) { + JsonObject req = JsonParser.parseString(line).getAsJsonObject(); + JsonElement id = req.get("id"); + String method = req.get("method").getAsString(); + JsonObject params = req.has("params") && req.get("params").isJsonObject() + ? req.getAsJsonObject("params") + : new JsonObject(); + + JsonObject response = new JsonObject(); + response.addProperty("jsonrpc", "2.0"); + response.add("id", id); + + try { + Object result = dispatch(method, params); + response.add("result", GSON.toJsonTree(result)); + } catch (Exception e) { + JsonObject error = new JsonObject(); + error.addProperty("code", -1); + error.addProperty("message", e.getMessage() == null ? "Unknown error" : e.getMessage()); + response.add("error", error); + } + + return GSON.toJson(response); + } + + private static JsonObject connectionObject(JsonObject params) { + JsonElement connection = params.get("connection"); + return connection != null && connection.isJsonObject() ? connection.getAsJsonObject() : params; + } + + private static KV requireKv() { + if (kv == null) { + throw new IllegalStateException("Not connected"); + } + return kv; + } + + private static void closeClient() { + if (kv != null) { + kv.close(); + kv = null; + } + if (client != null) { + client.close(); + client = null; + } + } + + private static ByteSequence byteSequence(String value) { + return ByteSequence.from(value, java.nio.charset.StandardCharsets.UTF_8); + } + + private static ByteSequence prefixEnd(ByteSequence prefix) { + byte[] bytes = prefix.getBytes(); + if (bytes.length == 0) { + return ByteSequence.from(new byte[] {0}); + } + byte[] end = Arrays.copyOf(bytes, bytes.length); + for (int i = end.length - 1; i >= 0; i--) { + if ((end[i] & 0xff) < 0xff) { + end[i]++; + return ByteSequence.from(Arrays.copyOf(end, i + 1)); + } + } + return ByteSequence.from(new byte[] {0}); + } + + private static ByteSequence prefixStart(String prefix) { + if (prefix.isEmpty()) { + return ByteSequence.from(new byte[] {0}); + } + return byteSequence(prefix); + } + + private static String nextContinuation(KeyValue item) { + byte[] key = item.getKey().getBytes(); + byte[] next = Arrays.copyOf(key, key.length + 1); + next[next.length - 1] = 0; + return Base64.getEncoder().encodeToString(next); + } + + private static Map metadata(KeyValue item) { + Map metadata = new LinkedHashMap<>(); + metadata.put("createRevision", item.getCreateRevision()); + metadata.put("modRevision", item.getModRevision()); + metadata.put("version", item.getVersion()); + metadata.put("lease", item.getLease()); + metadata.put("valueSize", item.getValue().size()); + return metadata; + } + + private static Map valueObject(byte[] bytes) { + Map value = new LinkedHashMap<>(); + String utf8 = strictUtf8(bytes); + if (utf8 != null) { + value.put("encoding", "utf8"); + value.put("data", utf8); + } else { + value.put("encoding", "base64"); + value.put("data", Base64.getEncoder().encodeToString(bytes)); + } + return value; + } + + private static byte[] parseValue(JsonObject value) { + String encoding = stringOrDefault(value, "encoding", "utf8"); + String data = stringOrDefault(value, "data", ""); + if ("base64".equals(encoding)) { + return Base64.getDecoder().decode(data); + } + if (!"utf8".equals(encoding)) { + throw new IllegalArgumentException("Unsupported value encoding: " + encoding); + } + return data.getBytes(java.nio.charset.StandardCharsets.UTF_8); + } + + private static String displayBytes(byte[] bytes) { + String utf8 = strictUtf8(bytes); + return utf8 == null ? Base64.getEncoder().encodeToString(bytes) : utf8; + } + + private static String strictUtf8(byte[] bytes) { + CharsetDecoder decoder = java.nio.charset.StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT); + try { + return decoder.decode(ByteBuffer.wrap(bytes)).toString(); + } catch (CharacterCodingException e) { + return null; + } + } + + private static String stringOrNull(JsonObject object, String key) { + JsonElement element = object.get(key); + return element == null || element.isJsonNull() ? null : element.getAsString(); + } + + private static String stringOrEmpty(JsonObject object, String key) { + return stringOrDefault(object, key, ""); + } + + private static String stringOrDefault(JsonObject object, String key, String fallback) { + String value = stringOrNull(object, key); + return value == null ? fallback : value; + } + + private static int intOrDefault(JsonObject object, String key, int fallback) { + JsonElement element = object.get(key); + return element == null || element.isJsonNull() ? fallback : element.getAsInt(); + } + + private static boolean boolOrDefault(JsonObject object, String key, boolean fallback) { + JsonElement element = object.get(key); + return element == null || element.isJsonNull() ? fallback : element.getAsBoolean(); + } + + private static String firstNonBlank(String... values) { + for (String value : values) { + if (value != null && !value.isBlank()) { + return value; + } + } + return null; + } + + public static void main(String[] args) throws Exception { + System.out.println("{\"ready\":true}"); + System.out.flush(); + + BufferedReader reader = new BufferedReader(new InputStreamReader(System.in)); + while (true) { + String line = reader.readLine(); + if (line == null) { + break; + } + + System.out.println(handleRequest(line)); + System.out.flush(); + } + } + + private static final class HandshakeResult { + private final int protocolVersion; + private final int agentProtocolVersion; + private final List capabilities; + + private HandshakeResult(int protocolVersion, int agentProtocolVersion, List capabilities) { + this.protocolVersion = protocolVersion; + this.agentProtocolVersion = agentProtocolVersion; + this.capabilities = capabilities; + } + } +} diff --git a/agents/drivers/etcd/src/test/java/com/dbx/agent/etcd/EtcdAgentTest.java b/agents/drivers/etcd/src/test/java/com/dbx/agent/etcd/EtcdAgentTest.java new file mode 100644 index 0000000000..65c205df8a --- /dev/null +++ b/agents/drivers/etcd/src/test/java/com/dbx/agent/etcd/EtcdAgentTest.java @@ -0,0 +1,55 @@ +package com.dbx.agent.etcd; + +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +final class EtcdAgentTest { + @Test + void handshakeAdvertisesKvCapability() { + String response = EtcdAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"handshake\",\"params\":{}}" + ); + + JsonObject result = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("result"); + + Assertions.assertEquals(1, result.get("protocolVersion").getAsInt()); + Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"kv\""))); + Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"connect\""))); + } + + @Test + void endpointsUseConfiguredListAndScheme() { + JsonObject connection = JsonParser.parseString( + "{\"endpoints\":\"etcd-1:2379,https://etcd-2:2379\",\"ssl\":true}" + ).getAsJsonObject(); + + Assertions.assertEquals( + List.of("https://etcd-1:2379", "https://etcd-2:2379"), + EtcdAgent.endpoints(connection) + ); + } + + @Test + void endpointsFallbackToHostPort() { + JsonObject connection = JsonParser.parseString( + "{\"host\":\"127.0.0.1\",\"port\":2379}" + ).getAsJsonObject(); + + Assertions.assertEquals(List.of("http://127.0.0.1:2379"), EtcdAgent.endpoints(connection)); + } + + @Test + void kvMethodDispatchReturnsJsonRpcErrorWhenDisconnected() { + String response = EtcdAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"kv_get\",\"params\":{\"key\":\"/app/name\"}}" + ); + + JsonObject error = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("error"); + + Assertions.assertEquals(-1, error.get("code").getAsInt()); + Assertions.assertEquals("Not connected", error.get("message").getAsString()); + } +} diff --git a/agents/drivers/exasol/build.gradle b/agents/drivers/exasol/build.gradle new file mode 100644 index 0000000000..71dda62bc5 --- /dev/null +++ b/agents/drivers/exasol/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.exasol:exasol-jdbc:26.2.7' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Exasol', 'Main-Class': 'com.dbx.agent.exasol.ExasolAgent') + } +} diff --git a/agents/drivers/exasol/src/main/java/com/dbx/agent/exasol/ExasolAgent.java b/agents/drivers/exasol/src/main/java/com/dbx/agent/exasol/ExasolAgent.java new file mode 100644 index 0000000000..8dd7bddee8 --- /dev/null +++ b/agents/drivers/exasol/src/main/java/com/dbx/agent/exasol/ExasolAgent.java @@ -0,0 +1,22 @@ +package com.dbx.agent.exasol; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +public final class ExasolAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile EXASOL_PROFILE = new JdbcAgentProfile( + "com.exasol.jdbc.EXADriver", + "jdbc:exa:{host}:{port};schema={database}", + 8563, + true + ); + + public ExasolAgent() { + super(EXASOL_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new ExasolAgent()).run(); + } +} diff --git a/agents/drivers/firebird/build.gradle b/agents/drivers/firebird/build.gradle new file mode 100644 index 0000000000..701aad5f99 --- /dev/null +++ b/agents/drivers/firebird/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'org.firebirdsql.jdbc:jaybird:6.0.5' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Firebird', 'Main-Class': 'com.dbx.agent.firebird.FirebirdAgent') + } +} diff --git a/agents/drivers/firebird/src/main/java/com/dbx/agent/firebird/FirebirdAgent.java b/agents/drivers/firebird/src/main/java/com/dbx/agent/firebird/FirebirdAgent.java new file mode 100644 index 0000000000..ed57d29ec7 --- /dev/null +++ b/agents/drivers/firebird/src/main/java/com/dbx/agent/firebird/FirebirdAgent.java @@ -0,0 +1,22 @@ +package com.dbx.agent.firebird; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +public final class FirebirdAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile FIREBIRD_PROFILE = new JdbcAgentProfile( + "org.firebirdsql.jdbc.FBDriver", + "jdbc:firebirdsql://{host}:{port}/{database}", + 3050, + true + ); + + public FirebirdAgent() { + super(FIREBIRD_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new FirebirdAgent()).run(); + } +} diff --git a/agents/drivers/gaussdb/build.gradle b/agents/drivers/gaussdb/build.gradle new file mode 100644 index 0000000000..8aab118786 --- /dev/null +++ b/agents/drivers/gaussdb/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'org.opengauss:opengauss-jdbc:6.0.0-og' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'GaussDB', 'Main-Class': 'com.dbx.agent.gaussdb.GaussdbAgent') + } +} diff --git a/agents/drivers/gaussdb/src/main/java/com/dbx/agent/gaussdb/GaussdbAgent.java b/agents/drivers/gaussdb/src/main/java/com/dbx/agent/gaussdb/GaussdbAgent.java new file mode 100644 index 0000000000..95a5de9a37 --- /dev/null +++ b/agents/drivers/gaussdb/src/main/java/com/dbx/agent/gaussdb/GaussdbAgent.java @@ -0,0 +1,432 @@ +package com.dbx.agent.gaussdb; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Types; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public final class GaussdbAgent extends BaseDatabaseAgent { + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("org.opengauss.Driver"); + connection = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword()); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("org.opengauss.Driver"); + try (Connection conn = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname" + ); ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString("datname"))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('pg_catalog','information_schema','pg_toast') ORDER BY schema_name" + ); ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(rs.getString("schema_name")); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + "SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = ? ORDER BY table_name" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String tableType = rs.getString("table_type"); + if ("BASE TABLE".equals(tableType)) { + tableType = "TABLE"; + } + result.add(new TableInfo(rs.getString("table_name"), tableType, null)); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + try (var stmt = requireConnected().prepareStatement( + "SELECT routine_name, routine_type FROM information_schema.routines WHERE routine_schema = ? ORDER BY routine_name" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String upperType = objectType.toUpperCase(Locale.ROOT); + String sql = switch (upperType) { + case "VIEW" -> "SELECT pg_get_viewdef(to_regclass(?), true)"; + case "FUNCTION" -> """ + SELECT pg_get_functiondef(p.oid) + FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE n.nspname = ? AND p.proname = ? AND p.prokind = 'f' + ORDER BY p.oid LIMIT 1 + """; + case "PROCEDURE" -> """ + SELECT pg_get_functiondef(p.oid) + FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE n.nspname = ? AND p.proname = ? AND p.prokind = 'p' + ORDER BY p.oid LIMIT 1 + """; + default -> throw new IllegalArgumentException("Unsupported object type: " + objectType); + }; + + String source; + if ("VIEW".equals(upperType)) { + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, quoteQualifiedIdentifier(schema, name)); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + } else { + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, name); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = new LinkedHashSet<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT kcu.column_name + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + WHERE tc.constraint_type = 'PRIMARY KEY' + AND tc.table_schema = ? + AND tc.table_name = ? + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + primaryKeys.add(rs.getString("column_name")); + } + } + } + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT column_name, data_type, is_nullable, column_default, + numeric_precision, numeric_scale, character_maximum_length + FROM information_schema.columns + WHERE table_schema = ? AND table_name = ? + ORDER BY ordinal_position + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String columnName = rs.getString("column_name"); + result.add(new ColumnInfo( + columnName, + rs.getString("data_type"), + "YES".equals(rs.getString("is_nullable")), + rs.getString("column_default"), + primaryKeys.contains(columnName), + null, + null, + intOrNull(rs, "numeric_precision"), + intOrNull(rs, "numeric_scale"), + intOrNull(rs, "character_maximum_length") + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT + i.relname AS index_name, + am.amname AS index_type, + ix.indisunique AS is_unique, + ix.indisprimary AS is_primary, + array_agg(a.attname ORDER BY k.n) AS columns + FROM pg_index ix + JOIN pg_class t ON t.oid = ix.indrelid + JOIN pg_class i ON i.oid = ix.indexrelid + JOIN pg_namespace n ON n.oid = t.relnamespace + JOIN pg_am am ON am.oid = i.relam + CROSS JOIN LATERAL unnest(ix.indkey) WITH ORDINALITY AS k(attnum, n) + JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum + WHERE n.nspname = ? AND t.relname = ? + GROUP BY i.relname, am.amname, ix.indisunique, ix.indisprimary + ORDER BY i.relname + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + Object[] columnArray = (Object[]) rs.getArray("columns").getArray(); + List columns = new ArrayList<>(); + for (Object column : columnArray) { + columns.add(String.valueOf(column)); + } + result.add(new IndexInfo( + rs.getString("index_name"), + columns, + rs.getBoolean("is_unique"), + rs.getBoolean("is_primary"), + null, + rs.getString("index_type"), + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT + tc.constraint_name, + kcu.column_name, + ccu.table_name AS ref_table, + ccu.column_name AS ref_column + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + JOIN information_schema.constraint_column_usage ccu + ON tc.constraint_name = ccu.constraint_name + AND tc.table_schema = ccu.table_schema + WHERE tc.constraint_type = 'FOREIGN KEY' + AND tc.table_schema = ? + AND tc.table_name = ? + ORDER BY tc.constraint_name + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("constraint_name"), + rs.getString("column_name"), + rs.getString("ref_table"), + rs.getString("ref_column") + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT trigger_name, event_manipulation, action_timing + FROM information_schema.triggers + WHERE trigger_schema = ? AND event_object_table = ? + ORDER BY trigger_name + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo( + rs.getString("trigger_name"), + rs.getString("event_manipulation"), + rs.getString("action_timing") + )); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::getResultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return "SET search_path TO " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + public static void main(String[] args) { + new JsonRpcServer(new GaussdbAgent()).run(); + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:opengauss://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private Object getResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value; + switch (sqlType) { + case Types.BIGINT: + value = rs.getLong(index); + break; + case Types.INTEGER: + case Types.SMALLINT: + case Types.TINYINT: + value = rs.getInt(index); + break; + case Types.FLOAT: + case Types.REAL: + value = rs.getFloat(index); + break; + case Types.DOUBLE: + value = rs.getDouble(index); + break; + case Types.DECIMAL: + case Types.NUMERIC: + value = rs.getBigDecimal(index); + break; + case Types.BOOLEAN: + case Types.BIT: + value = rs.getBoolean(index); + break; + default: + value = rs.getString(index); + break; + } + return rs.wasNull() ? null : value; + }); + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + private static String quoteQualifiedIdentifier(String schema, String name) { + return JdbcIdentifiers.INSTANCE.doubleQuote(schema) + "." + JdbcIdentifiers.INSTANCE.doubleQuote(name); + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } +} diff --git a/agents/drivers/gaussdb/src/test/java/com/dbx/agent/gaussdb/GaussdbAgentTest.java b/agents/drivers/gaussdb/src/test/java/com/dbx/agent/gaussdb/GaussdbAgentTest.java new file mode 100644 index 0000000000..497718ea8c --- /dev/null +++ b/agents/drivers/gaussdb/src/test/java/com/dbx/agent/gaussdb/GaussdbAgentTest.java @@ -0,0 +1,43 @@ +package com.dbx.agent.gaussdb; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import com.dbx.agent.test.TestSupport; +import java.util.Arrays; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class GaussdbAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new GaussdbAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } + + @Test + void agentExtendsBaseDatabaseAgent() { + Assertions.assertTrue(BaseDatabaseAgent.class.isAssignableFrom(GaussdbAgent.class)); + } + + @Test + void readsViewSourceWithQuotedRegclassParameter() { + GaussdbAgent agent = new GaussdbAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.getObjectSource("bad\"schema", "view's name", "VIEW"); + + Assertions.assertEquals( + Arrays.asList( + "SELECT pg_get_viewdef(to_regclass(?), true)", + "param:1=\"bad\"\"schema\".\"view's name\"" + ), + JdbcMetadataSqlFake.statements + ); + } +} diff --git a/agents/drivers/gbase8a/build.gradle b/agents/drivers/gbase8a/build.gradle new file mode 100644 index 0000000000..a066170eff --- /dev/null +++ b/agents/drivers/gbase8a/build.gradle @@ -0,0 +1,12 @@ +dependencies { + implementation files('libs/gbase-connector-java-9.5.0.10-build1-bin.jar') +} + +tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'GBase 8a', + 'Main-Class': 'com.dbx.agent.gbase8a.Gbase8aAgent' + ) + } +} diff --git a/agents/drivers/gbase8a/libs/gbase-connector-java-9.5.0.10-build1-bin.jar b/agents/drivers/gbase8a/libs/gbase-connector-java-9.5.0.10-build1-bin.jar new file mode 100644 index 0000000000..f8c6e016d9 Binary files /dev/null and b/agents/drivers/gbase8a/libs/gbase-connector-java-9.5.0.10-build1-bin.jar differ diff --git a/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java b/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java new file mode 100644 index 0000000000..ecc0b5d9c6 --- /dev/null +++ b/agents/drivers/gbase8a/src/main/java/com/dbx/agent/gbase8a/Gbase8aAgent.java @@ -0,0 +1,66 @@ +package com.dbx.agent.gbase8a; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.TableInfo; + +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; + +public final class Gbase8aAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile GBASE8A_PROFILE = new JdbcAgentProfile( + "com.gbase.jdbc.Driver", + "jdbc:gbase://{host}:{port}/{database}?useSSL=false", + 5258, + false, + java.util.Collections.emptySet(), + java.util.Arrays.asList("TABLE", "VIEW", "BASE TABLE"), + "`", + "USE", + true, + false, + false, + false + ); + + public Gbase8aAgent() { + super(GBASE8A_PROFILE); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql; + if (schema != null && !schema.trim().isEmpty()) { + sql = "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME"; + } else { + sql = "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA NOT IN ('information_schema', 'performance_schema', 'gclusterdb', 'gctmpdb') ORDER BY TABLE_SCHEMA, TABLE_NAME"; + } + try (PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + if (schema != null && !schema.trim().isEmpty()) { + stmt.setString(1, schema); + } + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String tableType = rs.getString("TABLE_TYPE"); + if ("BASE TABLE".equals(tableType)) { + tableType = "TABLE"; + } + result.add(new TableInfo(rs.getString("TABLE_NAME"), tableType, null)); + } + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + return result; + }); + } + + public static void main(String[] args) { + new JsonRpcServer(new Gbase8aAgent()).run(); + } +} diff --git a/agents/drivers/gbase8s/build.gradle b/agents/drivers/gbase8s/build.gradle new file mode 100644 index 0000000000..bc2f0fd973 --- /dev/null +++ b/agents/drivers/gbase8s/build.gradle @@ -0,0 +1,12 @@ +dependencies { + implementation 'com.gbasedbt:jdbc:3.6.3.2' +} + +tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'GBase 8s', + 'Main-Class': 'com.dbx.agent.gbase8s.Gbase8sAgent' + ) + } +} diff --git a/agents/drivers/gbase8s/libs/gbasedbt-jdbc.jar b/agents/drivers/gbase8s/libs/gbasedbt-jdbc.jar new file mode 100644 index 0000000000..5b0f899270 Binary files /dev/null and b/agents/drivers/gbase8s/libs/gbasedbt-jdbc.jar differ diff --git a/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java b/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java new file mode 100644 index 0000000000..6824fca345 --- /dev/null +++ b/agents/drivers/gbase8s/src/main/java/com/dbx/agent/gbase8s/Gbase8sAgent.java @@ -0,0 +1,374 @@ +package com.dbx.agent.gbase8s; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public final class Gbase8sAgent extends ConfiguredJdbcAgent { + private static final long METADATA_CACHE_TTL_MILLIS = 10_000L; + + public static final JdbcAgentProfile GBASE8S_PROFILE = new JdbcAgentProfile( + "com.gbasedbt.jdbc.Driver", + "jdbc:gbasedbt-sqli://{host}:{port}/{database}:GBASEDBTSERVER=gbase8s", + 9088, + true + ); + + private final Object metadataCacheLock = new Object(); + private long databaseCacheTimeMillis; + private List databaseCache = Collections.emptyList(); + private String schemaCacheCatalog = ""; + private long schemaCacheTimeMillis; + private List schemaCache = Collections.emptyList(); + private String tableCacheCatalog = ""; + private String tableCacheSchema = ""; + private long tableCacheTimeMillis; + private List tableCache = Collections.emptyList(); + + public Gbase8sAgent() { + super(GBASE8S_PROFILE); + } + + public static String buildUrl(ConnectParams params) { + if (!params.getConnection_string().trim().isEmpty()) { + return params.getConnection_string(); + } + String extraParams = trimEnd(trimStart(params.getUrl_params().trim(), ':', ';'), ';'); + String database = params.getDatabase().trim().isEmpty() ? "sysmaster" : params.getDatabase().trim(); + String serverParam = containsIgnoreCase(extraParams, "GBASEDBTSERVER=") + ? "" + : "GBASEDBTSERVER=" + getGbaseServer(params); + List jdbcParams = new ArrayList<>(); + if (!serverParam.isBlank()) { + jdbcParams.add(serverParam); + } + if (!extraParams.isBlank()) { + jdbcParams.add(extraParams); + } + return "jdbc:gbasedbt-sqli://" + params.getHost() + ":" + port(params) + "/" + database + ":" + + String.join(";", jdbcParams); + } + + private static String getGbaseServer(ConnectParams params) { + if (params.getGbase_server() != null && !params.getGbase_server().trim().isEmpty()) { + return params.getGbase_server().trim(); + } + return defaultGbaseServer(params.getHost()); + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + protected void afterConnect(ConnectParams params, Connection connection) { + super.afterConnect(params, connection); + clearMetadataCache(); + } + + @Override + public void disconnect() { + clearMetadataCache(); + super.disconnect(); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + QueryResult result = super.executeQuery(sql, schema, options); + if (mayChangeMetadata(sql)) { + clearMetadataCache(); + } + return result; + } + + @Override + public List listDatabases() { + List cached = cachedDatabases(); + if (cached != null) { + return cached; + } + List names = queryDatabaseNamesInCatalog("sysmaster", "SELECT name FROM sysdatabases ORDER BY name"); + if (names.isEmpty()) { + names = queryDatabaseNames("SELECT name FROM sysmaster:sysdatabases ORDER BY name"); + } + if (names.isEmpty()) { + names = queryDatabaseNames("SELECT name FROM sysdatabases ORDER BY name"); + } + if (names.isEmpty()) { + return super.listDatabases(); + } + List result = new ArrayList<>(); + for (String name : names) { + result.add(new DatabaseInfo(name)); + } + cacheDatabases(result); + return result; + } + + @Override + public List listSchemas() { + try { + String catalog = currentCatalog(); + List cached = cachedSchemas(catalog); + if (cached != null) { + return cached; + } + Set schemas = new LinkedHashSet<>(); + try (PreparedStatement stmt = requireConnection().prepareStatement( + "SELECT DISTINCT owner FROM systables WHERE tabid >= 100 AND tabtype IN ('T', 'V') ORDER BY owner" + ); ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String owner = trim(rs.getString("owner")); + if (!owner.isEmpty()) { + schemas.add(owner); + } + } + } + List result = new ArrayList<>(schemas); + Collections.sort(result); + cacheSchemas(catalog, result); + return result; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @Override + public List listTables(String schema) { + try { + String catalog = currentCatalog(); + List cached = cachedTables(catalog, schema); + if (cached != null) { + return cached; + } + List result = new ArrayList<>(); + String owner = trim(schema); + String sql = "SELECT tabname, tabtype FROM systables WHERE tabid >= 100 AND tabtype IN ('T', 'V')"; + if (!owner.isEmpty()) { + sql += " AND owner = ?"; + } + sql += " ORDER BY tabname"; + try (PreparedStatement stmt = requireConnection().prepareStatement(sql)) { + if (!owner.isEmpty()) { + stmt.setString(1, owner); + } + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + trim(rs.getString("tabname")), + tableType(rs.getString("tabtype")) + )); + } + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + cacheTables(catalog, schema, result); + return result; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + public static void main(String[] args) { + new JsonRpcServer(new Gbase8sAgent()).run(); + } + + private static int port(ConnectParams params) { + return params.getPort() > 0 ? params.getPort() : GBASE8S_PROFILE.getDefaultPort(); + } + + private static String defaultGbaseServer(String host) { + return isIpAddress(host) ? "gbase8s" : host; + } + + private static boolean isIpAddress(String host) { + return host.matches("\\d{1,3}(\\.\\d{1,3}){3}") || host.contains(":"); + } + + private static String trimStart(String value, char... chars) { + int start = 0; + while (start < value.length() && contains(chars, value.charAt(start))) { + start++; + } + return value.substring(start); + } + + private static String trimEnd(String value, char... chars) { + int end = value.length(); + while (end > 0 && contains(chars, value.charAt(end - 1))) { + end--; + } + return value.substring(0, end); + } + + private static boolean contains(char[] chars, char value) { + for (char ch : chars) { + if (ch == value) { + return true; + } + } + return false; + } + + private static boolean containsIgnoreCase(String value, String needle) { + return value.toLowerCase(Locale.ROOT).contains(needle.toLowerCase(Locale.ROOT)); + } + + private List queryDatabaseNames(String sql) { + try { + return queryDatabaseNames(requireConnection(), sql); + } catch (Exception ignored) { + return Collections.emptyList(); + } + } + + private List queryDatabaseNamesInCatalog(String catalog, String sql) { + try { + Connection connection = requireConnection(); + String previousCatalog = ""; + try { + previousCatalog = trim(connection.getCatalog()); + } catch (Exception ignored) { + } + connection.setCatalog(catalog); + try { + return queryDatabaseNames(connection, sql); + } finally { + if (!previousCatalog.isEmpty()) { + try { + connection.setCatalog(previousCatalog); + } catch (Exception ignored) { + } + } + } + } catch (Exception ignored) { + return Collections.emptyList(); + } + } + + private static List queryDatabaseNames(Connection connection, String sql) throws Exception { + Set names = new LinkedHashSet<>(); + try (PreparedStatement stmt = connection.prepareStatement(sql); ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String name = trim(rs.getString(1)); + if (!name.isEmpty()) { + names.add(name); + } + } + } + List result = new ArrayList<>(names); + Collections.sort(result); + return result; + } + + private static String tableType(String tabtype) { + return "V".equalsIgnoreCase(trim(tabtype)) ? "VIEW" : "TABLE"; + } + + private static String trim(String value) { + return value == null ? "" : value.trim(); + } + + private String currentCatalog() { + try { + return trim(requireConnection().getCatalog()); + } catch (Exception ignored) { + return ""; + } + } + + private List cachedDatabases() { + synchronized (metadataCacheLock) { + if (cacheFresh(databaseCacheTimeMillis) && !databaseCache.isEmpty()) { + return new ArrayList<>(databaseCache); + } + } + return null; + } + + private void cacheDatabases(List databases) { + synchronized (metadataCacheLock) { + databaseCache = new ArrayList<>(databases); + databaseCacheTimeMillis = System.currentTimeMillis(); + } + } + + private List cachedSchemas(String catalog) { + synchronized (metadataCacheLock) { + if (cacheFresh(schemaCacheTimeMillis) && schemaCacheCatalog.equals(catalog)) { + return new ArrayList<>(schemaCache); + } + } + return null; + } + + private void cacheSchemas(String catalog, List schemas) { + synchronized (metadataCacheLock) { + schemaCacheCatalog = catalog; + schemaCache = new ArrayList<>(schemas); + schemaCacheTimeMillis = System.currentTimeMillis(); + } + } + + private List cachedTables(String catalog, String schema) { + String owner = trim(schema); + synchronized (metadataCacheLock) { + if (cacheFresh(tableCacheTimeMillis) && tableCacheCatalog.equals(catalog) && tableCacheSchema.equals(owner)) { + return new ArrayList<>(tableCache); + } + } + return null; + } + + private void cacheTables(String catalog, String schema, List tables) { + synchronized (metadataCacheLock) { + tableCacheCatalog = catalog; + tableCacheSchema = trim(schema); + tableCache = new ArrayList<>(tables); + tableCacheTimeMillis = System.currentTimeMillis(); + } + } + + private boolean cacheFresh(long timeMillis) { + return timeMillis > 0 && System.currentTimeMillis() - timeMillis <= METADATA_CACHE_TTL_MILLIS; + } + + private void clearMetadataCache() { + synchronized (metadataCacheLock) { + databaseCacheTimeMillis = 0; + databaseCache = Collections.emptyList(); + schemaCacheCatalog = ""; + schemaCacheTimeMillis = 0; + schemaCache = Collections.emptyList(); + tableCacheCatalog = ""; + tableCacheSchema = ""; + tableCacheTimeMillis = 0; + tableCache = Collections.emptyList(); + } + } + + private static boolean mayChangeMetadata(String sql) { + String normalized = trim(sql).toLowerCase(Locale.ROOT); + return normalized.startsWith("create ") + || normalized.startsWith("drop ") + || normalized.startsWith("alter ") + || normalized.startsWith("rename ") + || normalized.startsWith("truncate "); + } +} diff --git a/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java b/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java new file mode 100644 index 0000000000..32183511a2 --- /dev/null +++ b/agents/drivers/gbase8s/src/test/java/com/dbx/agent/gbase8s/Gbase8sAgentTest.java @@ -0,0 +1,98 @@ +package com.dbx.agent.gbase8s; + +import com.dbx.agent.ConnectParams; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class Gbase8sAgentTest { + @Test + void declaresGbase8sProfile() { + Gbase8sAgent agent = new Gbase8sAgent(); + + Assertions.assertEquals("com.gbasedbt.jdbc.Driver", agent.getProfile().getDriverClass()); + Assertions.assertEquals("jdbc:gbasedbt-sqli://{host}:{port}/{database}:GBASEDBTSERVER=gbase8s", agent.getProfile().getUrlTemplate()); + Assertions.assertEquals(9088, agent.getProfile().getDefaultPort()); + Assertions.assertTrue(agent.getProfile().getSkipExecutionContext()); + } + + @Test + void buildsGbase8sJdbcUrlWithExplicitServerAndLocaleParameters() { + String url = Gbase8sAgent.buildUrl( + new ConnectParams( + "172.26.128.159", + 20013, + "testdb", + "", + "", + "GBASEDBTSERVER=gbase01;CLIENT_LOCALE=zh_cn.utf8;DB_LOCALE=zh_cn.utf8", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:gbasedbt-sqli://172.26.128.159:20013/testdb:GBASEDBTSERVER=gbase01;CLIENT_LOCALE=zh_cn.utf8;DB_LOCALE=zh_cn.utf8", + url + ); + } + + @Test + void fallsBackToHostAsGbaseServerWhenNoExplicitServerIsConfigured() { + String url = Gbase8sAgent.buildUrl( + new ConnectParams( + "gbase-host", + 9088, + "sysmaster", + "", + "", + "", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:gbasedbt-sqli://gbase-host:9088/sysmaster:GBASEDBTSERVER=gbase-host", + url + ); + } + + @Test + void fallsBackToGbaseServerNameWhenHostIsAnIpAddress() { + String url = Gbase8sAgent.buildUrl( + new ConnectParams( + "172.26.128.159", + 0, + "sysmaster", + "", + "", + "", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:gbasedbt-sqli://172.26.128.159:9088/sysmaster:GBASEDBTSERVER=gbase8s", + url + ); + } + + @Test + void usesConnectionStringWhenConfigured() { + String url = Gbase8sAgent.buildUrl( + new ConnectParams( + "ignored", + 0, + "", + "", + "", + "", + "jdbc:gbasedbt-sqli://db.example.com:20013/app:GBASEDBTSERVER=gbase01", + false + ) + ); + + Assertions.assertEquals("jdbc:gbasedbt-sqli://db.example.com:20013/app:GBASEDBTSERVER=gbase01", url); + } +} diff --git a/agents/drivers/goldendb/build.gradle b/agents/drivers/goldendb/build.gradle new file mode 100644 index 0000000000..a200788437 --- /dev/null +++ b/agents/drivers/goldendb/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.mysql:mysql-connector-j:9.2.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'GoldenDB', 'Main-Class': 'com.dbx.agent.goldendb.GoldendbAgent') + } +} diff --git a/agents/drivers/goldendb/libs/.gitkeep b/agents/drivers/goldendb/libs/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java b/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java new file mode 100644 index 0000000000..04590b0cc9 --- /dev/null +++ b/agents/drivers/goldendb/src/main/java/com/dbx/agent/goldendb/GoldendbAgent.java @@ -0,0 +1,409 @@ +package com.dbx.agent.goldendb; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Types; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.HashSet; + +public final class GoldendbAgent extends BaseDatabaseAgent { + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("com.mysql.cj.jdbc.Driver"); + connection = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword()); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("com.mysql.cj.jdbc.Driver"); + try (Connection conn = DriverManager.getConnection(buildUrl(params), params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + "SELECT TABLE_NAME, TABLE_TYPE, TABLE_COMMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String tableType = rs.getString("TABLE_TYPE"); + if ("BASE TABLE".equals(tableType)) { + tableType = "TABLE"; + } + result.add(new TableInfo(rs.getString("TABLE_NAME"), tableType, emptyToNull(rs.getString("TABLE_COMMENT")))); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + "SELECT ROUTINE_NAME, ROUTINE_TYPE FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = ? ORDER BY ROUTINE_NAME" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String quotedName = JdbcIdentifiers.INSTANCE.backtick(name); + String sql = switch (objectType.toUpperCase(Locale.ROOT)) { + case "VIEW" -> "SHOW CREATE VIEW " + quotedName; + case "PROCEDURE" -> "SHOW CREATE PROCEDURE " + quotedName; + case "FUNCTION" -> "SHOW CREATE FUNCTION " + quotedName; + default -> throw new IllegalArgumentException("Unsupported object type: " + objectType); + }; + + String source = ""; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + if (rs.next()) { + int sourceIndex = "VIEW".equals(objectType.toUpperCase(Locale.ROOT)) ? 2 : 3; + String value = rs.getString(sourceIndex); + source = value == null ? "" : value; + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = new HashSet<>(); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + """ + SELECT COLUMN_NAME + FROM information_schema.KEY_COLUMN_USAGE + WHERE TABLE_SCHEMA = ? + AND TABLE_NAME = ? + AND CONSTRAINT_NAME = 'PRIMARY' + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + primaryKeys.add(rs.getString("COLUMN_NAME")); + } + } + } + + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + """ + SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_DEFAULT, EXTRA, COLUMN_COMMENT, + NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH + FROM information_schema.COLUMNS + WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? + ORDER BY ORDINAL_POSITION + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String columnName = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + columnName, + rs.getString("COLUMN_TYPE"), + "YES".equals(rs.getString("IS_NULLABLE")), + rs.getString("COLUMN_DEFAULT"), + primaryKeys.contains(columnName), + rs.getString("EXTRA"), + emptyToNull(rs.getString("COLUMN_COMMENT")), + integerOrNull(rs, "NUMERIC_PRECISION"), + integerOrNull(rs, "NUMERIC_SCALE"), + numberToIntOrNull(rs, "CHARACTER_MAXIMUM_LENGTH") + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + Map> indexMap = new TreeMap<>(); + Map uniqueMap = new HashMap<>(); + Map typeMap = new HashMap<>(); + String quotedTable = JdbcIdentifiers.INSTANCE.backtick(table); + String quotedSchema = JdbcIdentifiers.INSTANCE.backtick(schema); + + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW INDEX FROM " + quotedTable + " FROM " + quotedSchema)) { + while (rs.next()) { + String indexName = rs.getString("Key_name"); + String columnName = rs.getString("Column_name"); + int seqInIndex = rs.getInt("Seq_in_index"); + int nonUnique = rs.getInt("Non_unique"); + String indexType = rs.getString("Index_type"); + + indexMap.computeIfAbsent(indexName, ignored -> new ArrayList<>()) + .add(new IndexColumn(seqInIndex, columnName)); + uniqueMap.put(indexName, nonUnique == 0); + typeMap.put(indexName, indexType); + } + } + + List result = new ArrayList<>(); + for (Map.Entry> entry : indexMap.entrySet()) { + List columns = entry.getValue(); + columns.sort((left, right) -> Integer.compare(left.ordinal, right.ordinal)); + List columnNames = new ArrayList<>(); + for (IndexColumn column : columns) { + columnNames.add(column.name); + } + + String indexName = entry.getKey(); + result.add(new IndexInfo( + indexName, + columnNames, + Boolean.TRUE.equals(uniqueMap.get(indexName)), + "PRIMARY".equals(indexName), + null, + typeMap.get(indexName), + null, + null + )); + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + """ + SELECT CONSTRAINT_NAME, COLUMN_NAME, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME + FROM information_schema.KEY_COLUMN_USAGE + WHERE TABLE_SCHEMA = ? + AND TABLE_NAME = ? + AND REFERENCED_TABLE_NAME IS NOT NULL + ORDER BY CONSTRAINT_NAME + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("CONSTRAINT_NAME"), + rs.getString("COLUMN_NAME"), + rs.getString("REFERENCED_TABLE_NAME"), + rs.getString("REFERENCED_COLUMN_NAME") + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement( + """ + SELECT TRIGGER_NAME, EVENT_MANIPULATION, ACTION_TIMING + FROM information_schema.TRIGGERS + WHERE TRIGGER_SCHEMA = ? AND EVENT_OBJECT_TABLE = ? + ORDER BY TRIGGER_NAME + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo( + rs.getString("TRIGGER_NAME"), + rs.getString("EVENT_MANIPULATION"), + rs.getString("ACTION_TIMING") + )); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::resultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + JdbcIdentifiers.INSTANCE.backtick(schema); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value; + switch (sqlType) { + case Types.BIGINT: + value = rs.getLong(index); + break; + case Types.INTEGER: + case Types.SMALLINT: + case Types.TINYINT: + value = rs.getInt(index); + break; + case Types.FLOAT: + case Types.REAL: + value = rs.getFloat(index); + break; + case Types.DOUBLE: + value = rs.getDouble(index); + break; + case Types.DECIMAL: + case Types.NUMERIC: + value = rs.getBigDecimal(index); + break; + case Types.BOOLEAN: + case Types.BIT: + value = rs.getBoolean(index); + break; + default: + value = rs.getString(index); + break; + } + return rs.wasNull() ? null : value; + }); + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:mysql://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase() + + "?useSSL=false&allowPublicKeyRetrieval=true"; + } + + private static String emptyToNull(String value) { + return value == null || value.isEmpty() ? null : value; + } + + private static Integer integerOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Integer ? (Integer) value : null; + } + + private static Integer numberToIntOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static final class IndexColumn { + private final int ordinal; + private final String name; + + private IndexColumn(int ordinal, String name) { + this.ordinal = ordinal; + this.name = name; + } + } + + public static void main(String[] args) { + new JsonRpcServer(new GoldendbAgent()).run(); + } +} diff --git a/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java b/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java new file mode 100644 index 0000000000..ba3d9e1ac6 --- /dev/null +++ b/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentMetadataTest.java @@ -0,0 +1,25 @@ +package com.dbx.agent.goldendb; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.test.JdbcAgentFake; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import java.sql.Connection; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class GoldendbAgentMetadataTest { + @Test + void quotesSchemaAndTableIdentifiersInIndexMetadataSql() { + GoldendbAgent agent = new GoldendbAgent(); + Connection fake = JdbcMetadataSqlFake.connection(); + TestSupport.setPrivateConnection(agent, fake); + + agent.listIndexes("bad`schema", "bad`table"); + + Assertions.assertEquals( + List.of("SHOW INDEX FROM `bad``table` FROM `bad``schema`"), + JdbcMetadataSqlFake.statements + ); + } +} diff --git a/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentTest.java b/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentTest.java new file mode 100644 index 0000000000..c22fa1d467 --- /dev/null +++ b/agents/drivers/goldendb/src/test/java/com/dbx/agent/goldendb/GoldendbAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.goldendb; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class GoldendbAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new GoldendbAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } +} diff --git a/agents/drivers/h2/build.gradle b/agents/drivers/h2/build.gradle new file mode 100644 index 0000000000..070380fd1e --- /dev/null +++ b/agents/drivers/h2/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.h2database:h2:2.3.232' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'H2', 'Main-Class': 'com.dbx.agent.h2.H2Agent') + } +} diff --git a/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java b/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java new file mode 100644 index 0000000000..88746621a7 --- /dev/null +++ b/agents/drivers/h2/src/main/java/com/dbx/agent/h2/H2Agent.java @@ -0,0 +1,349 @@ +package com.dbx.agent.h2; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class H2Agent extends AbstractJdbcAgent { + private String databaseName = ""; + + @Override + protected String driverClass() { + return "org.h2.Driver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + protected void afterConnect(ConnectParams params, Connection connection) { + databaseName = params.getDatabase(); + } + + @Override + public List listDatabases() { + return List.of(new DatabaseInfo(databaseName.isBlank() ? "default" : databaseName)); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA ORDER BY SCHEMA_NAME" + ); + ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(rs.getString(1)); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT TABLE_NAME, TABLE_TYPE + FROM INFORMATION_SCHEMA.TABLES + WHERE TABLE_SCHEMA = ? + ORDER BY TABLE_NAME + """ + )) { + stmt.setString(1, effectiveSchema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String tableType = rs.getString("TABLE_TYPE"); + if ("BASE TABLE".equals(tableType)) { + tableType = "TABLE"; + } + result.add(new TableInfo(rs.getString("TABLE_NAME"), tableType, null)); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + + try (var stmt = requireConnected().prepareStatement( + "SELECT ROUTINE_NAME, ROUTINE_TYPE FROM INFORMATION_SCHEMA.ROUTINES WHERE ROUTINE_SCHEMA = ? ORDER BY ROUTINE_NAME" + )) { + stmt.setString(1, effectiveSchema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + String sql = switch (objectType.toUpperCase(Locale.ROOT)) { + case "VIEW" -> "SELECT VIEW_DEFINITION FROM INFORMATION_SCHEMA.VIEWS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?"; + case "FUNCTION", "PROCEDURE" -> "SELECT ROUTINE_DEFINITION FROM INFORMATION_SCHEMA.ROUTINES WHERE ROUTINE_SCHEMA = ? AND ROUTINE_NAME = ?"; + default -> throw new IllegalArgumentException("Unsupported object type: " + objectType); + }; + + String source = ""; + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, effectiveSchema); + stmt.setString(2, name); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + String value = rs.getString(1); + source = value == null ? "" : value; + } + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + Set primaryKeys = new HashSet<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT ic.COLUMN_NAME + FROM INFORMATION_SCHEMA.INDEX_COLUMNS ic + JOIN INFORMATION_SCHEMA.INDEXES i + ON ic.INDEX_SCHEMA = i.INDEX_SCHEMA AND ic.INDEX_NAME = i.INDEX_NAME + WHERE ic.TABLE_SCHEMA = ? AND ic.TABLE_NAME = ? + AND i.INDEX_TYPE_NAME = 'PRIMARY KEY' + """ + )) { + stmt.setString(1, effectiveSchema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + primaryKeys.add(rs.getString("COLUMN_NAME")); + } + } + } + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT, + NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH + FROM INFORMATION_SCHEMA.COLUMNS + WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? + ORDER BY ORDINAL_POSITION + """ + )) { + stmt.setString(1, effectiveSchema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String columnName = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + columnName, + rs.getString("DATA_TYPE"), + "YES".equals(rs.getString("IS_NULLABLE")), + rs.getString("COLUMN_DEFAULT"), + primaryKeys.contains(columnName), + null, + null, + intOrNull(rs, "NUMERIC_PRECISION"), + intOrNull(rs, "NUMERIC_SCALE"), + intOrNull(rs, "CHARACTER_MAXIMUM_LENGTH") + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + Map> indexMap = new LinkedHashMap<>(); + Map uniqueMap = new HashMap<>(); + Map primaryMap = new HashMap<>(); + Map typeMap = new HashMap<>(); + + try (var stmt = requireConnected().prepareStatement( + """ + SELECT i.INDEX_NAME, ic.COLUMN_NAME, ic.IS_UNIQUE, i.INDEX_TYPE_NAME + FROM INFORMATION_SCHEMA.INDEX_COLUMNS ic + JOIN INFORMATION_SCHEMA.INDEXES i + ON ic.INDEX_SCHEMA = i.INDEX_SCHEMA AND ic.INDEX_NAME = i.INDEX_NAME + WHERE ic.TABLE_SCHEMA = ? AND ic.TABLE_NAME = ? + ORDER BY i.INDEX_NAME, ic.ORDINAL_POSITION + """ + )) { + stmt.setString(1, effectiveSchema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String indexName = rs.getString("INDEX_NAME"); + String columnName = rs.getString("COLUMN_NAME"); + String indexType = rs.getString("INDEX_TYPE_NAME"); + + indexMap.computeIfAbsent(indexName, ignored -> new ArrayList<>()).add(columnName); + uniqueMap.put(indexName, rs.getBoolean("IS_UNIQUE")); + primaryMap.put(indexName, "PRIMARY KEY".equals(indexType)); + typeMap.put(indexName, indexType == null ? "" : indexType); + } + } + } + + List result = new ArrayList<>(); + for (Map.Entry> entry : indexMap.entrySet()) { + String name = entry.getKey(); + result.add(new IndexInfo( + name, + entry.getValue(), + uniqueMap.getOrDefault(name, false), + primaryMap.getOrDefault(name, false), + null, + typeMap.get(name), + null, + null + )); + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT FK_NAME, FKCOLUMN_NAME, PKTABLE_NAME, PKCOLUMN_NAME + FROM INFORMATION_SCHEMA.CROSS_REFERENCES + WHERE FKTABLE_SCHEMA = ? AND FKTABLE_NAME = ? + ORDER BY FK_NAME + """ + )) { + stmt.setString(1, effectiveSchema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("FK_NAME"), + rs.getString("FKCOLUMN_NAME"), + rs.getString("PKTABLE_NAME"), + rs.getString("PKCOLUMN_NAME") + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + String effectiveSchema = resolveSchema(schema); + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT TRIGGER_NAME, EVENT_MANIPULATION, ACTION_TIMING + FROM INFORMATION_SCHEMA.TRIGGERS + WHERE TRIGGER_SCHEMA = ? AND EVENT_OBJECT_TABLE = ? + ORDER BY TRIGGER_NAME + """ + )) { + stmt.setString(1, effectiveSchema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo( + rs.getString("TRIGGER_NAME"), + rs.getString("EVENT_MANIPULATION"), + rs.getString("ACTION_TIMING") + )); + } + } + } + return result; + }); + } + + @Override + public String setSchemaSQL(String schema) { + return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> JdbcExecutor.INSTANCE.defaultResultValue(rs, index, sqlType)); + } + + static String buildUrl(ConnectParams params) { + String connectionString = params.getConnection_string(); + if (connectionString != null && !connectionString.trim().isEmpty()) { + return connectionString.trim(); + } + if (params.getHost().isBlank()) { + return "jdbc:h2:" + params.getDatabase(); + } + return "jdbc:h2:tcp://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private static String resolveSchema(String schema) { + if ("PUBLIC".equalsIgnoreCase(schema) || "INFORMATION_SCHEMA".equalsIgnoreCase(schema)) { + return schema.toUpperCase(Locale.ROOT); + } + return "PUBLIC"; + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + public static void main(String[] args) { + new JsonRpcServer(new H2Agent()).run(); + } +} diff --git a/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java b/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java new file mode 100644 index 0000000000..35f11cd023 --- /dev/null +++ b/agents/drivers/h2/src/test/java/com/dbx/agent/h2/H2AgentTest.java @@ -0,0 +1,114 @@ +package com.dbx.agent.h2; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcExecutionBehaviorTest; +import com.dbx.agent.test.JdbcMetadataBehaviorTest; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class H2AgentMigrationTest { + @Test + void agentExtendsBaseDatabaseAgent() { + Assertions.assertTrue(BaseDatabaseAgent.class.isAssignableFrom(H2Agent.class)); + } + + @Test + void buildUrlUsesExplicitConnectionString() { + ConnectParams params = new ConnectParams( + "127.0.0.1", + 9092, + "test", + "sa", + "", + "", + "jdbc:h2:file:/tmp/dbx-h2-test;AUTO_SERVER=TRUE", + false + ); + + Assertions.assertEquals("jdbc:h2:file:/tmp/dbx-h2-test;AUTO_SERVER=TRUE", H2Agent.buildUrl(params)); + } + + @Test + void buildUrlKeepsTcpModeWhenNoConnectionString() { + ConnectParams params = new ConnectParams("127.0.0.1", 9092, "test", "sa", "", "", "", false); + + Assertions.assertEquals("jdbc:h2:tcp://127.0.0.1:9092/test", H2Agent.buildUrl(params)); + } +} + +class H2ExecutionBehaviorTest extends JdbcExecutionBehaviorTest { + @Override + protected DatabaseAgent createConnectedAgent(String databaseName) { + return H2AgentTestSupport.createH2Agent(databaseName); + } + + @Override + protected String resultSetSql() { + return "CALL 42"; + } + + @Override + protected List expectedResultSetColumns() { + return List.of("42"); + } + + @Override + protected List> expectedResultSetRows() { + return List.of(List.of(42)); + } + + @Override + protected String rowsSql(int rowCount) { + return "SELECT X FROM SYSTEM_RANGE(1, " + rowCount + ")"; + } +} + +class H2MetadataBehaviorTest extends JdbcMetadataBehaviorTest { + @Override + protected DatabaseAgent createConnectedAgent(String databaseName) { + return H2AgentTestSupport.createH2Agent(databaseName); + } + + @Override + protected List metadataFixtureSql() { + return List.of( + "CREATE TABLE BETA_TABLE (ID INT PRIMARY KEY)", + "CREATE TABLE ALPHA_TABLE (ID INT PRIMARY KEY)", + "CREATE TABLE COLUMN_ORDER_SAMPLE (ID INT PRIMARY KEY, NAME VARCHAR(64), CREATED_AT TIMESTAMP)" + ); + } + + @Override + protected String metadataSchema() { + return "PUBLIC"; + } + + @Override + protected List expectedTablesInOrder() { + return List.of("ALPHA_TABLE", "BETA_TABLE", "COLUMN_ORDER_SAMPLE"); + } + + @Override + protected String metadataColumnsTable() { + return "COLUMN_ORDER_SAMPLE"; + } + + @Override + protected List expectedColumnsInOrder() { + return List.of("ID", "NAME", "CREATED_AT"); + } +} + +final class H2AgentTestSupport { + private H2AgentTestSupport() { + } + + static DatabaseAgent createH2Agent(String databaseName) { + H2Agent agent = new H2Agent(); + agent.connect(new ConnectParams("", 0, "mem:" + databaseName + ";DB_CLOSE_DELAY=-1", "", "", "", "", false)); + return agent; + } +} diff --git a/agents/drivers/highgo/build.gradle b/agents/drivers/highgo/build.gradle new file mode 100644 index 0000000000..1b944abe12 --- /dev/null +++ b/agents/drivers/highgo/build.gradle @@ -0,0 +1,16 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.highgo:HgdbJdbc:6.2.4' + // PostGIS JDBC is optional: EwkbWktDecoder reflectively probes for + // org.postgis.PGgeometry / net.postgis.jdbc.PGgeometry and falls back to + // its own EWKB parser when the class is absent. compileOnly keeps the + // shadowJar small while still letting users drop the jar into their + // user-driver path if they prefer the upstream parser. + compileOnly 'net.postgis:postgis-jdbc:2024.1.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': '瀚高 HighGo', 'Main-Class': 'com.dbx.agent.highgo.HighgoAgent') + } +} diff --git a/agents/drivers/highgo/src/main/java/com/dbx/agent/highgo/HighgoAgent.java b/agents/drivers/highgo/src/main/java/com/dbx/agent/highgo/HighgoAgent.java new file mode 100644 index 0000000000..ee9e6b17b4 --- /dev/null +++ b/agents/drivers/highgo/src/main/java/com/dbx/agent/highgo/HighgoAgent.java @@ -0,0 +1,20 @@ +package com.dbx.agent.highgo; + +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.PostgresLikeAgent; +import com.dbx.agent.PostgresLikeAgentProfile; + +public final class HighgoAgent extends PostgresLikeAgent { + public static final PostgresLikeAgentProfile HIGHGO_PROFILE = new PostgresLikeAgentProfile( + "com.highgo.jdbc.Driver", + "jdbc:highgo://{host}:{port}/{database}" + ); + + public HighgoAgent() { + super(HIGHGO_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new HighgoAgent()).run(); + } +} diff --git a/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java b/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java new file mode 100644 index 0000000000..2186312269 --- /dev/null +++ b/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java @@ -0,0 +1,46 @@ +package com.dbx.agent.highgo; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import com.dbx.agent.test.TestSupport; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; + +class HighgoAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new HighgoAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } + + @Test + void declaresHighgoPostgresLikeProfile() { + HighgoAgent agent = new HighgoAgent(); + + Assertions.assertEquals("com.highgo.jdbc.Driver", agent.getProfile().getDriverClass()); + Assertions.assertEquals("jdbc:highgo://{host}:{port}/{database}", agent.getProfile().getUrlTemplate()); + } + + @Test + void readsViewSourceWithQuotedRegclassParameter() { + HighgoAgent agent = new HighgoAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.getObjectSource("bad\"schema", "view's name", "VIEW"); + + Assertions.assertEquals( + Arrays.asList( + "SELECT pg_get_viewdef(to_regclass(?), true)", + "param:1=\"bad\"\"schema\".\"view's name\"" + ), + JdbcMetadataSqlFake.statements + ); + } +} diff --git a/agents/drivers/hive/build.gradle b/agents/drivers/hive/build.gradle new file mode 100644 index 0000000000..8787dd75aa --- /dev/null +++ b/agents/drivers/hive/build.gradle @@ -0,0 +1,11 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'org.apache.hive:hive-jdbc:4.0.1:standalone' + implementation 'org.slf4j:slf4j-nop:1.7.30' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Apache Hive', 'Main-Class': 'com.dbx.agent.hive.HiveAgent') + } +} diff --git a/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java b/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java new file mode 100644 index 0000000000..9a3b407621 --- /dev/null +++ b/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java @@ -0,0 +1,219 @@ +package com.dbx.agent.hive; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Locale; + +public final class HiveAgent extends AbstractJdbcAgent { + + @Override + protected String driverClass() { + return "org.apache.hive.jdbc.HiveDriver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + result.sort(Comparator.comparing(DatabaseInfo::getName)); + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + useSchema(schema); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW TABLES")) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), "TABLE", null)); + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + try { + return getColumnsFromDescribe(schema, table); + } catch (Exception ignored) { + return getColumnsFromMetadata(requireConnected(), schema, table); + } + }); + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + JdbcIdentifiers.INSTANCE.backtick(schema); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + private void useSchema(String schema) throws Exception { + try (java.sql.Statement stmt = requireConnected().createStatement()) { + stmt.execute(setSchemaSQL(schema)); + } + } + + private List getColumnsFromDescribe(String schema, String table) throws Exception { + List result = new ArrayList<>(); + useSchema(schema); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("DESCRIBE " + JdbcIdentifiers.INSTANCE.backtick(table))) { + while (rs.next()) { + String colName = trimToNull(rs.getString(1)); + if (colName == null || colName.startsWith("#")) { + continue; + } + result.add(new ColumnInfo( + colName, + trimToEmpty(rs.getString(2)), + true, + null, + false, + null, + optionalComment(rs), + null, + null, + null + )); + } + } + return result; + } + + private static List getColumnsFromMetadata(Connection conn, String schema, String table) throws Exception { + List result = new ArrayList<>(); + DatabaseMetaData meta = conn.getMetaData(); + try (ResultSet rs = meta.getColumns(null, trimToNull(schema), table, "%")) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + if (trimToNull(name) == null) { + continue; + } + result.add(new ColumnInfo( + name, + trimToEmpty(rs.getString("TYPE_NAME")), + rs.getInt("NULLABLE") != DatabaseMetaData.columnNoNulls, + rs.getString("COLUMN_DEF"), + false, + null, + rs.getString("REMARKS"), + intOrNull(rs, "COLUMN_SIZE"), + intOrNull(rs, "DECIMAL_DIGITS"), + characterLength(rs) + )); + } + } + return result; + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:hive2://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private static String trimToNull(String value) { + if (value == null) { + return null; + } + String trimmed = value.trim(); + return trimmed.isEmpty() ? null : trimmed; + } + + private static String trimToEmpty(String value) { + return value == null ? "" : value.trim(); + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static Integer characterLength(ResultSet rs) throws Exception { + String typeName = rs.getString("TYPE_NAME"); + if (typeName == null) { + return null; + } + String normalized = typeName.toLowerCase(Locale.ROOT); + if (!normalized.contains("char") && !normalized.contains("text") && !normalized.contains("string")) { + return null; + } + return intOrNull(rs, "COLUMN_SIZE"); + } + + private static String optionalComment(ResultSet rs) { + return unchecked(() -> { + try { + String comment = trimToNull(rs.getString(3)); + return comment; + } catch (Exception e) { + return null; + } + }); + } + + public static void main(String[] args) { + new JsonRpcServer(new HiveAgent()).run(); + } +} diff --git a/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentExecutionTest.java b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentExecutionTest.java new file mode 100644 index 0000000000..53f7bb7732 --- /dev/null +++ b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentExecutionTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.hive; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class HiveAgentExecutionTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new HiveAgent(); + } + + @Override + protected String resultSetSql() { + return "MSCK REPAIR TABLE sample"; + } +} diff --git a/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentMetadataTest.java b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentMetadataTest.java new file mode 100644 index 0000000000..0ec4fa8729 --- /dev/null +++ b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentMetadataTest.java @@ -0,0 +1,186 @@ +package com.dbx.agent.hive; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import com.dbx.agent.ColumnInfo; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.util.List; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.LinkedHashMap; +import java.util.Map; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class HiveAgentMetadataTest { + @Test + void quotesSchemaAndTableIdentifiersInMetadataSql() { + HiveAgent agent = new HiveAgent(); + Connection fake = JdbcMetadataSqlFake.connection(); + TestSupport.setPrivateConnection(agent, fake); + + agent.listTables("bad`schema"); + agent.getColumns("bad`schema", "bad`table"); + + Assertions.assertEquals( + List.of( + "USE `bad``schema`", + "SHOW TABLES", + "USE `bad``schema`", + "DESCRIBE `bad``table`" + ), + JdbcMetadataSqlFake.statements + ); + } + + @Test + void fallsBackToJdbcMetadataWhenDescribeCannotCreateResultSet() { + HiveAgent agent = new HiveAgent(); + Connection fake = describeFailingMetadataConnection(); + TestSupport.setPrivateConnection(agent, fake); + + List columns = agent.getColumns("default", "events"); + + Assertions.assertEquals(2, columns.size()); + Assertions.assertEquals("id", columns.get(0).getName()); + Assertions.assertEquals("bigint", columns.get(0).getData_type()); + Assertions.assertFalse(columns.get(0).getIs_nullable()); + Assertions.assertEquals("name", columns.get(1).getName()); + Assertions.assertEquals("string", columns.get(1).getData_type()); + Assertions.assertEquals(Integer.valueOf(255), columns.get(1).getCharacter_maximum_length()); + } + + private static Connection describeFailingMetadataConnection() { + DatabaseMetaData meta = proxy(DatabaseMetaData.class, (method, args) -> { + if ("getColumns".equals(method.getName())) { + return rows( + row( + "COLUMN_NAME", "id", + "TYPE_NAME", "bigint", + "NULLABLE", DatabaseMetaData.columnNoNulls, + "COLUMN_DEF", null, + "REMARKS", "identifier", + "COLUMN_SIZE", 19, + "DECIMAL_DIGITS", 0 + ), + row( + "COLUMN_NAME", "name", + "TYPE_NAME", "string", + "NULLABLE", DatabaseMetaData.columnNullable, + "COLUMN_DEF", null, + "REMARKS", null, + "COLUMN_SIZE", 255, + "DECIMAL_DIGITS", 0 + ) + ); + } + return defaultValue(method.getReturnType()); + }); + return proxy(Connection.class, (method, args) -> { + String name = method.getName(); + if ("createStatement".equals(name)) { + return describeFailingStatement(); + } + if ("getMetaData".equals(name)) { + return meta; + } + if ("isClosed".equals(name)) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static Statement describeFailingStatement() { + return proxy(Statement.class, (method, args) -> { + String name = method.getName(); + if ("execute".equals(name)) { + return false; + } + if ("executeQuery".equals(name)) { + throw new SQLException("Could not create ResultSet: Required field 'type' is unset! Struct:TPrimitiveTypeEntry(type:null)"); + } + return defaultValue(method.getReturnType()); + }); + } + + @SafeVarargs + private static ResultSet rows(Map... rows) { + return proxy(ResultSet.class, new MethodHandler() { + private int index = -1; + + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + index += 1; + return index < rows.length; + } + if ("close".equals(name)) { + return null; + } + Object value = rows[index].get(args[0]); + if ("getString".equals(name)) { + return value == null ? null : String.valueOf(value); + } + if ("getObject".equals(name)) { + return value; + } + if ("getInt".equals(name)) { + return value instanceof Number ? ((Number) value).intValue() : 0; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static Map row(Object... values) { + Map row = new LinkedHashMap<>(); + for (int i = 0; i < values.length; i += 2) { + row.put(String.valueOf(values[i]), values[i + 1]); + } + return row; + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = (proxy, method, args) -> handler.handle(method, args); + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } +} diff --git a/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java new file mode 100644 index 0000000000..36006e7271 --- /dev/null +++ b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java @@ -0,0 +1,11 @@ +package com.dbx.agent.hive; + +import org.junit.jupiter.api.Test; + +class HiveAgentTest { + @Test + void hiveJdbcStandaloneRuntimeClassesAreAvailable() throws ClassNotFoundException { + Class.forName("org.apache.hive.jdbc.HiveDriver"); + Class.forName("org.apache.hive.org.apache.thrift.protocol.TProtocol"); + } +} diff --git a/agents/drivers/informix/build.gradle b/agents/drivers/informix/build.gradle new file mode 100644 index 0000000000..1b38423ce7 --- /dev/null +++ b/agents/drivers/informix/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.ibm.informix:jdbc:4.50.10' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'IBM Informix', 'Main-Class': 'com.dbx.agent.informix.InformixAgent') + } +} diff --git a/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java b/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java new file mode 100644 index 0000000000..d4651ff971 --- /dev/null +++ b/agents/drivers/informix/src/main/java/com/dbx/agent/informix/InformixAgent.java @@ -0,0 +1,429 @@ +package com.dbx.agent.informix; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public final class InformixAgent extends BaseDatabaseAgent { + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + public static String buildJdbcUrl(ConnectParams params) { + String rawUrlParams = params.getUrl_params(); + String extraParams = rawUrlParams == null ? "" : trimEnd(trimStart(rawUrlParams.trim(), ':', ';'), ';'); + String rawDatabase = params.getDatabase(); + String database = rawDatabase == null || rawDatabase.trim().isEmpty() + ? "sysmaster" : rawDatabase.trim(); + + // Determine INFORMIXSERVER: from dedicated field, url_params, or default + String informixServer = params.getInformix_server(); + informixServer = informixServer == null ? "" : informixServer.trim(); + if (informixServer.isEmpty()) { + informixServer = containsIgnoreCase(extraParams, "INFORMIXSERVER=") ? "" : defaultInformixServer(params.getHost()); + } + String serverParam = informixServer.isEmpty() ? "" : "INFORMIXSERVER=" + informixServer; + + // Add CLIENT_LOCALE and DB_LOCALE defaults if not already specified + if (!containsIgnoreCase(extraParams, "CLIENT_LOCALE=")) { + String localeParam = "CLIENT_LOCALE=en_US.utf8"; + extraParams = extraParams.isEmpty() ? localeParam : extraParams + ";" + localeParam; + } + if (!containsIgnoreCase(extraParams, "DB_LOCALE=")) { + String localeParam = "DB_LOCALE=en_US.utf8"; + extraParams = extraParams.isEmpty() ? localeParam : extraParams + ";" + localeParam; + } + + List jdbcParams = new ArrayList<>(); + if (!serverParam.isEmpty()) { + jdbcParams.add(serverParam); + } + if (!extraParams.isEmpty()) { + jdbcParams.add(extraParams); + } + return "jdbc:informix-sqli://" + params.getHost() + ":" + params.getPort() + "/" + database + ":" + + String.join(";", jdbcParams); + } + + private static String defaultInformixServer(String host) { + return isIpAddress(host) ? "informix" : host; + } + + private static boolean isIpAddress(String host) { + return host.matches("\\d{1,3}(\\.\\d{1,3}){3}") || host.contains(":"); + } + + /** + * Map Informix coltype integer codes to type names. + * See Informix documentation for coltype values. + */ + public static String mapColType(int coltype) { + int baseType = coltype % 256; + return switch (baseType) { + case 0 -> "CHAR"; + case 1 -> "SMALLINT"; + case 2 -> "INTEGER"; + case 3 -> "FLOAT"; + case 4 -> "SMALLFLOAT"; + case 5 -> "DECIMAL"; + case 6 -> "SERIAL"; + case 7 -> "DATE"; + case 8 -> "MONEY"; + case 9 -> "NULL"; + case 10 -> "DATETIME"; + case 11 -> "BYTE"; + case 12 -> "TEXT"; + case 13 -> "VARCHAR"; + case 14 -> "INTERVAL"; + case 15 -> "NCHAR"; + case 16 -> "NVARCHAR"; + case 17 -> "INT8"; + case 18 -> "SERIAL8"; + case 19 -> "SET"; + case 20 -> "MULTISET"; + case 21 -> "LIST"; + case 22 -> "ROW"; + case 23 -> "COLLECTION"; + case 40 -> "LVARCHAR"; + case 41 -> "BOOLEAN"; + case 43, 52 -> "BIGINT"; + case 44, 53 -> "BIGSERIAL"; + default -> "UNKNOWN(" + baseType + ")"; + }; + } + + public static Set primaryKeyColumnNumbers(List parts) { + Set result = new HashSet<>(); + for (Integer part : parts) { + if (part == null) { + continue; + } + int value = Math.abs(part); + if (value > 0) { + result.add(value); + } + } + return result; + } + + public static String databaseCatalogSql() { + return "SELECT name FROM sysmaster:sysdatabases ORDER BY name"; + } + + @Override + public void connect(ConnectParams params) { + String url = buildJdbcUrl(params); + uncheckedVoid(() -> { + Class.forName("com.informix.jdbc.IfxDriver"); + try { + connection = DriverManager.getConnection(url, params.getUsername(), params.getPassword()); + } catch (SQLException e) { + throw new SQLException( + "Informix connection failed.\nURL: " + url.replaceAll("//[^@]+@", "//***@") + "\nError: " + e.getMessage(), + e.getSQLState(), e.getErrorCode() + ); + } + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + String url = buildJdbcUrl(params); + return unchecked(() -> { + Class.forName("com.informix.jdbc.IfxDriver"); + try (Connection conn = DriverManager.getConnection(url, params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(databaseCatalogSql())) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1).trim())); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = """ + SELECT tabname, + CASE tabtype WHEN 'T' THEN 'TABLE' WHEN 'V' THEN 'VIEW' ELSE tabtype END + FROM systables + WHERE tabid >= 100 + ORDER BY tabname + """; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql.stripIndent().trim())) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1).trim(), rs.getString(2).trim(), null)); + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery( + "SELECT procname FROM sysprocedures WHERE owner != 'informix' AND isproc = 'f' ORDER BY procname" + )) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1).trim(), "FUNCTION", schema, null)); + } + } + + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery( + "SELECT procname FROM sysprocedures WHERE owner != 'informix' AND isproc = 't' ORDER BY procname" + )) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1).trim(), "PROCEDURE", schema, null)); + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String sql = """ + SELECT b.data FROM sysprocbody b + JOIN sysprocedures p ON b.procid = p.procid + WHERE p.procname = ? AND b.datakey = 'T' + ORDER BY b.seqno + """; + StringBuilder sb = new StringBuilder(); + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.stripIndent().trim())) { + stmt.setString(1, name); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String value = rs.getString(1); + sb.append(value == null ? "" : value); + } + } + } + return new ObjectSource(name, objectType, schema, sb.toString()); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Connection conn = requireConnected(); + Set primaryKeyColumns = getPrimaryKeyColumnNumbers(conn, table); + List result = new ArrayList<>(); + String sql = """ + SELECT c.colname, c.coltype, c.colno + FROM syscolumns c + WHERE c.tabid = (SELECT tabid FROM systables WHERE tabname = ?) + ORDER BY c.colno + """; + try (java.sql.PreparedStatement stmt = conn.prepareStatement(sql.stripIndent().trim())) { + stmt.setString(1, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String colname = rs.getString(1).trim(); + int coltype = rs.getInt(2); + result.add(new ColumnInfo( + colname, + mapColType(coltype), + (coltype & 256) == 0, + null, + primaryKeyColumns.contains(rs.getInt(3)), + null, + null, + null, + null, + null + )); + } + } + } + return result; + }); + } + + private Set getPrimaryKeyColumnNumbers(Connection conn, String table) throws SQLException { + String sql = """ + SELECT i.part1, i.part2, i.part3, i.part4, i.part5, i.part6, i.part7, i.part8, + i.part9, i.part10, i.part11, i.part12, i.part13, i.part14, i.part15, i.part16 + FROM sysconstraints c + JOIN sysindexes i ON i.idxname = c.idxname AND i.tabid = c.tabid + JOIN systables t ON t.tabid = c.tabid + WHERE t.tabname = ? AND c.constrtype = 'P' + """; + + try (java.sql.PreparedStatement stmt = conn.prepareStatement(sql.stripIndent().trim())) { + stmt.setString(1, table); + try (ResultSet rs = stmt.executeQuery()) { + if (!rs.next()) { + return Collections.emptySet(); + } + List parts = new ArrayList<>(); + for (int index = 1; index <= 16; index++) { + int value = rs.getInt(index); + parts.add(rs.wasNull() ? null : value); + } + return primaryKeyColumnNumbers(parts); + } + } + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = """ + SELECT t.trigname, t.event, 'TRIGGER' + FROM systriggers t + JOIN systables s ON t.tabid = s.tabid + WHERE s.tabname = ? + """; + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.stripIndent().trim())) { + stmt.setString(1, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo(rs.getString(1).trim(), rs.getString(2).trim(), rs.getString(3).trim())); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + String normalizedSql = switch (trimEnd(sql.trim(), ';').toUpperCase(Locale.ROOT)) { + case "BEGIN WORK" -> "BEGIN"; + case "COMMIT WORK" -> "COMMIT"; + case "ROLLBACK WORK" -> "ROLLBACK"; + default -> sql; + }; + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + normalizedSql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::stringResultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return ""; + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Object stringResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + private static String trimStart(String value, char... chars) { + int index = 0; + while (index < value.length() && contains(chars, value.charAt(index))) { + index++; + } + return value.substring(index); + } + + private static String trimEnd(String value, char... chars) { + int end = value.length(); + while (end > 0 && contains(chars, value.charAt(end - 1))) { + end--; + } + return value.substring(0, end); + } + + private static boolean contains(char[] chars, char value) { + for (char candidate : chars) { + if (candidate == value) { + return true; + } + } + return false; + } + + private static boolean containsIgnoreCase(String value, String needle) { + return value.toLowerCase(Locale.ROOT).contains(needle.toLowerCase(Locale.ROOT)); + } + + public static void main(String[] args) { + new JsonRpcServer(new InformixAgent()).run(); + } +} diff --git a/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentExecutionTest.java b/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentExecutionTest.java new file mode 100644 index 0000000000..32364b50c7 --- /dev/null +++ b/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentExecutionTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.informix; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class InformixAgentExecutionTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new InformixAgent(); + } + + @Override + protected String resultSetSql() { + return "EXECUTE PROCEDURE sample_proc()"; + } +} diff --git a/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java b/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java new file mode 100644 index 0000000000..795cc8246d --- /dev/null +++ b/agents/drivers/informix/src/test/java/com/dbx/agent/informix/InformixAgentTest.java @@ -0,0 +1,128 @@ +package com.dbx.agent.informix; + +import com.dbx.agent.ConnectParams; +import java.util.Arrays; +import java.util.Set; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class InformixAgentTest { + @Test + void buildsJdbcUrlWithExplicitInformixServerAndLocaleParameters() { + String url = InformixAgent.buildJdbcUrl( + new ConnectParams( + "172.26.128.159", + 20013, + "testdb", + "", + "", + "INFORMIXSERVER=informix;CLIENT_LOCALE=en_US.utf8;DB_LOCALE=en_US.utf8", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:informix-sqli://172.26.128.159:20013/testdb:INFORMIXSERVER=informix;CLIENT_LOCALE=en_US.utf8;DB_LOCALE=en_US.utf8", + url + ); + } + + @Test + void fallsBackToHostAsInformixServerWhenNoExplicitServerIsConfigured() { + String url = InformixAgent.buildJdbcUrl( + new ConnectParams( + "informix-host", + 9088, + "sysmaster", + "", + "", + "", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:informix-sqli://informix-host:9088/sysmaster:INFORMIXSERVER=informix-host;CLIENT_LOCALE=en_US.utf8;DB_LOCALE=en_US.utf8", + url + ); + } + + @Test + void fallsBackToInformixServerNameWhenHostIsAnIpAddress() { + String url = InformixAgent.buildJdbcUrl( + new ConnectParams( + "172.26.128.159", + 20013, + "sysmaster", + "", + "", + "", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:informix-sqli://172.26.128.159:20013/sysmaster:INFORMIXSERVER=informix;CLIENT_LOCALE=en_US.utf8;DB_LOCALE=en_US.utf8", + url + ); + } + + @Test + void usesInformixServerFromDedicatedFieldWhenProvided() { + ConnectParams params = new ConnectParams( + "172.26.128.159", + 20013, + "testdb", + "", + "", + "CLIENT_LOCALE=en_US.utf8", + "", + false + ); + params.setInformix_server("ol_informix1410"); + + String url = InformixAgent.buildJdbcUrl(params); + + Assertions.assertEquals( + "jdbc:informix-sqli://172.26.128.159:20013/testdb:INFORMIXSERVER=ol_informix1410;CLIENT_LOCALE=en_US.utf8;DB_LOCALE=en_US.utf8", + url + ); + } + + @Test + void usesSysmasterWhenDatabaseIsBlank() { + String url = InformixAgent.buildJdbcUrl( + new ConnectParams( + "informix-host", + 9088, + "", + "", + "", + "INFORMIXSERVER=informix", + "", + false + ) + ); + + Assertions.assertEquals( + "jdbc:informix-sqli://informix-host:9088/sysmaster:INFORMIXSERVER=informix;CLIENT_LOCALE=en_US.utf8;DB_LOCALE=en_US.utf8", + url + ); + } + + @Test + void extractsPrimaryKeyColumnNumbersFromInformixIndexParts() { + Assertions.assertEquals( + Set.of(1, 3, 5), + InformixAgent.primaryKeyColumnNumbers(Arrays.asList(1, -3, 0, 5, null)) + ); + } + + @Test + void listsDatabasesFromSysmasterCatalog() { + Assertions.assertEquals("SELECT name FROM sysmaster:sysdatabases ORDER BY name", InformixAgent.databaseCatalogSql()); + } +} diff --git a/agents/drivers/iotdb/build.gradle b/agents/drivers/iotdb/build.gradle new file mode 100644 index 0000000000..adaf4bf828 --- /dev/null +++ b/agents/drivers/iotdb/build.gradle @@ -0,0 +1,10 @@ +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/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java b/agents/drivers/iotdb/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java new file mode 100644 index 0000000000..93dff6b599 --- /dev/null +++ b/agents/drivers/iotdb/src/main/java/com/dbx/agent/iotdb/IoTDBAgent.java @@ -0,0 +1,22 @@ +package com.dbx.agent.iotdb; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +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); + } + + public static void main(String[] args) { + new JsonRpcServer(new IoTDBAgent()).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 new file mode 100644 index 0000000000..acfc62f0d7 --- /dev/null +++ b/agents/drivers/iotdb/src/test/java/com/dbx/agent/iotdb/IoTDBAgentTest.java @@ -0,0 +1,32 @@ +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/drivers/iris/build.gradle b/agents/drivers/iris/build.gradle new file mode 100644 index 0000000000..0bcd80d99e --- /dev/null +++ b/agents/drivers/iris/build.gradle @@ -0,0 +1,9 @@ +dependencies { + implementation 'com.intersystems:intersystems-jdbc:3.10.5' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'IRIS', 'Main-Class': 'com.dbx.agent.iris.IrisAgent') + } +} diff --git a/agents/drivers/iris/src/main/java/com/dbx/agent/iris/IrisAgent.java b/agents/drivers/iris/src/main/java/com/dbx/agent/iris/IrisAgent.java new file mode 100644 index 0000000000..380ea6231b --- /dev/null +++ b/agents/drivers/iris/src/main/java/com/dbx/agent/iris/IrisAgent.java @@ -0,0 +1,159 @@ +package com.dbx.agent.iris; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.StandardJdbcMetadata; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class IrisAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile IRIS_PROFILE = new JdbcAgentProfile( + "com.intersystems.jdbc.IRISDriver", + "jdbc:IRIS://{host}:{port}/{database}", + 1972, + true + ); + + public IrisAgent() { + super(IRIS_PROFILE); + } + + @Override + public List listSchemas() { + return dedupeCaseInsensitiveSchemas(StandardJdbcMetadata.INSTANCE.listSchemas(requireConnection(), IRIS_PROFILE)); + } + + @Override + public List getColumns(String schema, String table) { + return irisColumns(requireConnection(), schema, table); + } + + static List irisColumns(Connection conn, String schema, String table) { + return unchecked(() -> { + Set primaryKeys = irisPrimaryKeys(conn, schema, table); + List columns = new ArrayList<>(); + try (PreparedStatement stmt = conn.prepareStatement(columnSql(schema))) { + int index = 1; + if (!isBlank(schema)) { + stmt.setString(index++, schema.trim()); + } + stmt.setString(index, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + if (isBlank(name)) { + continue; + } + columns.add(new ColumnInfo( + name, + stringOrDefault(rs, "DATA_TYPE", "Unknown"), + !"NO".equalsIgnoreCase(stringOrNull(rs, "IS_NULLABLE")), + stringOrNull(rs, "COLUMN_DEFAULT"), + primaryKeys.contains(name.toUpperCase(Locale.ROOT)), + null, + null, + intOrNull(rs, "NUMERIC_PRECISION"), + intOrNull(rs, "NUMERIC_SCALE"), + intOrNull(rs, "CHARACTER_MAXIMUM_LENGTH") + )); + } + } + } + return columns; + }); + } + + static List dedupeCaseInsensitiveSchemas(List schemas) { + Map byNormalizedName = new LinkedHashMap<>(); + for (String schema : schemas) { + if (schema == null || schema.trim().isEmpty()) { + continue; + } + String name = schema.trim(); + byNormalizedName.putIfAbsent(name.toUpperCase(Locale.ROOT), name); + } + return new ArrayList<>(byNormalizedName.values()); + } + + private static Set irisPrimaryKeys(Connection conn, String schema, String table) { + Set keys = new LinkedHashSet<>(); + try (PreparedStatement stmt = conn.prepareStatement(primaryKeySql(schema))) { + int index = 1; + if (!isBlank(schema)) { + stmt.setString(index++, schema.trim()); + } + stmt.setString(index, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String column = rs.getString("COLUMN_NAME"); + if (!isBlank(column)) { + keys.add(column.toUpperCase(Locale.ROOT)); + } + } + } + } catch (Exception ignored) { + // Primary key metadata is optional for display; column listing must keep working. + } + return keys; + } + + private static String columnSql(String schema) { + StringBuilder sql = new StringBuilder( + "SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT, " + + "NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH " + + "FROM INFORMATION_SCHEMA.COLUMNS WHERE " + ); + if (!isBlank(schema)) { + sql.append("UPPER(TABLE_SCHEMA) = UPPER(?) AND "); + } + sql.append("UPPER(TABLE_NAME) = UPPER(?) ORDER BY ORDINAL_POSITION"); + return sql.toString(); + } + + private static String primaryKeySql(String schema) { + StringBuilder sql = new StringBuilder( + "SELECT k.COLUMN_NAME " + + "FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE k " + + "WHERE k.CONSTRAINT_TYPE = 'PRIMARY KEY' AND " + ); + if (!isBlank(schema)) { + sql.append("UPPER(k.TABLE_SCHEMA) = UPPER(?) AND "); + } + sql.append("UPPER(k.TABLE_NAME) = UPPER(?) ORDER BY k.ORDINAL_POSITION"); + return sql.toString(); + } + + private static String stringOrDefault(ResultSet rs, String column, String fallback) throws Exception { + String value = stringOrNull(rs, column); + return isBlank(value) ? fallback : value; + } + + private static String stringOrNull(ResultSet rs, String column) throws Exception { + String value = rs.getString(column); + return rs.wasNull() ? null : value; + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static boolean isBlank(String value) { + return value == null || value.trim().isEmpty(); + } + + public static void main(String[] args) { + new JsonRpcServer(new IrisAgent()).run(); + } +} diff --git a/agents/drivers/iris/src/test/java/com/dbx/agent/iris/IrisAgentTest.java b/agents/drivers/iris/src/test/java/com/dbx/agent/iris/IrisAgentTest.java new file mode 100644 index 0000000000..d2aa31965d --- /dev/null +++ b/agents/drivers/iris/src/test/java/com/dbx/agent/iris/IrisAgentTest.java @@ -0,0 +1,154 @@ +package com.dbx.agent.iris; + +import com.dbx.agent.ColumnInfo; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class IrisAgentTest { + @Test + void skipsSchemaSwitchingBecauseIrisRejectsSetSchemaContext() { + IrisAgent agent = new IrisAgent(); + + assertEquals("", agent.setSchemaSQL("Ens")); + } + + @Test + void dedupesSchemasCaseInsensitively() { + assertEquals( + Arrays.asList("APP", "SQLUSER", "z_user"), + IrisAgent.dedupeCaseInsensitiveSchemas(Arrays.asList("APP", "SQLUSER", "SQLUser", "app", "z_user")) + ); + } + + @Test + void ignoresBlankSchemasWhenDeduping() { + assertEquals( + Collections.singletonList("SQLUSER"), + IrisAgent.dedupeCaseInsensitiveSchemas(Arrays.asList("", " ", null, "SQLUSER")) + ); + } + + @Test + void readsColumnsWithNativeInformationSchemaQueries() { + List sql = new ArrayList<>(); + Connection conn = fakeConnection(sql); + + List columns = IrisAgent.irisColumns(conn, "SQLUser", "People"); + + assertEquals(2, columns.size()); + assertEquals("ID", columns.get(0).getName()); + assertEquals("INTEGER", columns.get(0).getData_type()); + assertFalse(columns.get(0).getIs_nullable()); + assertTrue(columns.get(0).getIs_primary_key()); + assertEquals("NAME", columns.get(1).getName()); + assertEquals(Integer.valueOf(64), columns.get(1).getCharacter_maximum_length()); + assertFalse(sql.get(0).contains("LIMIT")); + assertTrue(sql.get(0).contains("INFORMATION_SCHEMA.KEY_COLUMN_USAGE")); + assertTrue(sql.get(1).contains("INFORMATION_SCHEMA.COLUMNS")); + } + + private static Connection fakeConnection(List sqlLog) { + return proxy(Connection.class, (method, args) -> { + if ("prepareStatement".equals(method.getName())) { + String sql = (String) args[0]; + sqlLog.add(sql); + return fakeStatement(sql); + } + if ("getMetaData".equals(method.getName())) { + throw new AssertionError("IRIS columns must not use JDBC DatabaseMetaData.getColumns()"); + } + return defaultValue(method.getReturnType()); + }); + } + + private static PreparedStatement fakeStatement(String sql) { + return proxy(PreparedStatement.class, (method, args) -> { + if ("executeQuery".equals(method.getName())) { + if (sql.contains("KEY_COLUMN_USAGE")) { + return fakeResultSet(List.of(Map.of("COLUMN_NAME", "ID"))); + } + return fakeResultSet(List.of( + Map.of( + "COLUMN_NAME", "ID", + "DATA_TYPE", "INTEGER", + "IS_NULLABLE", "NO", + "NUMERIC_PRECISION", 10 + ), + Map.of( + "COLUMN_NAME", "NAME", + "DATA_TYPE", "VARCHAR", + "IS_NULLABLE", "YES", + "CHARACTER_MAXIMUM_LENGTH", 64 + ) + )); + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet fakeResultSet(List> rows) { + final int[] index = {-1}; + final boolean[] wasNull = {false}; + return proxy(ResultSet.class, (method, args) -> { + String name = method.getName(); + if ("next".equals(name)) { + index[0] += 1; + return index[0] < rows.size(); + } + if ("getString".equals(name)) { + Object value = rows.get(index[0]).get((String) args[0]); + wasNull[0] = value == null; + return value == null ? null : String.valueOf(value); + } + if ("getObject".equals(name)) { + Object value = rows.get(index[0]).get((String) args[0]); + wasNull[0] = value == null; + return value; + } + if ("wasNull".equals(name)) { + return wasNull[0]; + } + return defaultValue(method.getReturnType()); + }); + } + + @SuppressWarnings("unchecked") + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + return handler.handle(method, args == null ? new Object[0] : args); + } + }; + return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } +} diff --git a/agents/drivers/kingbase/build.gradle b/agents/drivers/kingbase/build.gradle new file mode 100644 index 0000000000..7c9eef176f --- /dev/null +++ b/agents/drivers/kingbase/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'cn.com.kingbase:kingbase8:9.0.1.jre7' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': '人大金仓 KingbaseES', 'Main-Class': 'com.dbx.agent.kingbase.KingbaseAgent') + } +} diff --git a/agents/drivers/kingbase/libs/.gitkeep b/agents/drivers/kingbase/libs/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java new file mode 100644 index 0000000000..0cd723aa80 --- /dev/null +++ b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java @@ -0,0 +1,349 @@ +package com.dbx.agent.kingbase; + +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.PostgresLikeAgent; +import com.dbx.agent.PostgresLikeAgentProfile; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class KingbaseAgent extends PostgresLikeAgent { + public static final PostgresLikeAgentProfile KINGBASE_PROFILE = new PostgresLikeAgentProfile( + "com.kingbase8.Driver", + "jdbc:kingbase8://{host}:{port}/{database}" + ); + + public KingbaseAgent() { + super(KINGBASE_PROFILE); + } + + @Override + protected void afterConnect(ConnectParams params, Connection connection) { + if (params.isMysql_compat_mode()) { + setMysqlCompatMode(true); + } + } + + @Override + public List listDatabases() { + return unchecked(() -> { + try (PreparedStatement stmt = requireConnected().prepareStatement("SELECT current_database() AS database_name"); + ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + return Collections.singletonList(new DatabaseInfo(rs.getString("database_name"))); + } + } + return Collections.singletonList(new DatabaseInfo(getConfiguredDatabase())); + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT schema_name " + + "FROM information_schema.schemata " + + "WHERE UPPER(schema_name) <> 'INFORMATION_SCHEMA' " + + "AND UPPER(schema_name) NOT LIKE 'SYS%' " + + "AND UPPER(schema_name) NOT LIKE 'XLOG%' " + + "ORDER BY schema_name"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql); + ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(rs.getString("schema_name")); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return listTables(schema, "table_type = 'BASE TABLE'"); + } + + @Override + public List listObjects(String schema) { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), effectiveSchema(schema), table.getComment())); + } + return result; + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + if (!"VIEW".equalsIgnoreCase(objectType)) { + return new ObjectSource(name, objectType, effectiveSchema(schema), ""); + } + return unchecked(() -> { + String source = ""; + String sql = "SELECT view_definition " + + "FROM information_schema.views " + + "WHERE table_schema = " + sqlString(effectiveSchema(schema)) + + " AND table_name = " + sqlString(name); + try (Statement stmt = requireConnected().createStatement()) { + try (ResultSet rs = stmt.executeQuery(sql)) { + if (rs.next()) { + source = coalesce(rs.getString("view_definition")); + } + } + } + return new ObjectSource(name, objectType, effectiveSchema(schema), source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = primaryKeys(schema, table); + List result = new ArrayList<>(); + String sql = "SELECT column_name, data_type, is_nullable, column_default, " + + "numeric_precision, numeric_scale, character_maximum_length " + + "FROM information_schema.columns " + + "WHERE table_schema = " + sqlString(effectiveSchema(schema)) + + " AND table_name = " + sqlString(table) + " " + + "ORDER BY ordinal_position"; + try (Statement stmt = requireConnected().createStatement()) { + try (ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + String columnName = rs.getString("column_name"); + result.add(new ColumnInfo( + columnName, + rs.getString("data_type"), + "YES".equalsIgnoreCase(coalesce(rs.getString("is_nullable"))), + rs.getString("column_default"), + primaryKeys.contains(columnName), + null, + null, + intObject(rs, "numeric_precision"), + intObject(rs, "numeric_scale"), + intObject(rs, "character_maximum_length") + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + Map indexes = new LinkedHashMap<>(); + String sql = "SELECT tc.constraint_name, tc.constraint_type, kcu.column_name " + + "FROM information_schema.table_constraints tc " + + "JOIN information_schema.key_column_usage kcu " + + "ON kcu.constraint_schema = tc.constraint_schema " + + "AND kcu.constraint_name = tc.constraint_name " + + "AND kcu.table_schema = tc.table_schema " + + "AND kcu.table_name = tc.table_name " + + "WHERE tc.table_schema = " + sqlString(effectiveSchema(schema)) + + " AND tc.table_name = " + sqlString(table) + " " + + "AND tc.constraint_type IN ('PRIMARY KEY', 'UNIQUE') " + + "ORDER BY tc.constraint_name, kcu.ordinal_position"; + try (Statement stmt = requireConnected().createStatement()) { + try (ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + String name = rs.getString("constraint_name"); + String type = rs.getString("constraint_type"); + ConstraintIndexBuilder builder = indexes.get(name); + if (builder == null) { + builder = new ConstraintIndexBuilder(name, "PRIMARY KEY".equalsIgnoreCase(type)); + indexes.put(name, builder); + } + builder.columns.add(rs.getString("column_name")); + } + } + } + List result = new ArrayList<>(); + for (ConstraintIndexBuilder index : indexes.values()) { + result.add(new IndexInfo(index.name, index.columns, true, index.primary, null, index.primary ? "PRIMARY KEY" : "UNIQUE", null, null)); + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT fk.constraint_name, fk.column_name, pk.table_name AS ref_table, pk.column_name AS ref_column " + + "FROM information_schema.table_constraints tc " + + "JOIN information_schema.key_column_usage fk " + + "ON fk.constraint_schema = tc.constraint_schema " + + "AND fk.constraint_name = tc.constraint_name " + + "AND fk.table_schema = tc.table_schema " + + "AND fk.table_name = tc.table_name " + + "JOIN information_schema.referential_constraints rc " + + "ON rc.constraint_schema = tc.constraint_schema " + + "AND rc.constraint_name = tc.constraint_name " + + "JOIN information_schema.key_column_usage pk " + + "ON pk.constraint_schema = rc.unique_constraint_schema " + + "AND pk.constraint_name = rc.unique_constraint_name " + + "AND pk.ordinal_position = fk.position_in_unique_constraint " + + "WHERE tc.table_schema = " + sqlString(effectiveSchema(schema)) + + " AND tc.table_name = " + sqlString(table) + " " + + "AND tc.constraint_type = 'FOREIGN KEY' " + + "ORDER BY fk.constraint_name, fk.ordinal_position"; + try (Statement stmt = requireConnected().createStatement()) { + try (ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("constraint_name"), + rs.getString("column_name"), + rs.getString("ref_table"), + rs.getString("ref_column") + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return "SET search_path TO " + JdbcIdentifiers.INSTANCE.doubleQuote(effectiveSchema(schema)); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType, String columnTypeName) { + if (isTemporalType(sqlType, columnTypeName)) { + return unchecked(() -> { + Object value = rs.getTimestamp(index); + return rs.wasNull() ? null : value.toString(); + }); + } + return super.resultValue(rs, index, sqlType, columnTypeName); + } + + private static boolean isTemporalType(int sqlType, String columnTypeName) { + switch (sqlType) { + case Types.DATE: + case Types.TIME: + case Types.TIME_WITH_TIMEZONE: + case Types.TIMESTAMP: + case Types.TIMESTAMP_WITH_TIMEZONE: + return true; + default: + break; + } + if (columnTypeName == null) { + return false; + } + String normalized = columnTypeName.trim().toLowerCase(Locale.ROOT); + return normalized.equals("date") + || normalized.equals("time") + || normalized.equals("datetime") + || normalized.startsWith("timestamp"); + } + + private Set primaryKeys(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = new LinkedHashSet<>(); + String sql = "SELECT kcu.column_name " + + "FROM information_schema.table_constraints tc " + + "JOIN information_schema.key_column_usage kcu " + + "ON kcu.constraint_schema = tc.constraint_schema " + + "AND kcu.constraint_name = tc.constraint_name " + + "AND kcu.table_schema = tc.table_schema " + + "AND kcu.table_name = tc.table_name " + + "WHERE tc.table_schema = " + sqlString(effectiveSchema(schema)) + + " AND tc.table_name = " + sqlString(table) + " " + + "AND tc.constraint_type = 'PRIMARY KEY' " + + "ORDER BY kcu.ordinal_position"; + try (Statement stmt = requireConnected().createStatement()) { + try (ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + primaryKeys.add(rs.getString("column_name")); + } + } + } + return primaryKeys; + }); + } + + private List listTables(String schema, String tableTypePredicate) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT table_name, table_type " + + "FROM information_schema.tables " + + "WHERE table_schema = " + sqlString(effectiveSchema(schema)) + " AND " + tableTypePredicate + " " + + "ORDER BY table_name"; + try (Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), normalizeTableType(rs.getString(2)))); + } + } + return result; + }); + } + + private String effectiveSchema(String schema) { + if (schema != null && !schema.trim().isEmpty()) { + return schema; + } + return "PUBLIC"; + } + + private static Integer intObject(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static String normalizeTableType(String type) { + if (type == null || type.trim().isEmpty()) return "TABLE"; + if ("BASE TABLE".equalsIgnoreCase(type)) return "TABLE"; + return type; + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + private static String sqlString(String value) { + return "'" + coalesce(value).replace("'", "''") + "'"; + } + + private static final class ConstraintIndexBuilder { + final String name; + final List columns = new ArrayList<>(); + final boolean primary; + + ConstraintIndexBuilder(String name, boolean primary) { + this.name = name; + this.primary = primary; + } + } + + public static void main(String[] args) { + new JsonRpcServer(new KingbaseAgent()).run(); + } +} diff --git a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java new file mode 100644 index 0000000000..6d2e542bd2 --- /dev/null +++ b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java @@ -0,0 +1,184 @@ +package com.dbx.agent.kingbase; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.TestSupport; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.sql.Timestamp; +import java.sql.Types; +import java.util.ArrayList; +import java.util.List; + +class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new KingbaseAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } + + @Test + void declaresKingbasePostgresLikeProfile() { + KingbaseAgent agent = new KingbaseAgent(); + + Assertions.assertEquals("com.kingbase8.Driver", agent.getProfile().getDriverClass()); + Assertions.assertEquals("jdbc:kingbase8://{host}:{port}/{database}", agent.getProfile().getUrlTemplate()); + } + + @Test + void mysqlCompatListDatabasesUsesCurrentDatabase() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + agent.setMysqlCompatMode(true); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"database_name"}, + new Object[][]{{"TEST"}} + ))); + + Assertions.assertEquals("TEST", agent.listDatabases().get(0).getName()); + Assertions.assertEquals("SELECT current_database() AS database_name", sql.get(0)); + } + + @Test + void mysqlCompatListTablesUsesInformationSchema() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + agent.setMysqlCompatMode(true); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"table_name", "table_type"}, + new Object[][]{{"test_timestamps", "BASE TABLE"}} + ))); + + Assertions.assertEquals("test_timestamps", agent.listTables("PUBLIC").get(0).getName()); + Assertions.assertTrue(sql.get(0).contains("FROM information_schema.tables")); + Assertions.assertFalse(sql.get(0).contains("SHOW")); + } + + @Test + void mysqlCompatTimestampTypeNameIsReadAsTimestampText() throws Exception { + Timestamp timestamp = Timestamp.valueOf("2026-06-22 11:29:00"); + KingbaseAgent agent = new KingbaseAgent(); + agent.setMysqlCompatMode(true); + + Object value = readResultValue(agent, timestampResultSet(timestamp), Types.BINARY, "timestamp"); + + Assertions.assertEquals("2026-06-22 11:29:00.0", value); + } + + private static Connection preparedConnection(List sql, ResultSet rs) { + PreparedStatement statement = proxy(PreparedStatement.class, (method, args) -> { + if ("executeQuery".equals(method.getName())) { + return rs; + } + if ("close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + Statement plainStatement = proxy(Statement.class, (method, args) -> { + if ("executeQuery".equals(method.getName())) { + sql.add(String.valueOf(args[0])); + return rs; + } + if ("close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + return proxy(Connection.class, (method, args) -> { + if ("prepareStatement".equals(method.getName())) { + sql.add(String.valueOf(args[0])); + return statement; + } + if ("createStatement".equals(method.getName())) { + return plainStatement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet resultSet(String[] columns, Object[][] rows) { + int[] index = {-1}; + return proxy(ResultSet.class, (method, args) -> { + switch (method.getName()) { + case "next": + index[0] += 1; + return index[0] < rows.length; + case "getString": + Object key = args[0]; + if (key instanceof Number) { + return rows[index[0]][((Number) key).intValue() - 1]; + } + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(String.valueOf(key))) { + return rows[index[0]][i]; + } + } + return null; + case "wasNull": + return false; + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet timestampResultSet(Timestamp timestamp) { + return proxy(ResultSet.class, (method, args) -> { + switch (method.getName()) { + case "getTimestamp": + return timestamp; + case "getBytes": + throw new AssertionError("timestamp should not be read as bytes"); + case "wasNull": + return false; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = (Object unused, Method method, Object[] args) -> handler.handle(method, args); + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) return false; + if (Byte.TYPE.equals(type)) return (byte) 0; + if (Short.TYPE.equals(type)) return (short) 0; + if (Integer.TYPE.equals(type)) return 0; + if (Long.TYPE.equals(type)) return 0L; + if (Float.TYPE.equals(type)) return 0f; + if (Double.TYPE.equals(type)) return 0.0d; + if (Character.TYPE.equals(type)) return '\0'; + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } + + private static Object readResultValue(KingbaseAgent agent, ResultSet rs, int sqlType, String columnTypeName) throws Exception { + Method method = KingbaseAgent.class.getDeclaredMethod("resultValue", ResultSet.class, int.class, int.class, String.class); + method.setAccessible(true); + return method.invoke(agent, rs, 1, sqlType, columnTypeName); + } +} diff --git a/agents/drivers/kylin/build.gradle b/agents/drivers/kylin/build.gradle new file mode 100644 index 0000000000..7ec3cd652b --- /dev/null +++ b/agents/drivers/kylin/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'org.apache.kylin:kylin-jdbc:5.0.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Apache Kylin', 'Main-Class': 'com.dbx.agent.kylin.KylinAgent') + } +} diff --git a/agents/drivers/kylin/src/main/java/com/dbx/agent/kylin/KylinAgent.java b/agents/drivers/kylin/src/main/java/com/dbx/agent/kylin/KylinAgent.java new file mode 100644 index 0000000000..8af270cdb7 --- /dev/null +++ b/agents/drivers/kylin/src/main/java/com/dbx/agent/kylin/KylinAgent.java @@ -0,0 +1,144 @@ +package com.dbx.agent.kylin; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +public final class KylinAgent extends AbstractJdbcAgent { + + @Override + protected String driverClass() { + return "org.apache.kylin.jdbc.Driver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (ResultSet rs = requireConnected().getMetaData().getCatalogs()) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString("TABLE_CAT"))); + } + } + result.sort(Comparator.comparing(DatabaseInfo::getName)); + return result; + }); + } + + @Override + public List listSchemas() { + return Collections.emptyList(); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (ResultSet rs = requireConnected().getMetaData().getTables(null, schema, null, null)) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("TABLE_NAME"), + normalizeTableType(rs.getString("TABLE_TYPE")), + null + )); + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = new LinkedHashSet<>(); + try (ResultSet rs = requireConnected().getMetaData().getPrimaryKeys(null, schema, table)) { + while (rs.next()) { + primaryKeys.add(rs.getString("COLUMN_NAME")); + } + } + + List result = new ArrayList<>(); + try (ResultSet rs = requireConnected().getMetaData().getColumns(null, schema, table, null)) { + while (rs.next()) { + String colName = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + colName, + rs.getString("TYPE_NAME"), + "YES".equals(rs.getString("IS_NULLABLE")), + rs.getString("COLUMN_DEF"), + primaryKeys.contains(colName), + null, + emptyToNull(rs.getString("REMARKS")), + intOrNull(rs, "COLUMN_SIZE"), + intOrNull(rs, "DECIMAL_DIGITS"), + null + )); + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return ""; + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:kylin://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private static String normalizeTableType(String type) { + if ("BASE TABLE".equals(type)) { + return "TABLE"; + } + return type; + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static String emptyToNull(String value) { + return value == null || value.isEmpty() ? null : value; + } + + public static void main(String[] args) { + new JsonRpcServer(new KylinAgent()).run(); + } +} diff --git a/agents/drivers/kylin/src/test/java/com/dbx/agent/kylin/KylinAgentTest.java b/agents/drivers/kylin/src/test/java/com/dbx/agent/kylin/KylinAgentTest.java new file mode 100644 index 0000000000..4c3e88be9b --- /dev/null +++ b/agents/drivers/kylin/src/test/java/com/dbx/agent/kylin/KylinAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.kylin; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class KylinAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new KylinAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } +} diff --git a/agents/drivers/mongodb/build.gradle b/agents/drivers/mongodb/build.gradle new file mode 100644 index 0000000000..125b760f20 --- /dev/null +++ b/agents/drivers/mongodb/build.gradle @@ -0,0 +1,11 @@ +dependencies { + implementation project(':common') + implementation 'org.mongodb:mongodb-driver-sync:3.12.14' + implementation 'com.google.code.gson:gson:2.11.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'MongoDB (Legacy)', 'Main-Class': 'com.dbx.agent.mongodb.MongoAgent') + } +} diff --git a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java new file mode 100644 index 0000000000..93f2fae868 --- /dev/null +++ b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java @@ -0,0 +1,586 @@ +package com.dbx.agent.mongodb; + +import com.dbx.agent.AgentProtocol; +import com.google.gson.Gson; +import com.google.gson.JsonElement; +import com.google.gson.JsonNull; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import com.mongodb.ConnectionString; +import com.mongodb.MongoCredential; +import com.mongodb.MongoClientSettings; +import com.mongodb.ServerAddress; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import java.io.BufferedReader; +import java.io.FileInputStream; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.net.URLDecoder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.security.KeyFactory; +import java.security.KeyStore; +import java.security.PrivateKey; +import java.security.SecureRandom; +import java.security.cert.Certificate; +import java.security.cert.CertificateFactory; +import java.security.spec.PKCS8EncodedKeySpec; +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collection; +import java.util.Collections; +import java.util.Date; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import javax.net.ssl.KeyManager; +import javax.net.ssl.KeyManagerFactory; +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManager; +import javax.net.ssl.TrustManagerFactory; +import org.bson.Document; +import org.bson.types.ObjectId; + +public final class MongoAgent { + private static final Gson GSON = new Gson(); + private static final DateTimeFormatter DATE_FORMAT = DateTimeFormatter.ISO_INSTANT.withZone(ZoneOffset.UTC); + private static final long JS_MAX_SAFE_INTEGER = 9_007_199_254_740_991L; + private static MongoClient client; + + private MongoAgent() { + } + + private static String stringOrNull(JsonObject object, String key) { + JsonElement element = object.get(key); + return element == null || element instanceof JsonNull ? null : element.getAsString(); + } + + static MongoClientSettings.Builder configureBuilder(JsonObject connObj) { + String host = connObj.has("host") ? connObj.get("host").getAsString() : "127.0.0.1"; + int port = connObj.has("port") ? connObj.get("port").getAsInt() : 27017; + String username = coalesce(stringOrNull(connObj, "username")); + String password = coalesce(stringOrNull(connObj, "password")); + String authDatabase = authenticationDatabase(connObj); + String connectionString = stringOrNull(connObj, "connection_string"); + boolean ssl = connObj.has("ssl") && !connObj.get("ssl").isJsonNull() && connObj.get("ssl").getAsBoolean(); + String caCertPath = stringOrNull(connObj, "ca_cert_path"); + String clientCertPath = firstNonBlank( + stringOrNull(connObj, "client_cert_path"), stringOrNull(connObj, "cert_path")); + String clientKeyPath = firstNonBlank( + stringOrNull(connObj, "client_key_path"), stringOrNull(connObj, "key_path")); + + if ((clientCertPath == null) != (clientKeyPath == null)) { + throw new IllegalArgumentException("Client certificate and key must be provided together"); + } + + MongoClientSettings.Builder builder = MongoClientSettings.builder(); + if (connectionString != null && !connectionString.isBlank()) { + builder.applyConnectionString(new ConnectionString(connectionString)); + } else { + builder.applyToClusterSettings( + settings -> settings.hosts(Collections.singletonList(new ServerAddress(host, port)))); + if (!username.isBlank()) { + builder.credential(MongoCredential.createCredential(username, authDatabase, password.toCharArray())); + } + } + + if (ssl) { + applyTlsSettings(builder, caCertPath, clientCertPath, clientKeyPath); + } + + return builder; + } + + private static Object connect(JsonObject params) { + JsonObject connObj = params.has("connection") && params.get("connection").isJsonObject() + ? params.getAsJsonObject("connection") + : params; + String database = defaultString(stringOrNull(connObj, "database"), "admin"); + + MongoClientSettings.Builder builder = configureBuilder(connObj); + + if (client != null) { + client.close(); + } + client = MongoClients.create(builder.build()); + client.getDatabase(database).runCommand(new Document("ping", 1)); + return Collections.singletonMap("ok", true); + } + + private static void applyTlsSettings(MongoClientSettings.Builder builder, + String caCertPath, String clientCertPath, String clientKeyPath) { + builder.applyToSslSettings(sslBuilder -> { + sslBuilder.enabled(true); + if (caCertPath != null && !caCertPath.isBlank() + || clientCertPath != null && !clientCertPath.isBlank()) { + try { + sslBuilder.context(createTlsSslContext(caCertPath, clientCertPath, clientKeyPath)); + } catch (Exception e) { + throw new RuntimeException("Failed to configure TLS: " + e.getMessage(), e); + } + } + }); + } + + static SSLContext createTlsSslContext(String caCertPath, String clientCertPath, String clientKeyPath) + throws Exception { + TrustManager[] trustManagers = null; + if (caCertPath != null && !caCertPath.isBlank()) { + trustManagers = loadTrustManagersFromPem(caCertPath); + } + + KeyManager[] keyManagers = null; + if (clientCertPath != null && !clientCertPath.isBlank() + && clientKeyPath != null && !clientKeyPath.isBlank()) { + keyManagers = loadKeyManagersFromPem(clientCertPath, clientKeyPath); + } + + SSLContext ctx = SSLContext.getInstance("TLS"); + ctx.init(keyManagers, trustManagers, new SecureRandom()); + return ctx; + } + + static TrustManager[] loadTrustManagersFromPem(String caCertPath) throws Exception { + CertificateFactory cf = CertificateFactory.getInstance("X.509"); + KeyStore trustStore = KeyStore.getInstance(KeyStore.getDefaultType()); + trustStore.load(null, null); + int i = 0; + try (InputStream is = new FileInputStream(caCertPath)) { + for (Certificate cert : (Collection) cf.generateCertificates(is)) { + trustStore.setCertificateEntry("ca-" + i, cert); + i++; + } + } + TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm()); + tmf.init(trustStore); + return tmf.getTrustManagers(); + } + + static KeyManager[] loadKeyManagersFromPem(String certPath, String keyPath) throws Exception { + CertificateFactory cf = CertificateFactory.getInstance("X.509"); + Certificate cert; + try (InputStream is = new FileInputStream(certPath)) { + cert = cf.generateCertificate(is); + } + + PrivateKey key = loadPrivateKeyFromPem(keyPath); + + KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType()); + keyStore.load(null, null); + keyStore.setCertificateEntry("client", cert); + keyStore.setKeyEntry("client", key, new char[0], new Certificate[] {cert}); + + KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); + kmf.init(keyStore, new char[0]); + return kmf.getKeyManagers(); + } + + static PrivateKey loadPrivateKeyFromPem(String keyPath) throws Exception { + String content = new String(Files.readAllBytes(Paths.get(keyPath)), StandardCharsets.UTF_8); + content = content.replace("-----BEGIN PRIVATE KEY-----", "") + .replace("-----END PRIVATE KEY-----", "") + .replace("-----BEGIN RSA PRIVATE KEY-----", "") + .replace("-----END RSA PRIVATE KEY-----", "") + .replace("-----BEGIN EC PRIVATE KEY-----", "") + .replace("-----END EC PRIVATE KEY-----", ""); + byte[] keyBytes = Base64.getDecoder().decode(content.replaceAll("\\s", "")); + + // PKCS#8 (standard format, "-----BEGIN PRIVATE KEY-----") + try { + return KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(keyBytes)); + } catch (Exception e) { + // ignore — try next format + } + try { + return KeyFactory.getInstance("EC").generatePrivate(new PKCS8EncodedKeySpec(keyBytes)); + } catch (Exception e) { + // ignore — try next format + } + + // PKCS#1 RSA — add PKCS#8 AlgorithmIdentifier prefix + // The prefix is: SEQUENCE { INTEGER 0, SEQUENCE { OID 1.2.840.113549.1.1.1, NULL }, OCTET STRING } + try { + byte[] pkcs8Header = { + 0x30, (byte) 0x82, 0, 0, // SEQUENCE (length filled in below) + 0x02, 0x01, 0x00, // INTEGER 0 + 0x30, 0x0d, // SEQUENCE (AlgorithmIdentifier) + 0x06, 0x09, 0x2a, (byte) 0x86, 0x48, (byte) 0x86, (byte) 0xf7, 0x0d, 0x01, 0x01, 0x01, // OID 1.2.840.113549.1.1.1 + 0x05, 0x00, // NULL + 0x04 // OCTET STRING (length filled in below) + }; + int totalLen = pkcs8Header.length + keyBytes.length - 4; // subtract placeholder SEQUENCE length + pkcs8Header[2] = (byte) ((totalLen >> 8) & 0xff); + pkcs8Header[3] = (byte) (totalLen & 0xff); + // OCTET STRING length + int octetLen = keyBytes.length; + byte[] octetLenBytes; + if (octetLen < 128) { + octetLenBytes = new byte[] {(byte) octetLen}; + } else if (octetLen < 256) { + octetLenBytes = new byte[] {(byte) 0x81, (byte) octetLen}; + } else { + octetLenBytes = new byte[] {(byte) 0x82, (byte) (octetLen >> 8), (byte) (octetLen & 0xff)}; + } + byte[] pkcs8Key = new byte[pkcs8Header.length + octetLenBytes.length - 1 + keyBytes.length]; + int pos = 0; + System.arraycopy(pkcs8Header, 0, pkcs8Key, pos, pkcs8Header.length - 1); // exclude placeholder OCTET STRING length + pos += pkcs8Header.length - 1; + System.arraycopy(octetLenBytes, 0, pkcs8Key, pos, octetLenBytes.length); + pos += octetLenBytes.length; + System.arraycopy(keyBytes, 0, pkcs8Key, pos, keyBytes.length); + return KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(pkcs8Key)); + } catch (Exception e) { + throw new IllegalArgumentException( + "Unsupported private key format in " + keyPath + + ". Use PKCS#8 (-----BEGIN PRIVATE KEY-----) or PKCS#1 RSA (-----BEGIN RSA PRIVATE KEY-----).", + e); + } + } + + static String firstNonBlank(String... values) { + for (String value : values) { + if (value != null && !value.isBlank()) { + return value; + } + } + return null; + } + + static String authenticationDatabase(JsonObject connObj) { + String authSource = urlParam(stringOrNull(connObj, "url_params"), "authSource"); + if (authSource != null && !authSource.isBlank()) { + return authSource; + } + return defaultString(stringOrNull(connObj, "database"), "admin"); + } + + private static String urlParam(String urlParams, String key) { + if (urlParams == null || urlParams.isBlank()) { + return null; + } + String normalized = urlParams.startsWith("?") ? urlParams.substring(1) : urlParams; + for (String pair : normalized.split("&")) { + if (pair.isBlank()) continue; + String[] parts = pair.split("=", 2); + if (decode(parts[0]).equals(key)) { + return parts.length > 1 ? decode(parts[1]) : ""; + } + } + return null; + } + + private static String decode(String value) { + return URLDecoder.decode(value, StandardCharsets.UTF_8); + } + + private static Object listDatabases() { + MongoClient c = requireClient(); + List> result = new ArrayList<>(); + for (String name : c.listDatabaseNames()) { + result.add(Collections.singletonMap("name", name)); + } + return result; + } + + private static Object listCollections(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + List result = new ArrayList<>(); + for (String name : c.getDatabase(database).listCollectionNames()) { + result.add(name); + } + return result; + } + + private static Object findDocuments(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + long skip = params.has("skip") ? params.get("skip").getAsLong() : 0; + int limit = params.has("limit") ? params.get("limit").getAsInt() : 50; + String filterJson = stringOrNull(params, "filter"); + String sortJson = stringOrNull(params, "sort"); + + var col = c.getDatabase(database).getCollection(collection); + Document filterDoc = filterJson != null && !filterJson.isBlank() ? Document.parse(filterJson) : new Document(); + long total = col.countDocuments(filterDoc); + + var iterable = col.find(filterDoc).skip((int) skip).limit(limit); + if (sortJson != null && !sortJson.isBlank()) { + iterable = iterable.sort(Document.parse(sortJson)); + } + + List> documents = new ArrayList<>(); + for (Document document : iterable) { + documents.add(bsonToJson(document)); + } + Map result = new LinkedHashMap<>(); + result.put("documents", documents); + result.put("total", total); + return result; + } + + private static Object insertDocument(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + String docJson = params.get("doc_json").getAsString(); + + Document doc = Document.parse(docJson); + c.getDatabase(database).getCollection(collection).insertOne(doc); + Object insertedId = convertValue(doc.get("_id")); + return Collections.singletonMap("inserted_id", insertedId); + } + + private static Object parseId(String id) { + try { + return new ObjectId(id); + } catch (Exception e) { + return id; + } + } + + private static Object updateDocument(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + String id = params.get("id").getAsString(); + String docJson = params.get("doc_json").getAsString(); + + var col = c.getDatabase(database).getCollection(collection); + Document newDoc = documentForWrite(docJson); + var filter = new Document("_id", parseId(id)); + var result = isUpdateOperatorDocument(newDoc) + ? col.updateOne(filter, newDoc) + : col.replaceOne(filter, replacementDocument(newDoc)); + return Collections.singletonMap("modified_count", result.getModifiedCount()); + } + + static Document documentForWrite(String docJson) { + Document doc = Document.parse(docJson); + convertMongoShellDates(doc); + return doc; + } + + private static Document replacementDocument(Document doc) { + Document replacement = new Document(doc); + replacement.remove("_id"); + return replacement; + } + + static boolean isUpdateOperatorDocument(Document doc) { + if (doc.isEmpty()) { + return false; + } + for (String key : doc.keySet()) { + if (!key.startsWith("$")) { + return false; + } + } + return true; + } + + private static Object deleteDocument(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + String id = params.get("id").getAsString(); + + var col = c.getDatabase(database).getCollection(collection); + var result = col.deleteOne(new Document("_id", parseId(id))); + return Collections.singletonMap("deleted_count", result.getDeletedCount()); + } + + private static Map bsonToJson(Document doc) { + Map result = new LinkedHashMap<>(); + for (Map.Entry entry : doc.entrySet()) { + result.put(entry.getKey(), convertValue(entry.getValue())); + } + return result; + } + + static Object convertValue(Object value) { + if (value == null) { + return null; + } + if (value instanceof ObjectId objectId) { + return objectId.toHexString(); + } + if (value instanceof Document document) { + return bsonToJson(document); + } + if (value instanceof List values) { + List result = new ArrayList<>(); + for (Object item : values) { + result.add(convertValue(item)); + } + return result; + } + if (value instanceof java.util.Date date) { + Instant instant = Instant.ofEpochMilli(date.getTime()); + return "ISODate(\"" + DATE_FORMAT.format(instant) + "\")"; + } + if (value instanceof Long longValue) { + return longValue < -JS_MAX_SAFE_INTEGER || longValue > JS_MAX_SAFE_INTEGER ? longValue.toString() : longValue; + } + if (value instanceof Boolean || value instanceof Integer || value instanceof Double || value instanceof String) { + return value; + } + return value.toString(); + } + + @SuppressWarnings("unchecked") + private static Object convertMongoShellDates(Object value) { + if (value instanceof Document document) { + for (String key : new ArrayList<>(document.keySet())) { + document.put(key, convertMongoShellDates(document.get(key))); + } + return document; + } + if (value instanceof List values) { + List converted = (List) values; + for (int i = 0; i < converted.size(); i++) { + converted.set(i, convertMongoShellDates(converted.get(i))); + } + return converted; + } + if (value instanceof String text) { + Date date = parseMongoShellDate(text); + if (date == null) { + date = parseLegacyDateDisplay(text); + } + return date == null ? value : date; + } + return value; + } + + static Date parseMongoShellDate(String value) { + String trimmed = value.trim(); + String inner = null; + if (trimmed.startsWith("ISODate(") && trimmed.endsWith(")")) { + inner = trimmed.substring("ISODate(".length(), trimmed.length() - 1).trim(); + } else if (trimmed.startsWith("new Date(") && trimmed.endsWith(")")) { + inner = trimmed.substring("new Date(".length(), trimmed.length() - 1).trim(); + } + if (inner == null || inner.length() < 2) { + return null; + } + char quote = inner.charAt(0); + if ((quote != '"' && quote != '\'') || inner.charAt(inner.length() - 1) != quote) { + return null; + } + try { + Instant instant = Instant.parse(inner.substring(1, inner.length() - 1)); + return Date.from(instant); + } catch (Exception e) { + return null; + } + } + + static Date parseLegacyDateDisplay(String value) { + String trimmed = value.trim(); + if (!trimmed.matches("\\d{4}-\\d{2}-\\d{2}[ T]\\d{2}:\\d{2}:\\d{2}(\\.\\d{1,3})?")) { + return null; + } + String normalized = trimmed.replace(' ', 'T'); + int dot = normalized.indexOf('.'); + if (dot < 0) { + normalized = normalized + ".000"; + } else { + int millisStart = dot + 1; + int millisEnd = normalized.length(); + normalized = normalized.substring(0, millisStart) + + String.format("%-3s", normalized.substring(millisStart, millisEnd)).replace(' ', '0'); + } + try { + return Date.from(Instant.parse(normalized + "Z")); + } catch (Exception e) { + return null; + } + } + + private static Object dispatch(String method, JsonObject params) { + return switch (method) { + case AgentProtocol.METHOD_HANDSHAKE -> AgentProtocol.handshakeResult(); + case AgentProtocol.METHOD_CONNECT -> connect(params); + case AgentProtocol.MONGO_METHOD_LIST_DATABASES -> listDatabases(); + case AgentProtocol.MONGO_METHOD_LIST_COLLECTIONS -> listCollections(params); + case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS -> findDocuments(params); + case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENT -> insertDocument(params); + case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params); + case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENT -> deleteDocument(params); + case AgentProtocol.METHOD_DISCONNECT, AgentProtocol.METHOD_SHUTDOWN -> { + if (client != null) { + client.close(); + client = null; + } + if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) { + System.exit(0); + } + yield Collections.singletonMap("ok", true); + } + default -> throw new IllegalArgumentException("Unknown method: " + method); + }; + } + + private static MongoClient requireClient() { + if (client == null) { + throw new IllegalStateException("Not connected"); + } + return client; + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + private static String defaultString(String value, String fallback) { + return value == null ? fallback : value; + } + + static String handleRequest(String line) { + JsonObject req = JsonParser.parseString(line).getAsJsonObject(); + JsonElement id = req.get("id"); + String method = req.get("method").getAsString(); + JsonObject params = req.has("params") && req.get("params").isJsonObject() + ? req.getAsJsonObject("params") + : new JsonObject(); + + JsonObject response = new JsonObject(); + response.addProperty("jsonrpc", "2.0"); + response.add("id", id); + + try { + Object result = dispatch(method, params); + response.add("result", GSON.toJsonTree(result)); + } catch (Exception e) { + JsonObject error = new JsonObject(); + error.addProperty("code", -1); + error.addProperty("message", e.getMessage() == null ? "Unknown error" : e.getMessage()); + response.add("error", error); + } + + return GSON.toJson(response); + } + + public static void main(String[] args) throws Exception { + System.out.println("{\"ready\":true}"); + System.out.flush(); + + BufferedReader reader = new BufferedReader(new InputStreamReader(System.in)); + while (true) { + String line = reader.readLine(); + if (line == null) { + break; + } + + System.out.println(handleRequest(line)); + System.out.flush(); + } + } +} diff --git a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java new file mode 100644 index 0000000000..7c1458827a --- /dev/null +++ b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java @@ -0,0 +1,322 @@ +package com.dbx.agent.mongodb; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.dbx.agent.AgentProtocol; +import com.google.gson.JsonArray; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import com.mongodb.MongoClientSettings; +import java.io.FileInputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.security.KeyStore; +import java.security.PrivateKey; +import java.util.Base64; +import java.util.Date; +import org.bson.Document; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class MongoAgentTest { + @TempDir + static Path tempDir; + + private static Path caPemPath; + private static Path clientPemPath; + private static Path clientKeyPath; + + @BeforeAll + static void setUpCerts() throws Exception { + Path keystore = tempDir.resolve("keystore.jks"); + caPemPath = tempDir.resolve("ca.pem"); + clientPemPath = tempDir.resolve("client.pem"); + clientKeyPath = tempDir.resolve("client-key.pem"); + + // Generate a key pair in a JKS keystore using keytool + ProcessBuilder pb = new ProcessBuilder( + "keytool", "-genkeypair", "-alias", "test", "-keyalg", "RSA", "-keysize", "2048", + "-keystore", keystore.toString(), "-storepass", "pass123", "-keypass", "pass123", + "-dname", "CN=Test TLS Cert", "-validity", "365" + ); + pb.inheritIO(); + int rc = pb.start().waitFor(); + if (rc != 0) { + throw new RuntimeException("keytool -genkeypair failed with exit code " + rc); + } + + // Export the certificate as PEM (for ca_cert_path / client_cert_path) + for (Path pem : new Path[] {caPemPath, clientPemPath}) { + ProcessBuilder exportPb = new ProcessBuilder( + "keytool", "-exportcert", "-alias", "test", + "-keystore", keystore.toString(), "-storepass", "pass123", "-rfc" + ); + exportPb.redirectOutput(pem.toFile()); + exportPb.redirectError(ProcessBuilder.Redirect.INHERIT); + int exportRc = exportPb.start().waitFor(); + if (exportRc != 0) { + throw new RuntimeException("keytool -exportcert failed with exit code " + exportRc); + } + } + + // Extract the private key as PKCS#8 PEM + KeyStore ks = KeyStore.getInstance("JKS"); + try (FileInputStream fis = new FileInputStream(keystore.toFile())) { + ks.load(fis, "pass123".toCharArray()); + } + PrivateKey pk = (PrivateKey) ks.getKey("test", "pass123".toCharArray()); + String pkcs8Pem = "-----BEGIN PRIVATE KEY-----\n" + + Base64.getEncoder().encodeToString(pk.getEncoded()) + + "\n-----END PRIVATE KEY-----\n"; + Files.writeString(clientKeyPath, pkcs8Pem); + } + + // ─── existing tests ─── + + @Test + void exposesProtocolHandshakeOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"handshake\"," + + "\"params\":{\"appVersion\":\"0.5.13\",\"supportedProtocolVersions\":[1]}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + JsonObject result = json.getAsJsonObject("result"); + assertEquals("2.0", json.get("jsonrpc").getAsString()); + assertEquals(7, json.get("id").getAsInt()); + assertEquals(AgentProtocol.PROTOCOL_VERSION, result.get("protocolVersion").getAsInt()); + assertEquals(AgentProtocol.PROTOCOL_VERSION, result.get("agentProtocolVersion").getAsInt()); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_CONNECT)); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_QUERY)); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_METADATA)); + } + + @Test + void usesAuthSourceFromUrlParamsAsAuthenticationDatabase() { + JsonObject connection = new JsonObject(); + connection.addProperty("database", "gray_lite_twin_fat"); + connection.addProperty("url_params", "authSource=admin&authMechanism=SCRAM-SHA-1"); + + assertEquals("admin", MongoAgent.authenticationDatabase(connection)); + } + + @Test + void fallsBackToDefaultDatabaseWhenAuthSourceIsMissing() { + JsonObject connection = new JsonObject(); + connection.addProperty("database", "gray_lite_twin_fat"); + + assertEquals("gray_lite_twin_fat", MongoAgent.authenticationDatabase(connection)); + } + + // ─── TLS: configureBuilder JSON parsing ─── + + @Test + void sslTrueFromConnectionObject() { + JsonObject connection = minimalConnection(); + connection.addProperty("ssl", true); + + MongoClientSettings.Builder builder = MongoAgent.configureBuilder(connection); + + assertNotNull(builder); + } + + @Test + void sslFalseByDefault() { + JsonObject connection = minimalConnection(); + // ssl is not set — should default to false + + MongoClientSettings.Builder builder = MongoAgent.configureBuilder(connection); + + assertNotNull(builder); + } + + @Test + void sslTrueFromTopLevelParams() { + JsonObject connObj = minimalConnection(); + connObj.addProperty("ssl", true); + JsonObject params = new JsonObject(); + params.add("connection", connObj); + + // connect() unwraps the connection sub-object; verify configureBuilder reads ssl from it + JsonObject extracted = params.has("connection") && params.get("connection").isJsonObject() + ? params.getAsJsonObject("connection") + : params; + assertEquals(true, extracted.get("ssl").getAsBoolean()); + } + + @Test + void readsCaCertPathFromConnection() { + JsonObject connection = minimalConnection(); + connection.addProperty("ssl", true); + connection.addProperty("ca_cert_path", caPemPath.toString()); + + MongoClientSettings.Builder builder = MongoAgent.configureBuilder(connection); + + assertNotNull(builder); + } + + @Test + void readsClientCertAndKeyPathsFromConnection() { + JsonObject connection = minimalConnection(); + connection.addProperty("ssl", true); + connection.addProperty("client_cert_path", clientPemPath.toString()); + connection.addProperty("client_key_path", clientKeyPath.toString()); + + MongoClientSettings.Builder builder = MongoAgent.configureBuilder(connection); + + assertNotNull(builder); + } + + @Test + void certPathAndKeyPathFallbackNames() { + JsonObject connection = minimalConnection(); + connection.addProperty("ssl", true); + connection.addProperty("cert_path", clientPemPath.toString()); + connection.addProperty("key_path", clientKeyPath.toString()); + + // Should not throw — cert_path/key_path are fallback names for client_cert_path/client_key_path + MongoClientSettings.Builder builder = MongoAgent.configureBuilder(connection); + + assertNotNull(builder); + } + + @Test + void rejectsMismatchedClientCertAndKey() { + JsonObject connection = minimalConnection(); + connection.addProperty("ssl", true); + connection.addProperty("client_cert_path", clientPemPath.toString()); + // client_key_path is missing + + assertThrows(IllegalArgumentException.class, () -> MongoAgent.configureBuilder(connection)); + } + + // ─── TLS: SSLContext creation ─── + + @Test + void createsSslContextWithCaCert() throws Exception { + var ctx = MongoAgent.createTlsSslContext(caPemPath.toString(), null, null); + + assertNotNull(ctx); + } + + @Test + void createsSslContextWithClientCertAndKey() throws Exception { + var ctx = MongoAgent.createTlsSslContext(null, clientPemPath.toString(), clientKeyPath.toString()); + + assertNotNull(ctx); + } + + @Test + void createsSslContextWithAllCertPaths() throws Exception { + var ctx = MongoAgent.createTlsSslContext( + caPemPath.toString(), clientPemPath.toString(), clientKeyPath.toString()); + + assertNotNull(ctx); + } + + // ─── TLS: trust manager loading ─── + + @Test + void loadsTrustManagersFromPemFile() throws Exception { + var trustManagers = MongoAgent.loadTrustManagersFromPem(caPemPath.toString()); + + assertNotNull(trustManagers); + assertTrue(trustManagers.length > 0); + } + + // ─── TLS: key manager loading ─── + + @Test + void loadsKeyManagersFromPemFiles() throws Exception { + var keyManagers = MongoAgent.loadKeyManagersFromPem( + clientPemPath.toString(), clientKeyPath.toString()); + + assertNotNull(keyManagers); + assertTrue(keyManagers.length > 0); + } + + // ─── TLS: private key format support ─── + + @Test + void loadsPkcs8PrivateKeyFromPem() throws Exception { + var key = MongoAgent.loadPrivateKeyFromPem(clientKeyPath.toString()); + + assertNotNull(key); + assertEquals("RSA", key.getAlgorithm()); + } + + // ─── utility ─── + + @Test + void firstNonBlankReturnsFirstNonBlankValue() { + assertEquals("b", MongoAgent.firstNonBlank(null, "", "b", "c")); + assertEquals("a", MongoAgent.firstNonBlank("a", "b")); + } + + @Test + void firstNonBlankReturnsNullWhenAllBlank() { + assertEquals(null, MongoAgent.firstNonBlank(null, "", " ")); + } + + @Test + void convertValuePreservesUnsafeLongForJsonClients() { + assertEquals("2326645729978441729", MongoAgent.convertValue(2_326_645_729_978_441_729L)); + assertEquals("-2326645729978441729", MongoAgent.convertValue(-2_326_645_729_978_441_729L)); + } + + @Test + void convertValueKeepsSafeLongAsNumber() { + assertEquals(42L, MongoAgent.convertValue(42L)); + } + + @Test + void convertValueFormatsDatesAsMongoShellIsoDate() { + assertEquals("ISODate(\"2026-06-10T13:59:31.287Z\")", MongoAgent.convertValue(Date.from(java.time.Instant.parse("2026-06-10T13:59:31.287Z")))); + } + + @Test + void documentForWriteParsesMongoShellIsoDateStrings() { + Document doc = MongoAgent.documentForWrite("{\"$set\":{\"CreateDate\":\"ISODate(\\\"2026-06-10T13:59:31.287Z\\\")\"}}"); + + assertTrue(MongoAgent.isUpdateOperatorDocument(doc)); + Document set = (Document) doc.get("$set"); + assertTrue(set.get("CreateDate") instanceof Date); + } + + @Test + void documentForWriteParsesNestedMongoShellIsoDateStrings() { + Document doc = MongoAgent.documentForWrite("{\"items\":[{\"created\":\"new Date(\\\"2026-06-10T13:59:31.287Z\\\")\"}]}"); + + assertTrue(((Document) ((java.util.List) doc.get("items")).get(0)).get("created") instanceof Date); + } + + @Test + void documentForWriteParsesLegacyDateDisplayStrings() { + Document doc = MongoAgent.documentForWrite("{\"$set\":{\"CreateDate\":\"2025-08-14 02:25:43.718\"}}"); + + Document set = (Document) doc.get("$set"); + assertTrue(set.get("CreateDate") instanceof Date); + assertEquals(1_755_138_343_718L, ((Date) set.get("CreateDate")).getTime()); + } + + // ─── helpers ─── + + private static JsonObject minimalConnection() { + JsonObject conn = new JsonObject(); + conn.addProperty("host", "127.0.0.1"); + conn.addProperty("port", 27017); + return conn; + } + + private static boolean containsCapability(JsonArray capabilities, String expected) { + for (int i = 0; i < capabilities.size(); i++) { + if (expected.equals(capabilities.get(i).getAsString())) { + return true; + } + } + return false; + } +} diff --git a/agents/drivers/neo4j/build.gradle b/agents/drivers/neo4j/build.gradle new file mode 100644 index 0000000000..232f5958d5 --- /dev/null +++ b/agents/drivers/neo4j/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'org.neo4j:neo4j-jdbc-full-bundle:6.10.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Neo4j', 'Main-Class': 'com.dbx.agent.neo4j.Neo4jAgent') + } +} diff --git a/agents/drivers/neo4j/src/main/java/com/dbx/agent/neo4j/Neo4jAgent.java b/agents/drivers/neo4j/src/main/java/com/dbx/agent/neo4j/Neo4jAgent.java new file mode 100644 index 0000000000..9447b38dc0 --- /dev/null +++ b/agents/drivers/neo4j/src/main/java/com/dbx/agent/neo4j/Neo4jAgent.java @@ -0,0 +1,350 @@ +package com.dbx.agent.neo4j; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TransactionExecutor; +import com.dbx.agent.TriggerInfo; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public class Neo4jAgent extends BaseDatabaseAgent { + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public String setSchemaSQL(String schema) { + return ""; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("org.neo4j.jdbc.Neo4jDriver"); + String url = buildNeo4jUrl(params); + connection = DriverManager.getConnection(url, params.getUsername(), params.getPassword()); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("org.neo4j.jdbc.Neo4jDriver"); + String url = buildNeo4jUrl(params); + try (Connection conn = DriverManager.getConnection(url, params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + private static String buildNeo4jUrl(ConnectParams params) { + StringBuilder url = new StringBuilder(); + url.append("jdbc:neo4j://").append(params.getHost()).append(":").append(params.getPort()); + String database = params.getDatabase(); + if (database != null && !database.trim().isEmpty()) { + url.append("?database=").append(database.trim()); + } + return url.toString(); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + Connection conn = requireConnected(); + List result = new ArrayList<>(); + try (Statement stmt = conn.createStatement(); + ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString("name"))); + } + } + result.sort(Comparator.comparing(DatabaseInfo::getName)); + return result; + }); + } + + @Override + public List listSchemas() { + return Collections.emptyList(); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + Connection conn = requireConnected(); + // Try JDBC metadata first — the driver handles version compatibility. + List result = new ArrayList<>(); + DatabaseMetaData meta = conn.getMetaData(); + try (ResultSet rs = meta.getTables(null, blankToNull(schema), "%", new String[]{"TABLE", "VIEW"})) { + while (rs.next()) { + String name = rs.getString("TABLE_NAME"); + String type = rs.getString("TABLE_TYPE"); + String comment = getStringOrNull(rs, "REMARKS"); + if (name != null && !name.isEmpty()) { + result.add(new TableInfo(name, type, comment)); + } + } + } catch (Exception ignored) { + // Fall through to Cypher fallback + } + // Fallback: use Cypher for drivers that don't implement getTables(). + if (result.isEmpty()) { + try (Statement stmt = conn.createStatement(); + ResultSet rs = stmt.executeQuery("CALL db.labels()")) { + while (rs.next()) { + String label = getStringOrDefault(rs, "label", ""); + if (label.isEmpty()) { + label = rs.getString(1); + } + if (label != null && !label.isEmpty()) { + result.add(new TableInfo(label, "TABLE", null)); + } + } + } catch (Exception ignored) { + // Cypher fallback also failed + } + } + result.sort(Comparator.comparing(TableInfo::getName)); + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Connection conn = requireConnected(); + List result = new ArrayList<>(); + // Try JDBC metadata first. + DatabaseMetaData meta = conn.getMetaData(); + Set primaryKeys = new LinkedHashSet<>(); + try (ResultSet rs = meta.getPrimaryKeys(null, blankToNull(schema), table)) { + while (rs.next()) { + String pkColumn = rs.getString("COLUMN_NAME"); + if (pkColumn != null) { + primaryKeys.add(pkColumn); + } + } + } catch (Exception ignored) { + // Fall through + } + try (ResultSet rs = meta.getColumns(null, blankToNull(schema), table, "%")) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + if (name == null || name.isEmpty()) { + continue; + } + String type = rs.getString("TYPE_NAME"); + if (type == null) { + type = "Unknown"; + } + boolean nullable = true; + try { + nullable = "YES".equals(rs.getString("IS_NULLABLE")); + } catch (Exception ignored) { + // keep default + } + String defaultValue = getStringOrNull(rs, "COLUMN_DEF"); + String comment = getStringOrNull(rs, "REMARKS"); + result.add(new ColumnInfo( + name, type, nullable, defaultValue, + primaryKeys.contains(name), + comment, null, null, null, null + )); + } + } catch (Exception ignored) { + // Fall through to Cypher fallback + } + // Fallback: use Cypher for drivers that don't implement getColumns(). + if (result.isEmpty()) { + try (Statement stmt = conn.createStatement(); + ResultSet rs = stmt.executeQuery("CALL db.schema.nodeTypeProperties()")) { + while (rs.next()) { + String nodeLabels = rs.getString("nodeLabels"); + if (nodeLabels == null) { + nodeLabels = ""; + } + if (!nodeLabels.contains(table)) { + continue; + } + String propertyName = rs.getString("propertyName"); + if (propertyName == null) { + continue; + } + String propertyTypes = rs.getString("propertyTypes"); + if (propertyTypes == null) { + propertyTypes = "Unknown"; + } + boolean isNullable; + try { + isNullable = !rs.getBoolean("mandatory"); + } catch (Exception ignored) { + isNullable = true; + } + result.add(new ColumnInfo( + propertyName, propertyTypes, isNullable, + null, false, null, null, null, null, null + )); + } + } catch (Exception ignored) { + // Cypher fallback also failed + } + } + result.sort(Comparator.comparing(ColumnInfo::getName)); + return result; + }); + } + + private static String blankToNull(String value) { + if (value == null) return null; + String trimmed = value.trim(); + return trimmed.isEmpty() ? null : trimmed; + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + Connection conn = requireConnected(); + List result = new ArrayList<>(); + try (Statement stmt = conn.createStatement(); + ResultSet rs = stmt.executeQuery("SHOW INDEXES")) { + while (rs.next()) { + String labelsOrTypes = getStringOrDefault(rs, "labelsOrTypes", ""); + if (!labelsOrTypes.contains(table)) { + continue; + } + + String properties = getStringOrDefault(rs, "properties", ""); + String uniqueness = getStringOrDefault(rs, "uniqueness", ""); + result.add(new IndexInfo( + getStringOrDefault(rs, "name", ""), + parseProperties(properties), + "UNIQUE".equals(uniqueness.toUpperCase(Locale.ROOT)), + false, + null, + getStringOrNull(rs, "type"), + null, + null + )); + } + } + result.sort(Comparator.comparing(IndexInfo::getName)); + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public QueryResult executeTransaction(List statements, String schema) { + return TransactionExecutor.executeStatements(requireConnected(), statements, schema, this::setSchemaSQL, (stmt, sql) -> { + stmt.execute(sql); + return Math.max(stmt.getUpdateCount(), 0); + }); + } + + public QueryResult executeQuery(String sql, String schema) { + return executeQuery(sql, schema, new ExecuteQueryOptions()); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::stringResultValue + ); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Object stringResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + private static List parseProperties(String properties) { + String trimmed = properties; + if (trimmed.startsWith("[") && trimmed.endsWith("]") && trimmed.length() >= 2) { + trimmed = trimmed.substring(1, trimmed.length() - 1); + } + + List result = new ArrayList<>(); + for (String property : trimmed.split(",")) { + String value = property.trim(); + if (value.startsWith("\"") && value.endsWith("\"") && value.length() >= 2) { + value = value.substring(1, value.length() - 1); + } + if (!value.isEmpty()) { + result.add(value); + } + } + return result; + } + + private static String getStringOrDefault(ResultSet rs, String column, String defaultValue) { + try { + String value = rs.getString(column); + return value == null ? defaultValue : value; + } catch (Exception ignored) { + return defaultValue; + } + } + + private static String getStringOrNull(ResultSet rs, String column) { + try { + return rs.getString(column); + } catch (Exception ignored) { + return null; + } + } + + public static void main(String[] args) { + new JsonRpcServer(new Neo4jAgent()).run(); + } +} diff --git a/agents/drivers/neo4j/src/test/java/com/dbx/agent/neo4j/Neo4jAgentTest.java b/agents/drivers/neo4j/src/test/java/com/dbx/agent/neo4j/Neo4jAgentTest.java new file mode 100644 index 0000000000..0573ea3d69 --- /dev/null +++ b/agents/drivers/neo4j/src/test/java/com/dbx/agent/neo4j/Neo4jAgentTest.java @@ -0,0 +1,183 @@ +package com.dbx.agent.neo4j; + +import com.dbx.agent.QueryResult; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class Neo4jAgentTest { + @Test + void executesCypherWritesWithReturnThroughStatementExecute() throws Exception { + Neo4jAgent agent = new Neo4jAgent(); + List calls = new ArrayList<>(); + setConnectionForTest(agent, fakeConnection(calls)); + + QueryResult result = agent.executeQuery("CREATE (n:Person {name: 'Ada'}) RETURN n", null); + + Assertions.assertEquals(Collections.singletonList("n"), result.getColumns()); + Assertions.assertEquals( + Collections.singletonList(Collections.singletonList("(:Person {name: Ada})")), + result.getRows() + ); + Assertions.assertTrue(calls.contains("execute")); + Assertions.assertFalse(calls.contains("executeUpdate")); + } + + @Test + void executesTransactionsThroughStatementExecute() throws Exception { + Neo4jAgent agent = new Neo4jAgent(); + List calls = new ArrayList<>(); + setConnectionForTest(agent, fakeConnection(calls)); + + QueryResult result = agent.executeTransaction( + Arrays.asList( + "MATCH (n:Employee) WHERE elementId(n) = '4:abc:7' SET n.name = 'Grace'", + "CREATE (n:Employee {name: 'Linus'})" + ), + null + ); + + Assertions.assertEquals(0L, result.getAffected_rows()); + Assertions.assertEquals( + Arrays.asList("setAutoCommit:false", "execute", "execute", "commit", "setAutoCommit:true"), + calls + ); + Assertions.assertFalse(calls.contains("executeUpdate")); + } + + private static void setConnectionForTest(Neo4jAgent agent, Connection connection) throws Exception { + java.lang.reflect.Field field = Neo4jAgent.class.getDeclaredField("connection"); + field.setAccessible(true); + field.set(agent, connection); + } + + private static Connection fakeConnection(List calls) { + Statement statement = fakeStatement(calls); + return proxy(Connection.class, (unused, method, args) -> { + switch (method.getName()) { + case "createStatement": + return statement; + case "setAutoCommit": + calls.add("setAutoCommit:" + args[0]); + return null; + case "getAutoCommit": + return true; + case "commit": + calls.add("commit"); + return null; + case "rollback": + calls.add("rollback"); + return null; + case "close": + return null; + case "isClosed": + return false; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static Statement fakeStatement(List calls) { + ResultSet resultSet = fakeResultSet(); + return proxy(Statement.class, (unused, method, args) -> { + switch (method.getName()) { + case "execute": + calls.add("execute"); + return true; + case "executeUpdate": + calls.add("executeUpdate"); + throw new SQLException("syntax error or access rule violation - invalid syntax"); + case "executeQuery": + calls.add("executeQuery"); + return resultSet; + case "getResultSet": + return resultSet; + case "getUpdateCount": + return 0; + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet fakeResultSet() { + int[] index = {-1}; + ResultSetMetaData metadata = fakeMetadata(); + return proxy(ResultSet.class, (unused, method, args) -> { + switch (method.getName()) { + case "next": + index[0] += 1; + return index[0] == 0; + case "getMetaData": + return metadata; + case "getObject": + return "(:Person {name: Ada})"; + case "wasNull": + return false; + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSetMetaData fakeMetadata() { + return proxy(ResultSetMetaData.class, (unused, method, args) -> { + switch (method.getName()) { + case "getColumnCount": + return 1; + case "getColumnLabel": + return "n"; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + @SuppressWarnings("unchecked") + private static T proxy(Class type, InvocationHandler handler) { + return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, handler); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) { + return false; + } + if (type == Byte.TYPE) { + return (byte) 0; + } + if (type == Short.TYPE) { + return (short) 0; + } + if (type == Integer.TYPE) { + return 0; + } + if (type == Long.TYPE) { + return 0L; + } + if (type == Float.TYPE) { + return 0F; + } + if (type == Double.TYPE) { + return 0.0D; + } + if (type == Character.TYPE) { + return '\0'; + } + return null; + } +} diff --git a/agents/drivers/oceanbase-oracle/build.gradle b/agents/drivers/oceanbase-oracle/build.gradle new file mode 100644 index 0000000000..ef3e6f3ed9 --- /dev/null +++ b/agents/drivers/oceanbase-oracle/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.oceanbase:oceanbase-client:2.4.18' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'OceanBase Oracle Mode', 'Main-Class': 'com.dbx.agent.oceanbaseoracle.OceanBaseOracleAgent') + } +} diff --git a/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java b/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java new file mode 100644 index 0000000000..683a19f887 --- /dev/null +++ b/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java @@ -0,0 +1,449 @@ +package com.dbx.agent.oceanbaseoracle; + +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { + private static final String COMPATIBLE_OJDBC_VERSION = "compatibleOjdbcVersion"; + private static final String DEFAULT_COMPATIBLE_OJDBC_VERSION = "compatibleOjdbcVersion=8"; + private static final Set SYSTEM_SCHEMAS = Set.of( + "SYS", "SYSTEM", "SYSMAN", "DBSNMP", "SYSBACKUP", "SYSDG", "SYSKM", "OUTLN", + "AUDSYS", "LBACSYS", "DVF", "DVSYS", "APPQOSSYS", "CTXSYS", "MDSYS", "MDDATA", + "ORDSYS", "ORDDATA", "ORDPLUGINS", "XDB", "ANONYMOUS", "DIP", "EXFSYS", + "GSMADMIN_INTERNAL", "GSMCATUSER", "GSMUSER", "OJVMSYS", "OLAPSYS", + "ORACLE_OCM", "SI_INFORMTN_SCHEMA", "WMSYS", "XS$NULL", "DBSFWUSER", + "REMOTE_SCHEDULER_AGENT", "PDBADMIN", "DGPDB_INT", "OPS$ORACLE", + "GGSYS", "FLOWS_FILES", "APEX_PUBLIC_USER", "GSMROOTUSER", "SYSRAC" + ); + + public static final JdbcAgentProfile OCEANBASE_ORACLE_PROFILE = new JdbcAgentProfile( + "com.oceanbase.jdbc.Driver", + "jdbc:oceanbase://{host}:{port}/{database}", + 2881, + false, + SYSTEM_SCHEMAS, + List.of("TABLE", "VIEW", "BASE TABLE") + ) { + @Override + public String schemaSwitchSql(String schema, String quote) { + return "ALTER SESSION SET CURRENT_SCHEMA = " + quote + schema.replace(quote, quote + quote) + quote; + } + }; + + public OceanBaseOracleAgent() { + super(OCEANBASE_ORACLE_PROFILE); + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + static String buildUrl(ConnectParams params) { + return appendDefaultCompatibilityOption(OCEANBASE_ORACLE_PROFILE.buildUrl(params)); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + for (String schema : querySchemas()) { + result.add(new DatabaseInfo(schema)); + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(this::querySchemas); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + String owner = normalizeSchema(schema); + String sql = """ + SELECT o.OBJECT_NAME, + CASE o.OBJECT_TYPE WHEN 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS TABLE_TYPE, + c.COMMENTS + FROM ALL_OBJECTS o + LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME + WHERE o.OWNER = ? AND o.OBJECT_TYPE IN ('TABLE', 'VIEW') + ORDER BY o.OBJECT_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, owner); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + String owner = normalizeSchema(schema); + String sql = """ + SELECT OBJECT_NAME, OBJECT_TYPE + FROM ALL_OBJECTS + WHERE OWNER = ? AND OBJECT_TYPE IN ('TABLE', 'VIEW', 'PROCEDURE', 'FUNCTION', 'PACKAGE', 'SEQUENCE', 'SYNONYM') + ORDER BY CASE OBJECT_TYPE + WHEN 'TABLE' THEN 0 + WHEN 'VIEW' THEN 1 + WHEN 'PROCEDURE' THEN 2 + WHEN 'FUNCTION' THEN 3 + WHEN 'PACKAGE' THEN 4 + WHEN 'SEQUENCE' THEN 5 + ELSE 6 + END, OBJECT_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, owner); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), owner, null)); + } + } + } + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + String owner = normalizeSchema(schema); + String sql = """ + SELECT c.COLUMN_NAME, c.DATA_TYPE, c.NULLABLE, c.DATA_PRECISION, c.DATA_SCALE, + c.DATA_LENGTH, c.CHAR_LENGTH, cc.COMMENTS, + CASE WHEN pk.COLUMN_NAME IS NULL THEN 0 ELSE 1 END AS IS_PK + FROM ALL_TAB_COLUMNS c + LEFT JOIN ALL_COL_COMMENTS cc + ON cc.OWNER = c.OWNER AND cc.TABLE_NAME = c.TABLE_NAME AND cc.COLUMN_NAME = c.COLUMN_NAME + LEFT JOIN ( + SELECT cols.COLUMN_NAME + FROM ALL_CONS_COLUMNS cols + JOIN ALL_CONSTRAINTS cons + ON cols.CONSTRAINT_NAME = cons.CONSTRAINT_NAME AND cols.OWNER = cons.OWNER + WHERE cons.CONSTRAINT_TYPE = 'P' AND cons.OWNER = ? AND cons.TABLE_NAME = ? + ) pk ON pk.COLUMN_NAME = c.COLUMN_NAME + WHERE c.OWNER = ? AND c.TABLE_NAME = ? + ORDER BY c.COLUMN_ID + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, owner); + stmt.setString(2, table); + stmt.setString(3, owner); + stmt.setString(4, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + String baseType = rs.getString("DATA_TYPE"); + Integer numPrec = intOrNull(rs, "DATA_PRECISION"); + Integer numScale = intOrNull(rs, "DATA_SCALE"); + Integer dataLen = intOrNull(rs, "DATA_LENGTH"); + Integer charLen = intOrNull(rs, "CHAR_LENGTH"); + result.add(new ColumnInfo( + name, + formatDataType(baseType, numPrec, numScale, dataLen, charLen), + "Y".equalsIgnoreCase(rs.getString("NULLABLE")), + null, + rs.getInt("IS_PK") == 1, + null, + rs.getString("COMMENTS"), + numPrec, + numScale, + charLen + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + String owner = normalizeSchema(schema); + String sql = """ + SELECT i.INDEX_NAME, ic.COLUMN_NAME, ic.COLUMN_POSITION, i.UNIQUENESS, + c.CONSTRAINT_TYPE, i.INDEX_TYPE + FROM ALL_INDEXES i + JOIN ALL_IND_COLUMNS ic + ON i.INDEX_NAME = ic.INDEX_NAME + AND i.OWNER = ic.INDEX_OWNER + AND i.TABLE_OWNER = ic.TABLE_OWNER + AND i.TABLE_NAME = ic.TABLE_NAME + LEFT JOIN ALL_CONSTRAINTS c + ON i.INDEX_NAME = c.INDEX_NAME + AND i.TABLE_OWNER = c.OWNER + AND i.TABLE_NAME = c.TABLE_NAME + AND c.CONSTRAINT_TYPE = 'P' + WHERE i.TABLE_OWNER = ? AND i.TABLE_NAME = ? + ORDER BY i.INDEX_NAME, ic.COLUMN_POSITION + """.stripIndent().trim(); + + Map> columnsByIndex = new LinkedHashMap<>(); + Map uniqueByIndex = new LinkedHashMap<>(); + Map primaryByIndex = new LinkedHashMap<>(); + Map typeByIndex = new LinkedHashMap<>(); + try (var stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, owner); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String indexName = rs.getString("INDEX_NAME"); + columnsByIndex.computeIfAbsent(indexName, ignored -> new ArrayList<>()).add(rs.getString("COLUMN_NAME")); + uniqueByIndex.put(indexName, "UNIQUE".equalsIgnoreCase(rs.getString("UNIQUENESS"))); + primaryByIndex.put(indexName, "P".equalsIgnoreCase(rs.getString("CONSTRAINT_TYPE"))); + typeByIndex.put(indexName, rs.getString("INDEX_TYPE")); + } + } + } + + List result = new ArrayList<>(); + for (Map.Entry> entry : columnsByIndex.entrySet()) { + String name = entry.getKey(); + result.add(new IndexInfo( + name, + entry.getValue(), + Boolean.TRUE.equals(uniqueByIndex.get(name)), + Boolean.TRUE.equals(primaryByIndex.get(name)), + null, + typeByIndex.get(name), + null, + null + )); + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + String owner = normalizeSchema(schema); + String sql = """ + SELECT c.CONSTRAINT_NAME, cc.COLUMN_NAME, rc.TABLE_NAME, rcc.COLUMN_NAME + FROM ALL_CONSTRAINTS c + JOIN ALL_CONS_COLUMNS cc ON c.CONSTRAINT_NAME = cc.CONSTRAINT_NAME AND c.OWNER = cc.OWNER + JOIN ALL_CONSTRAINTS rc ON c.R_CONSTRAINT_NAME = rc.CONSTRAINT_NAME AND c.R_OWNER = rc.OWNER + JOIN ALL_CONS_COLUMNS rcc + ON rc.CONSTRAINT_NAME = rcc.CONSTRAINT_NAME + AND rc.OWNER = rcc.OWNER + AND cc.POSITION = rcc.POSITION + WHERE c.CONSTRAINT_TYPE = 'R' AND c.OWNER = ? AND c.TABLE_NAME = ? + ORDER BY c.CONSTRAINT_NAME, cc.POSITION + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, owner); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString(1), + rs.getString(2), + rs.getString(3), + rs.getString(4) + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + String owner = normalizeSchema(schema); + String sql = """ + SELECT TRIGGER_NAME, TRIGGERING_EVENT, TRIGGER_TYPE + FROM ALL_TRIGGERS + WHERE OWNER = ? AND TABLE_NAME = ? + ORDER BY TRIGGER_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnection().prepareStatement(sql)) { + stmt.setString(1, owner); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public String setSchemaSQL(String schema) { + return "ALTER SESSION SET CURRENT_SCHEMA = " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + private List querySchemas() throws SQLException { + try { + return querySchemaNames(); + } catch (SQLException primaryError) { + String current; + try { + current = currentSchema(); + } catch (SQLException ignored) { + throw primaryError; + } + if (current == null || current.isBlank()) { + throw primaryError; + } + return List.of(current); + } + } + + private List querySchemaNames() throws SQLException { + String placeholders = SYSTEM_SCHEMAS.stream() + .map(schema -> "'" + schema + "'") + .collect(Collectors.joining(",")); + String sql = """ + SELECT username + FROM ALL_USERS + WHERE username IS NOT NULL + AND username NOT IN (%s) + AND username NOT LIKE 'APEX_%%' + AND username NOT LIKE 'FLOWS_%%' + AND username NOT LIKE '%%$%%' + ORDER BY CASE + WHEN username = SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') THEN 0 + WHEN username = SYS_CONTEXT('USERENV', 'SESSION_USER') THEN 1 + ELSE 2 + END, username + """.formatted(placeholders).stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnection().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + String schema = rs.getString(1); + if (schema != null && !schema.isBlank()) { + result.add(schema); + } + } + } + return result; + } + + private String normalizeSchema(String schema) throws SQLException { + String value = schema == null ? "" : schema.trim(); + return value.isEmpty() ? currentSchema() : value; + } + + private String currentSchema() throws SQLException { + try (var stmt = requireConnection().createStatement(); + ResultSet rs = stmt.executeQuery("SELECT SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') FROM DUAL")) { + if (rs.next()) { + String schema = rs.getString(1); + if (schema != null && !schema.isBlank()) { + return schema; + } + } + } + return ""; + } + + private static String appendDefaultCompatibilityOption(String url) { + if (hasQueryKey(url, COMPATIBLE_OJDBC_VERSION)) { + return url; + } + return url + (url.contains("?") ? "&" : "?") + DEFAULT_COMPATIBLE_OJDBC_VERSION; + } + + private static boolean hasQueryKey(String url, String key) { + int queryStart = url.indexOf('?'); + if (queryStart < 0) { + return false; + } + String query = url.substring(queryStart + 1); + int fragmentStart = query.indexOf('#'); + if (fragmentStart >= 0) { + query = query.substring(0, fragmentStart); + } + for (String part : query.split("[&;]")) { + String normalized = part.trim(); + if (normalized.isEmpty()) { + continue; + } + int equals = normalized.indexOf('='); + String paramKey = equals >= 0 ? normalized.substring(0, equals) : normalized; + if (paramKey.trim().equalsIgnoreCase(key)) { + return true; + } + } + return false; + } + + private static String formatDataType(String base, Integer numPrec, Integer numScale, Integer dataLen, Integer charLen) { + if (base == null || base.isBlank()) { + return ""; + } + return switch (base.toUpperCase(Locale.ROOT)) { + case "VARCHAR2", "NVARCHAR2", "CHAR", "NCHAR" -> { + Integer len = charLen == null ? dataLen : charLen; + yield len == null ? base : base + "(" + len + ")"; + } + case "NUMBER" -> { + if (numPrec != null && numScale != null && numScale > 0) { + yield base + "(" + numPrec + "," + numScale + ")"; + } + if (numPrec != null && numPrec > 0) { + yield base + "(" + numPrec + ")"; + } + yield base; + } + case "RAW" -> dataLen == null ? "RAW" : "RAW(" + dataLen + ")"; + default -> base; + }; + } + + private static Integer intOrNull(ResultSet rs, String column) throws SQLException { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + public static void main(String[] args) { + new JsonRpcServer(new OceanBaseOracleAgent()).run(); + } +} diff --git a/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java b/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java new file mode 100644 index 0000000000..1fd9b0901a --- /dev/null +++ b/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java @@ -0,0 +1,59 @@ +package com.dbx.agent.oceanbaseoracle; + +import com.dbx.agent.ConnectParams; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class OceanBaseOracleAgentTest { + @Test + void buildsOceanBaseJdbcUrl() { + ConnectParams params = new ConnectParams(); + params.setHost("oceanbase.example.com"); + params.setPort(0); + params.setDatabase("sys"); + + Assertions.assertEquals( + "jdbc:oceanbase://oceanbase.example.com:2881/sys?compatibleOjdbcVersion=8", + OceanBaseOracleAgent.buildUrl(params) + ); + } + + @Test + void appendsQueryParametersToJdbcUrl() { + ConnectParams params = new ConnectParams(); + params.setHost("oceanbase.example.com"); + params.setPort(2881); + params.setDatabase("sys"); + params.setUrl_params("useSSL=false"); + + Assertions.assertEquals( + "jdbc:oceanbase://oceanbase.example.com:2881/sys?useSSL=false&compatibleOjdbcVersion=8", + OceanBaseOracleAgent.buildUrl(params) + ); + } + + @Test + void keepsExplicitCompatibleOjdbcVersion() { + ConnectParams params = new ConnectParams(); + params.setHost("oceanbase.example.com"); + params.setPort(2881); + params.setDatabase("sys"); + params.setUrl_params("compatibleOjdbcVersion=6&useSSL=false"); + + Assertions.assertEquals( + "jdbc:oceanbase://oceanbase.example.com:2881/sys?compatibleOjdbcVersion=6&useSSL=false", + OceanBaseOracleAgent.buildUrl(params) + ); + } + + @Test + void appendsCompatibleOjdbcVersionToCustomJdbcUrl() { + ConnectParams params = new ConnectParams(); + params.setConnection_string("jdbc:oceanbase://custom-host:2881/sys?useSSL=false"); + + Assertions.assertEquals( + "jdbc:oceanbase://custom-host:2881/sys?useSSL=false&compatibleOjdbcVersion=8", + OceanBaseOracleAgent.buildUrl(params) + ); + } +} diff --git a/agents/drivers/oracle-10g/build.gradle b/agents/drivers/oracle-10g/build.gradle new file mode 100644 index 0000000000..a38986ab84 --- /dev/null +++ b/agents/drivers/oracle-10g/build.gradle @@ -0,0 +1,11 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.oracle.database.jdbc:ojdbc6:11.2.0.4' + implementation 'com.oracle.database.nls:orai18n:11.2.0.4' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Oracle 10g', 'Main-Class': 'com.dbx.agent.oracle10g.Oracle10gAgent') + } +} diff --git a/agents/drivers/oracle-10g/src/main/java/com/dbx/agent/oracle10g/Oracle10gAgent.java b/agents/drivers/oracle-10g/src/main/java/com/dbx/agent/oracle10g/Oracle10gAgent.java new file mode 100644 index 0000000000..5c74825645 --- /dev/null +++ b/agents/drivers/oracle-10g/src/main/java/com/dbx/agent/oracle10g/Oracle10gAgent.java @@ -0,0 +1,515 @@ +package com.dbx.agent.oracle10g; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryPageOptions; +import com.dbx.agent.QueryPageResult; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Locale; +import java.util.Properties; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +public final class Oracle10gAgent extends BaseDatabaseAgent { + private static final List SYSTEM_SCHEMAS = Arrays.asList( + "SYS", "SYSTEM", "SYSMAN", "DBSNMP", "SYSBACKUP", "SYSDG", "SYSKM", "OUTLN", + "AUDSYS", "LBACSYS", "DVF", "DVSYS", "APPQOSSYS", "CTXSYS", "MDSYS", "MDDATA", + "ORDSYS", "ORDDATA", "ORDPLUGINS", "XDB", "ANONYMOUS", "DIP", "EXFSYS", + "GSMADMIN_INTERNAL", "GSMCATUSER", "GSMUSER", "OJVMSYS", "OLAPSYS", + "ORACLE_OCM", "SI_INFORMTN_SCHEMA", "WMSYS", "XS$NULL", "DBSFWUSER", + "REMOTE_SCHEDULER_AGENT", "PDBADMIN", "DGPDB_INT", "OPS$ORACLE", + "GGSYS", "FLOWS_FILES", "APEX_PUBLIC_USER" + ); + private static final Pattern OFFSET_FETCH_RE = Pattern.compile( + "(.+?)\\s+OFFSET\\s+(\\d+)\\s+ROWS?\\s+FETCH\\s+(FIRST|NEXT)\\s+(\\d+)\\s+ROWS?\\s+ONLY", + Pattern.CASE_INSENSITIVE + ); + private static final Pattern FETCH_ONLY_RE = Pattern.compile( + "(.+?)\\s+FETCH\\s+(FIRST|NEXT)\\s+(\\d+)\\s+ROWS?\\s+ONLY", + Pattern.CASE_INSENSITIVE + ); + private static final Pattern PLSQL_OBJECT_DDL_RE = Pattern.compile( + "^\\s*CREATE\\s+(?:OR\\s+REPLACE\\s+)?(?:(?:NON)?EDITIONABLE\\s+)?(?:PROCEDURE|FUNCTION|PACKAGE(?:\\s+BODY)?|TRIGGER|TYPE(?:\\s+BODY)?)\\b", + Pattern.CASE_INSENSITIVE + ); + + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + public static String rewriteFetchFirst(String sql) { + Matcher offsetFetch = OFFSET_FETCH_RE.matcher(sql); + if (offsetFetch.matches()) { + String innerSql = offsetFetch.group(1); + long offset = Long.parseLong(offsetFetch.group(2)); + long limit = Long.parseLong(offsetFetch.group(4)); + long upper = offset + limit; + return "SELECT * FROM (SELECT a.*, ROWNUM rn__ FROM (" + innerSql + ") a WHERE ROWNUM <= " + upper + ") WHERE rn__ > " + offset; + } + + Matcher fetchOnly = FETCH_ONLY_RE.matcher(sql); + if (fetchOnly.matches()) { + String innerSql = fetchOnly.group(1); + long limit = Long.parseLong(fetchOnly.group(3)); + return "SELECT * FROM (" + innerSql + ") WHERE ROWNUM <= " + limit; + } + + return sql; + } + + @Override + public String setSchemaSQL(String schema) { + return "ALTER SESSION SET CURRENT_SCHEMA = " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("oracle.jdbc.OracleDriver"); + connection = DriverManager.getConnection(buildUrl(params), connectionProperties(params)); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("oracle.jdbc.OracleDriver"); + try (Connection conn = DriverManager.getConnection(buildUrl(params), connectionProperties(params))) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + String placeholders = quotedSystemSchemas(); + String sql = "SELECT owner FROM (" + + " SELECT SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') AS owner FROM DUAL" + + " UNION" + + " SELECT DISTINCT owner FROM all_tables" + + " UNION" + + " SELECT DISTINCT owner FROM all_views" + + " )" + + " WHERE owner IS NOT NULL" + + " AND owner NOT IN (" + placeholders + ")" + + " AND owner NOT LIKE 'APEX_%'" + + " AND owner NOT LIKE 'FLOWS_%'" + + " ORDER BY owner"; + try (Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT o.OBJECT_NAME," + + " CASE o.OBJECT_TYPE WHEN 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS TABLE_TYPE," + + " c.COMMENTS" + + " FROM ALL_OBJECTS o" + + " LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME" + + " WHERE o.OWNER = ? AND o.OBJECT_TYPE IN ('TABLE','VIEW')" + + " ORDER BY o.OBJECT_NAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT OBJECT_NAME, OBJECT_TYPE FROM ALL_OBJECTS" + + " WHERE OWNER = ? AND OBJECT_TYPE IN ('TABLE', 'VIEW', 'PROCEDURE', 'FUNCTION')" + + " ORDER BY CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 ELSE 3 END, OBJECT_NAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String upperType = objectType.toUpperCase(Locale.ROOT); + String source; + if ("VIEW".equals(upperType)) { + // ALL_VIEWS is more reliable than DBMS_METADATA.GET_DDL, + // which fails on XE editions where XSL stylesheets are missing. + String sql = "SELECT TEXT FROM ALL_VIEWS WHERE OWNER = ? AND VIEW_NAME = ?"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, name); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + } else { + String dbmsType; + switch (upperType) { + case "PROCEDURE": + dbmsType = "PROCEDURE"; + break; + case "FUNCTION": + dbmsType = "FUNCTION"; + break; + default: + throw new IllegalArgumentException("Unsupported object type: " + objectType); + } + String sql = "SELECT DBMS_METADATA.GET_DDL(?, ?, ?) FROM DUAL"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, dbmsType); + stmt.setString(2, name); + stmt.setString(3, schema); + try (ResultSet rs = stmt.executeQuery()) { + source = rs.next() ? coalesce(rs.getString(1)) : ""; + } + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public String getTableDdl(String schema, String table) { + return unchecked(() -> { + String sql = "SELECT DBMS_METADATA.GET_DDL(?, ?, ?) FROM DUAL"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, "TABLE"); + stmt.setString(2, table); + stmt.setString(3, schema); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + return coalesce(rs.getString(1)); + } + } + } + throw new IllegalArgumentException("Table not found: " + schema + "." + table); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT c.COLUMN_NAME, c.DATA_TYPE, c.NULLABLE, c.DATA_PRECISION, c.DATA_SCALE," + + " c.DATA_LENGTH, c.CHAR_LENGTH, cc.COMMENTS," + + " CASE WHEN pk.COLUMN_NAME IS NULL THEN 0 ELSE 1 END AS IS_PK" + + " FROM ALL_TAB_COLUMNS c" + + " LEFT JOIN ALL_COL_COMMENTS cc" + + " ON cc.OWNER = c.OWNER AND cc.TABLE_NAME = c.TABLE_NAME AND cc.COLUMN_NAME = c.COLUMN_NAME" + + " LEFT JOIN (" + + " SELECT cols.COLUMN_NAME FROM ALL_CONS_COLUMNS cols" + + " JOIN ALL_CONSTRAINTS cons" + + " ON cols.CONSTRAINT_NAME = cons.CONSTRAINT_NAME AND cols.OWNER = cons.OWNER" + + " WHERE cons.CONSTRAINT_TYPE = 'P' AND cons.OWNER = ? AND cons.TABLE_NAME = ?" + + " ) pk ON pk.COLUMN_NAME = c.COLUMN_NAME" + + " WHERE c.OWNER = ? AND c.TABLE_NAME = ?" + + " ORDER BY c.COLUMN_ID"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + stmt.setString(3, schema); + stmt.setString(4, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + String baseType = rs.getString("DATA_TYPE"); + Integer numPrec = intObject(rs, "DATA_PRECISION"); + Integer numScale = intObject(rs, "DATA_SCALE"); + Integer dataLen = intObject(rs, "DATA_LENGTH"); + Integer charLen = intObject(rs, "CHAR_LENGTH"); + String dataType = formatDataType(baseType, numPrec, numScale, dataLen, charLen); + + result.add(new ColumnInfo( + name, + dataType, + "Y".equals(rs.getString("NULLABLE")), + null, + rs.getInt("IS_PK") == 1, + null, + rs.getString("COMMENTS"), + numPrec, + numScale, + charLen + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT i.INDEX_NAME," + + " RTRIM(XMLAGG(XMLELEMENT(e, ic.COLUMN_NAME || ',') ORDER BY ic.COLUMN_POSITION).EXTRACT('//text()').GETSTRINGVAL(), ',') AS COLUMNS," + + " i.UNIQUENESS," + + " CASE WHEN c.CONSTRAINT_TYPE = 'P' THEN 1 ELSE 0 END AS IS_PK," + + " i.INDEX_TYPE" + + " FROM ALL_INDEXES i" + + " JOIN ALL_IND_COLUMNS ic ON i.INDEX_NAME = ic.INDEX_NAME AND i.OWNER = ic.INDEX_OWNER AND i.TABLE_OWNER = ic.TABLE_OWNER" + + " LEFT JOIN ALL_CONSTRAINTS c ON i.INDEX_NAME = c.INDEX_NAME AND i.TABLE_OWNER = c.OWNER" + + " AND c.CONSTRAINT_TYPE = 'P'" + + " WHERE i.TABLE_OWNER = ? AND i.TABLE_NAME = ?" + + " GROUP BY i.INDEX_NAME, i.UNIQUENESS, c.CONSTRAINT_TYPE, i.INDEX_TYPE" + + " ORDER BY i.INDEX_NAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new IndexInfo( + rs.getString(1), + splitNonEmpty(coalesce(rs.getString(2)), ","), + "UNIQUE".equals(rs.getString(3)), + "1".equals(rs.getString(4)), + null, + rs.getString(5), + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT c.CONSTRAINT_NAME, cc.COLUMN_NAME, rc.TABLE_NAME, rcc.COLUMN_NAME" + + " FROM ALL_CONSTRAINTS c" + + " JOIN ALL_CONS_COLUMNS cc ON c.CONSTRAINT_NAME = cc.CONSTRAINT_NAME AND c.OWNER = cc.OWNER" + + " JOIN ALL_CONSTRAINTS rc ON c.R_CONSTRAINT_NAME = rc.CONSTRAINT_NAME AND c.R_OWNER = rc.OWNER" + + " JOIN ALL_CONS_COLUMNS rcc ON rc.CONSTRAINT_NAME = rcc.CONSTRAINT_NAME AND rc.OWNER = rcc.OWNER" + + " WHERE c.CONSTRAINT_TYPE = 'R' AND c.OWNER = ? AND c.TABLE_NAME = ?" + + " ORDER BY c.CONSTRAINT_NAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo(rs.getString(1), rs.getString(2), rs.getString(3), rs.getString(4))); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT TRIGGER_NAME, TRIGGERING_EVENT, TRIGGER_TYPE" + + " FROM ALL_TRIGGERS" + + " WHERE OWNER = ? AND TABLE_NAME = ?" + + " ORDER BY TRIGGER_NAME"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + prepareExecutableSql(sql), + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::stringResultValue + ); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + return JdbcExecutor.INSTANCE.executePage( + requireConnected(), + prepareExecutableSql(sql), + schema, + this::setSchemaSQL, + options, + this::stringResultValue + ); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private Object stringResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + return JdbcExecutor.stringResultValue(rs, index, sqlType); + }); + } + + static String buildUrl(ConnectParams params) { + String connectionString = params.getConnection_string(); + if (connectionString != null && !connectionString.trim().isEmpty()) { + return connectionString; + } + return "jdbc:oracle:thin:@" + params.getHost() + ":" + params.getPort() + "/" + serviceName(params); + } + + private static Properties connectionProperties(ConnectParams params) { + Properties props = new Properties(); + props.put("user", params.getUsername()); + props.put("password", params.getPassword()); + if (params.getDatabase().toUpperCase(Locale.ROOT).startsWith("SYSDBA:")) { + props.put("internal_logon", "SYSDBA"); + } + return props; + } + + private static String serviceName(ConnectParams params) { + String serviceName = params.getDatabase(); + if (serviceName.toUpperCase(Locale.ROOT).startsWith("SYSDBA:")) { + return serviceName.substring(7); + } + return serviceName; + } + + private static String quotedSystemSchemas() { + List quoted = new ArrayList<>(); + for (String schema : SYSTEM_SCHEMAS) { + quoted.add("'" + schema + "'"); + } + return String.join(",", quoted); + } + + private static String trimTrailingSemicolons(String sql) { + String trimmed = sql; + while (trimmed.endsWith(";")) { + trimmed = trimmed.substring(0, trimmed.length() - 1); + } + return trimmed; + } + + static String prepareExecutableSql(String sql) { + String trimmed = sql.trim(); + if (PLSQL_OBJECT_DDL_RE.matcher(trimmed).find()) { + return trimmed; + } + return rewriteFetchFirst(trimTrailingSemicolons(trimmed)); + } + + private static String formatDataType(String base, Integer numPrec, Integer numScale, Integer dataLen, Integer charLen) { + switch (base.toUpperCase(Locale.ROOT)) { + case "VARCHAR2": + case "NVARCHAR2": + case "CHAR": + case "NCHAR": + Integer length = charLen != null ? charLen : dataLen; + return length != null ? base + "(" + length + ")" : base; + case "NUMBER": + if (numPrec != null && numScale != null && numScale > 0) { + return base + "(" + numPrec + "," + numScale + ")"; + } + return numPrec != null && numPrec > 0 ? base + "(" + numPrec + ")" : base; + case "RAW": + return dataLen != null ? "RAW(" + dataLen + ")" : "RAW"; + default: + return base; + } + } + + private static Integer intObject(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value == null ? null : ((Number) value).intValue(); + } + + private static List splitNonEmpty(String value, String delimiter) { + List result = new ArrayList<>(); + for (String part : value.split(delimiter)) { + if (!part.isEmpty()) { + result.add(part); + } + } + return result; + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + public static void main(String[] args) { + new JsonRpcServer(new Oracle10gAgent()).run(); + } +} diff --git a/agents/drivers/oracle-10g/src/test/java/com/dbx/agent/oracle10g/Oracle10gAgentTest.java b/agents/drivers/oracle-10g/src/test/java/com/dbx/agent/oracle10g/Oracle10gAgentTest.java new file mode 100644 index 0000000000..b6d9cd3e78 --- /dev/null +++ b/agents/drivers/oracle-10g/src/test/java/com/dbx/agent/oracle10g/Oracle10gAgentTest.java @@ -0,0 +1,205 @@ +package com.dbx.agent.oracle10g; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.test.JdbcAgentFake; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class Oracle10gAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new Oracle10gAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL DBMS_XPLAN.DISPLAY_CURSOR()"; + } + + @Test + void agentInheritsBaseDatabaseAgent() { + Assertions.assertTrue(createAgent() instanceof BaseDatabaseAgent); + } + + @Test + void buildUrlUsesExplicitConnectionString() { + ConnectParams params = new ConnectParams( + "oracle.example.com", + 1521, + "ORCL", + "scott", + "tiger", + "", + "jdbc:oracle:thin:@oracle.example.com:1521:ORCL", + false + ); + + Assertions.assertEquals("jdbc:oracle:thin:@oracle.example.com:1521:ORCL", Oracle10gAgent.buildUrl(params)); + } + + @Test + void prepareExecutableSqlKeepsPlsqlObjectTerminator() { + String sql = "CREATE OR REPLACE PROCEDURE APP_PROC AS BEGIN NULL; END;"; + + Assertions.assertEquals(sql, Oracle10gAgent.prepareExecutableSql(sql)); + } + + @Test + void prepareExecutableSqlTrimsPlainStatementTerminator() { + Assertions.assertEquals("SELECT 1 FROM DUAL", Oracle10gAgent.prepareExecutableSql("SELECT 1 FROM DUAL;")); + } + + @Test + void prepareExecutableSqlStillRewritesFetchFirst() { + Assertions.assertEquals( + "SELECT * FROM (SELECT * FROM EMP) WHERE ROWNUM <= 10", + Oracle10gAgent.prepareExecutableSql("SELECT * FROM EMP FETCH FIRST 10 ROWS ONLY;") + ); + } + + @Test + void listsTablesViewsProceduresAndFunctions() { + Oracle10gAgent agent = new Oracle10gAgent(); + TestSupport.setPrivateConnection(agent, objectListConnection()); + + List objects = agent.listObjects("APP"); + + Assertions.assertEquals( + Arrays.asList("TABLE", "VIEW", "PROCEDURE", "FUNCTION"), + objects.stream().map(ObjectInfo::getObject_type).collect(Collectors.toList()) + ); + Assertions.assertEquals( + Arrays.asList("APP_TABLE", "APP_VIEW", "APP_PROC", "APP_FUNC"), + objects.stream().map(ObjectInfo::getName).collect(Collectors.toList()) + ); + } + + @Test + void loadsRoutineSourceFromDbmsMetadata() { + Oracle10gAgent agent = new Oracle10gAgent(); + TestSupport.setPrivateConnection( + agent, + metadataConnection(sql -> Arrays.asList( + Arrays.asList("CREATE OR REPLACE PROCEDURE APP_PROC AS BEGIN NULL; END;") + )) + ); + + ObjectSource source = agent.getObjectSource("APP", "APP_PROC", "PROCEDURE"); + + Assertions.assertEquals("APP_PROC", source.getName()); + Assertions.assertEquals("PROCEDURE", source.getObject_type()); + Assertions.assertEquals("APP", source.getSchema()); + Assertions.assertEquals("CREATE OR REPLACE PROCEDURE APP_PROC AS BEGIN NULL; END;", source.getSource()); + } + + private static Connection objectListConnection() { + return metadataConnection(sql -> { + if (sql.contains("'PROCEDURE'")) { + return Arrays.asList( + Arrays.asList("APP_TABLE", "TABLE"), + Arrays.asList("APP_VIEW", "VIEW"), + Arrays.asList("APP_PROC", "PROCEDURE"), + Arrays.asList("APP_FUNC", "FUNCTION") + ); + } + return Arrays.asList( + Arrays.asList("APP_TABLE", "TABLE", ""), + Arrays.asList("APP_VIEW", "VIEW", "") + ); + }); + } + + private static Connection metadataConnection(Function>> rowsForSql) { + return proxy(Connection.class, (unused, method, args) -> { + switch (method.getName()) { + case "prepareStatement": + return metadataStatement(rowsForSql.apply((String) args[0])); + case "close": + return null; + case "isClosed": + return false; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static PreparedStatement metadataStatement(List> rows) { + return proxy(PreparedStatement.class, (unused, method, args) -> { + switch (method.getName()) { + case "executeQuery": + return metadataResultSet(rows); + case "setString": + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet metadataResultSet(List> rows) { + int[] index = {-1}; + return proxy(ResultSet.class, (unused, method, args) -> { + switch (method.getName()) { + case "next": + index[0] += 1; + return index[0] < rows.size(); + case "getString": + return rows.get(index[0]).get(((Integer) args[0]) - 1); + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + @SuppressWarnings("unchecked") + private static T proxy(Class type, InvocationHandler handler) { + return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, handler); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) { + return false; + } + if (type == Byte.TYPE) { + return (byte) 0; + } + if (type == Short.TYPE) { + return (short) 0; + } + if (type == Integer.TYPE) { + return 0; + } + if (type == Long.TYPE) { + return 0L; + } + if (type == Float.TYPE) { + return 0F; + } + if (type == Double.TYPE) { + return 0.0D; + } + if (type == Character.TYPE) { + return '\0'; + } + return null; + } +} diff --git a/agents/drivers/oracle-go/.gitignore b/agents/drivers/oracle-go/.gitignore new file mode 100644 index 0000000000..8aa3c07867 --- /dev/null +++ b/agents/drivers/oracle-go/.gitignore @@ -0,0 +1,2 @@ +agent +agent.exe diff --git a/agents/drivers/oracle-go/README.md b/agents/drivers/oracle-go/README.md new file mode 100644 index 0000000000..847b78d5ab --- /dev/null +++ b/agents/drivers/oracle-go/README.md @@ -0,0 +1,36 @@ +# DBX Oracle Go Agent + +Experimental native Oracle agent for DBX using `github.com/sijms/go-ora/v2`. + +## Build + +```bash +go build -o agent . +``` + +Cross-compile release builds use pure Go output: + +```bash +CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o dbx-agent-oracle-linux-x64 . +CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 go build -trimpath -ldflags="-s -w" -o dbx-agent-oracle-macos-aarch64 . +CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o dbx-agent-oracle-windows-x64.exe . +``` + +## Local DBX Test + +Build the binary, then copy it into DBX's installed Oracle driver directory: + +```bash +mkdir -p ~/.dbx/agents/drivers/oracle +cp agent ~/.dbx/agents/drivers/oracle/agent +chmod +x ~/.dbx/agents/drivers/oracle/agent +``` + +DBX prefers `agent` over `agent.jar`, so Oracle connections will use this Go +agent until the file is removed. + +To restore the Java agent: + +```bash +rm ~/.dbx/agents/drivers/oracle/agent +``` diff --git a/agents/drivers/oracle-go/go.mod b/agents/drivers/oracle-go/go.mod new file mode 100644 index 0000000000..569ba62d6a --- /dev/null +++ b/agents/drivers/oracle-go/go.mod @@ -0,0 +1,5 @@ +module github.com/t8y2/dbx/agents/drivers/oracle-go + +go 1.22 + +require github.com/sijms/go-ora/v2 v2.9.0 diff --git a/agents/drivers/oracle-go/go.sum b/agents/drivers/oracle-go/go.sum new file mode 100644 index 0000000000..c8b6880459 --- /dev/null +++ b/agents/drivers/oracle-go/go.sum @@ -0,0 +1,2 @@ +github.com/sijms/go-ora/v2 v2.9.0 h1:+iQbUeTeCOFMb5BsOMgUhV8KWyrv9yjKpcK4x7+MFrg= +github.com/sijms/go-ora/v2 v2.9.0/go.mod h1:QgFInVi3ZWyqAiJwzBQA+nbKYKH77tdp1PYoCqhR2dU= diff --git a/agents/drivers/oracle-go/main.go b/agents/drivers/oracle-go/main.go new file mode 100644 index 0000000000..2db1c4fc7e --- /dev/null +++ b/agents/drivers/oracle-go/main.go @@ -0,0 +1,1400 @@ +package main + +import ( + "bufio" + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "io" + "net/url" + "os" + "regexp" + "strconv" + "strings" + "time" + + _ "github.com/sijms/go-ora/v2" + go_ora "github.com/sijms/go-ora/v2" +) + +const protocolVersion = 1 +const defaultMaxRows = 1000 +const oracleListDatabasesSQL = ` +SELECT username AS owner +FROM all_users +WHERE username IS NOT NULL + AND username NOT IN ( + 'SYS','SYSTEM','SYSMAN','DBSNMP','SYSBACKUP','SYSDG','SYSKM','SYSRAC','OUTLN', + 'AUDSYS','LBACSYS','DVF','DVSYS','APPQOSSYS','CTXSYS','MDSYS','MDDATA', + 'ORDSYS','ORDDATA','ORDPLUGINS','XDB','ANONYMOUS','DIP','EXFSYS', + 'GSMADMIN_INTERNAL','GSMCATUSER','GSMROOTUSER','GSMUSER','OJVMSYS','OLAPSYS', + 'ORACLE_OCM','SI_INFORMTN_SCHEMA','WMSYS','XS$NULL','DBSFWUSER', + 'REMOTE_SCHEDULER_AGENT','PDBADMIN','DGPDB_INT','OPS$ORACLE', + 'GGSYS','FLOWS_FILES','APEX_PUBLIC_USER' + ) + AND username NOT LIKE 'APEX_%' + AND username NOT LIKE 'FLOWS_%' + AND username NOT LIKE '%$%' +ORDER BY CASE + WHEN username = SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') THEN 0 + WHEN username = SYS_CONTEXT('USERENV', 'SESSION_USER') THEN 1 + ELSE 2 +END, username` + +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 rpcError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +type connectParams struct { + Host string `json:"host"` + Port int `json:"port"` + Database string `json:"database"` + Username string `json:"username"` + Password string `json:"password"` + SysDBA bool `json:"sysdba"` + URLParams string `json:"url_params"` + ConnectionString string `json:"connection_string"` +} + +type queryOptions struct { + SQL string `json:"sql"` + Database string `json:"database"` + Schema string `json:"schema"` + MaxRows int `json:"maxRows"` + FetchSize int `json:"fetchSize"` +} + +type queryResult struct { + Columns []string `json:"columns"` + Rows [][]any `json:"rows"` + AffectedRows int64 `json:"affected_rows"` + ExecutionTimeMS int64 `json:"execution_time_ms"` + Truncated bool `json:"truncated"` +} + +func (r queryResult) MarshalJSON() ([]byte, error) { + type alias queryResult + value := alias(r) + if value.Columns == nil { + value.Columns = []string{} + } + if value.Rows == nil { + value.Rows = [][]any{} + } + return json.Marshal(value) +} + +type queryPageResult struct { + Columns []string `json:"columns"` + 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"` +} + +func (r queryPageResult) MarshalJSON() ([]byte, error) { + type alias queryPageResult + value := alias(r) + if value.Columns == nil { + value.Columns = []string{} + } + if value.Rows == nil { + value.Rows = [][]any{} + } + return json.Marshal(value) +} + +type querySession struct { + rows *sql.Rows + columns []string + pending []any + remaining int +} + +type databaseInfo struct { + Name string `json:"name"` +} + +type tableInfo struct { + Name string `json:"name"` + TableType string `json:"table_type"` + Comment *string `json:"comment"` +} + +type objectInfo struct { + Name string `json:"name"` + ObjectType string `json:"object_type"` + Schema string `json:"schema"` + Comment *string `json:"comment"` +} + +type columnInfo struct { + Name string `json:"name"` + DataType string `json:"data_type"` + IsNullable bool `json:"is_nullable"` + ColumnDefault *string `json:"column_default"` + IsPrimaryKey bool `json:"is_primary_key"` + Extra *string `json:"extra"` + Comment *string `json:"comment"` + NumericPrecision *int `json:"numeric_precision"` + NumericScale *int `json:"numeric_scale"` + CharacterMaximumLength *int `json:"character_maximum_length"` +} + +type indexInfo struct { + Name string `json:"name"` + Columns []string `json:"columns"` + IsUnique bool `json:"is_unique"` + IsPrimary bool `json:"is_primary"` + Filter *string `json:"filter"` + IndexType *string `json:"index_type"` + IncludedColumns []string `json:"included_columns"` + Comment *string `json:"comment"` +} + +func (i indexInfo) MarshalJSON() ([]byte, error) { + type alias indexInfo + value := alias(i) + if value.Columns == nil { + value.Columns = []string{} + } + if value.IncludedColumns == nil { + value.IncludedColumns = []string{} + } + return json.Marshal(value) +} + +type foreignKeyInfo struct { + Name string `json:"name"` + Column string `json:"column"` + RefTable string `json:"ref_table"` + RefColumn string `json:"ref_column"` +} + +type triggerInfo struct { + Name string `json:"name"` + Event string `json:"event"` + Timing string `json:"timing"` +} + +type server struct { + db *sql.DB + params connectParams + sessions map[string]*querySession + nextSessionID int64 +} + +func main() { + s := newServer() + encoder := json.NewEncoder(os.Stdout) + 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 + } + resp, shutdown := s.handleLine(line) + if err := encoder.Encode(resp); err != nil { + fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err) + return + } + if shutdown { + return + } + } + if err := scanner.Err(); err != nil && !errors.Is(err, io.EOF) { + fmt.Fprintf(os.Stderr, "failed to read stdin: %v\n", err) + } +} + +func newServer() *server { + return &server{sessions: map[string]*querySession{}} +} + +func (s *server) 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 := s.dispatch(req.Method, req.Params) + if err != nil { + return errorResponse(req.ID, err), false + } + return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown +} + +func (s *server) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) { + switch method { + case "handshake": + return map[string]any{ + "protocolVersion": protocolVersion, + "agentProtocolVersion": protocolVersion, + "capabilities": []string{"connect", "test_connection", "metadata", "query", "ddl"}, + }, false, nil + case "connect": + var cp connectParams + if err := decodeParams(params, &cp); err != nil { + return nil, false, err + } + return map[string]bool{"ok": true}, false, s.connect(cp) + case "test_connection": + var cp connectParams + if err := decodeParams(params, &cp); err != nil { + return nil, false, err + } + db, err := openDB(cp) + if err != nil { + return nil, false, err + } + defer db.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := db.PingContext(ctx); err != nil { + return nil, false, err + } + return map[string]bool{"ok": true}, false, nil + 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": + schema := stringParam(params, "schema") + result, err := s.listTables(schema) + return result, false, err + case "list_objects": + schema := stringParam(params, "schema") + result, err := s.listObjects(schema) + return result, false, err + case "get_columns": + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.getColumns(schema, table) + return result, false, err + case "get_object_source": + schema := stringParam(params, "schema") + name := stringParam(params, "name") + objectType := stringParam(params, "object_type") + source, err := s.getObjectSource(schema, name, objectType) + return source, false, err + case "get_table_ddl": + schema := stringParam(params, "schema") + table := stringParam(params, "table") + objectType := stringParam(params, "object_type") + ddl, err := s.getTableDDL(schema, table, objectType) + return ddl, false, err + case "execute_query": + var opts queryOptions + if err := decodeParams(params, &opts); err != nil { + return nil, false, err + } + result, err := s.executeQuery(opts) + return result, false, err + case "execute_query_page": + var opts queryOptions + if err := decodeParams(params, &opts); err != nil { + return nil, false, err + } + result, err := s.executeQueryPage(opts, intParam(params, "pageSize")) + return result, false, err + case "fetch_query_page": + result, err := s.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize")) + return result, false, err + case "close_query_session": + return s.closeQuerySession(stringParam(params, "sessionId")), false, nil + case "list_indexes": + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.listIndexes(schema, table) + return result, false, err + case "list_foreign_keys": + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.listForeignKeys(schema, table) + return result, false, err + case "list_triggers": + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.listTriggers(schema, table) + return result, false, err + case "get_explain_info": + sqlText := stringParam(params, "sql") + plan, err := s.getExplainInfo(sqlText) + return map[string]any{"plan": plan, "has_actual_stats": false}, false, err + case "execute_transaction": + result, err := s.executeTransaction(params) + return result, false, err + case "disconnect": + return map[string]bool{"ok": true}, false, s.disconnect() + case "shutdown": + _ = s.disconnect() + return map[string]bool{"ok": true}, true, nil + default: + return nil, false, fmt.Errorf("unknown method: %s", method) + } +} + +func (s *server) connect(params connectParams) error { + _ = s.disconnect() + db, err := openDB(params) + if err != nil { + return err + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := db.PingContext(ctx); err != nil { + db.Close() + return err + } + if _, err := db.ExecContext(ctx, "ALTER SESSION SET NLS_LANGUAGE='AMERICAN'"); err != nil { + db.Close() + return err + } + s.db = db + s.params = params + return nil +} + +func (s *server) disconnect() error { + s.closeAllQuerySessions() + if s.db == nil { + return nil + } + err := s.db.Close() + s.db = nil + return err +} + +func openDB(params connectParams) (*sql.DB, error) { + dsn := buildDSN(params) + db, err := sql.Open("oracle", dsn) + if err != nil { + return nil, err + } + db.SetMaxOpenConns(4) + db.SetMaxIdleConns(1) + db.SetConnMaxLifetime(30 * time.Minute) + return db, nil +} + +func buildDSN(params connectParams) string { + connectionString := strings.TrimSpace(params.ConnectionString) + if strings.HasPrefix(strings.ToLower(connectionString), "oracle://") { + return connectionString + } + options := parseURLParams(params.URLParams) + if params.SysDBA { + options["AUTH TYPE"] = "SYSDBA" + } + + if jdbc := parseOracleJDBCURL(connectionString); jdbc.Kind != "" { + if jdbc.Descriptor != "" { + return go_ora.BuildJDBC(params.Username, params.Password, jdbc.Descriptor, options) + } + host := jdbc.Host + port := jdbc.Port + if port == 0 { + port = 1521 + } + if jdbc.Kind == "sid" { + options["SID"] = jdbc.Database + return go_ora.BuildUrl(host, port, "", params.Username, params.Password, options) + } + return go_ora.BuildUrl(host, port, jdbc.Database, params.Username, params.Password, options) + } + + service := strings.TrimSpace(params.Database) + if strings.HasPrefix(strings.ToUpper(service), "SYSDBA:") { + service = strings.TrimSpace(service[len("SYSDBA:"):]) + } + port := params.Port + if port == 0 { + port = 1521 + } + return go_ora.BuildUrl(params.Host, port, service, params.Username, params.Password, options) +} + +type jdbcURLInfo struct { + Kind string + Host string + Port int + Database string + Descriptor string +} + +var ( + oracleJDBCServiceRegexp = regexp.MustCompile(`(?i)^jdbc:oracle:thin:@//([^/:]+):([0-9]+)/([^?]+)`) + oracleJDBCSIDRegexp = regexp.MustCompile(`(?i)^jdbc:oracle:thin:@([^/:]+):([0-9]+):([^?]+)`) + oracleJDBCLegacyRegexp = regexp.MustCompile(`(?i)^jdbc:oracle:thin:@([^/:]+):([0-9]+)/([^?]+)`) +) + +func parseOracleJDBCURL(value string) jdbcURLInfo { + value = strings.TrimSpace(value) + lower := strings.ToLower(value) + if !strings.HasPrefix(lower, "jdbc:oracle:thin:@") { + return jdbcURLInfo{} + } + descriptor := strings.TrimSpace(value[len("jdbc:oracle:thin:@"):]) + if strings.HasPrefix(descriptor, "(") { + return jdbcURLInfo{Kind: "descriptor", Descriptor: descriptor} + } + if match := oracleJDBCServiceRegexp.FindStringSubmatch(value); len(match) == 4 { + return jdbcURLInfo{Kind: "service", Host: match[1], Port: parsePort(match[2]), Database: match[3]} + } + if match := oracleJDBCSIDRegexp.FindStringSubmatch(value); len(match) == 4 { + return jdbcURLInfo{Kind: "sid", Host: match[1], Port: parsePort(match[2]), Database: match[3]} + } + if match := oracleJDBCLegacyRegexp.FindStringSubmatch(value); len(match) == 4 { + return jdbcURLInfo{Kind: "service", Host: match[1], Port: parsePort(match[2]), Database: match[3]} + } + return jdbcURLInfo{} +} + +func parsePort(value string) int { + port, _ := strconv.Atoi(value) + return port +} + +func (s *server) requireDB() (*sql.DB, error) { + if s.db == nil { + return nil, errors.New("agent is not connected") + } + return s.db, nil +} + +func (s *server) listDatabases() ([]databaseInfo, error) { + rows, err := s.queryRows(oracleListDatabasesSQL, nil) + if err != nil { + if isOraclePGALimitError(err) { + return s.currentSchemaDatabase() + } + return nil, err + } + defer rows.Close() + var result []databaseInfo + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return nil, err + } + result = append(result, databaseInfo{Name: name}) + } + if err := rows.Err(); err != nil { + if isOraclePGALimitError(err) { + return s.currentSchemaDatabase() + } + return nil, err + } + return emptyIfNil(result), nil +} + +func isOraclePGALimitError(err error) bool { + return err != nil && strings.Contains(strings.ToUpper(err.Error()), "ORA-04036") +} + +func (s *server) currentSchemaDatabase() ([]databaseInfo, error) { + schema, err := s.currentSchema() + if err != nil { + return nil, err + } + if strings.TrimSpace(schema) == "" { + return []databaseInfo{}, nil + } + return []databaseInfo{{Name: schema}}, nil +} + +func (s *server) listSchemas() ([]string, error) { + databases, err := s.listDatabases() + if err != nil { + return nil, err + } + result := make([]string, 0, len(databases)) + for _, database := range databases { + result = append(result, database.Name) + } + return emptyIfNil(result), nil +} + +func (s *server) currentSchema() (string, error) { + db, err := s.requireDB() + if err != nil { + return "", err + } + var schema string + if err := db.QueryRow("SELECT SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') FROM DUAL").Scan(&schema); err != nil { + return "", err + } + return strings.ToUpper(schema), nil +} + +func (s *server) normalizeSchema(schema string) (string, error) { + schema = strings.TrimSpace(schema) + if schema == "" { + return s.currentSchema() + } + return strings.ToUpper(schema), nil +} + +func (s *server) listTables(schema string) ([]tableInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + rows, err := s.queryRows(` +SELECT /*+ NO_QUERY_TRANSFORMATION */ o.OBJECT_NAME, + CASE o.OBJECT_TYPE WHEN 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS TABLE_TYPE, + c.COMMENTS +FROM ALL_OBJECTS o +LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME +WHERE o.OWNER = :1 + AND o.OBJECT_TYPE IN ('TABLE', 'VIEW') +ORDER BY o.OBJECT_NAME`, []any{schema}) + if err != nil { + if isOraclePGALimitError(err) { + return []tableInfo{}, nil + } + return nil, err + } + defer rows.Close() + var result []tableInfo + for rows.Next() { + var item tableInfo + if err := rows.Scan(&item.Name, &item.TableType, &item.Comment); err != nil { + return nil, err + } + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) listObjects(schema string) ([]objectInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + rows, err := s.queryRows(` +SELECT /*+ NO_QUERY_TRANSFORMATION */ o.OBJECT_NAME, o.OBJECT_TYPE, c.COMMENTS +FROM ALL_OBJECTS o +LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME +WHERE o.OWNER = :1 + AND o.OBJECT_TYPE IN ('TABLE', 'VIEW', 'PROCEDURE', 'FUNCTION', 'PACKAGE') +ORDER BY o.OBJECT_TYPE, o.OBJECT_NAME`, []any{schema}) + if err != nil { + if isOraclePGALimitError(err) { + return []objectInfo{}, nil + } + return nil, err + } + defer rows.Close() + var result []objectInfo + for rows.Next() { + var item objectInfo + item.Schema = schema + if err := rows.Scan(&item.Name, &item.ObjectType, &item.Comment); err != nil { + return nil, err + } + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) getColumns(schema, table string) ([]columnInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT c.COLUMN_NAME, + c.DATA_TYPE, + c.NULLABLE, + c.DATA_DEFAULT, + CASE WHEN pk.COLUMN_NAME IS NULL THEN 0 ELSE 1 END AS IS_PRIMARY_KEY, + cc.COMMENTS, + c.DATA_PRECISION, + c.DATA_SCALE, + c.CHAR_LENGTH +FROM ALL_TAB_COLUMNS c +LEFT JOIN ( + SELECT acc.OWNER, acc.TABLE_NAME, acc.COLUMN_NAME + FROM ALL_CONSTRAINTS ac + JOIN ALL_CONS_COLUMNS acc ON acc.OWNER = ac.OWNER AND acc.CONSTRAINT_NAME = ac.CONSTRAINT_NAME + WHERE ac.CONSTRAINT_TYPE = 'P' +) pk ON pk.OWNER = c.OWNER AND pk.TABLE_NAME = c.TABLE_NAME AND pk.COLUMN_NAME = c.COLUMN_NAME +LEFT JOIN ALL_COL_COMMENTS cc ON cc.OWNER = c.OWNER AND cc.TABLE_NAME = c.TABLE_NAME AND cc.COLUMN_NAME = c.COLUMN_NAME +WHERE c.OWNER = :1 AND c.TABLE_NAME = :2 +ORDER BY c.COLUMN_ID`, []any{schema, table}) + if err != nil { + return nil, err + } + defer rows.Close() + var result []columnInfo + for rows.Next() { + var item columnInfo + var nullable string + var primary int + if err := rows.Scan( + &item.Name, + &item.DataType, + &nullable, + &item.ColumnDefault, + &primary, + &item.Comment, + &item.NumericPrecision, + &item.NumericScale, + &item.CharacterMaximumLength, + ); err != nil { + return nil, err + } + item.IsNullable = nullable == "Y" + item.IsPrimaryKey = primary != 0 + item.DataType = oracleColumnTypeDDL(item) + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) listIndexes(schema, table string) ([]indexInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT i.INDEX_NAME, + ic.COLUMN_NAME, + i.UNIQUENESS, + CASE WHEN pk.CONSTRAINT_NAME IS NULL THEN 0 ELSE 1 END AS IS_PRIMARY, + i.INDEX_TYPE, + ic.COLUMN_POSITION +FROM ALL_INDEXES i +JOIN ALL_IND_COLUMNS ic ON ic.INDEX_OWNER = i.OWNER AND ic.INDEX_NAME = i.INDEX_NAME +LEFT JOIN ALL_CONSTRAINTS pk ON pk.OWNER = i.TABLE_OWNER + AND pk.TABLE_NAME = i.TABLE_NAME + AND pk.CONSTRAINT_TYPE = 'P' + AND pk.INDEX_NAME = i.INDEX_NAME +WHERE i.TABLE_OWNER = :1 AND i.TABLE_NAME = :2 +ORDER BY i.INDEX_NAME, ic.COLUMN_POSITION`, []any{schema, table}) + if err != nil { + return nil, err + } + defer rows.Close() + + byName := map[string]*indexInfo{} + order := []string{} + for rows.Next() { + var name, column, uniqueness, indexType string + var primary int + var position int + if err := rows.Scan(&name, &column, &uniqueness, &primary, &indexType, &position); err != nil { + return nil, err + } + item := byName[name] + if item == nil { + item = &indexInfo{ + Name: name, + Columns: []string{}, + IsUnique: uniqueness == "UNIQUE", + IsPrimary: primary != 0, + IndexType: &indexType, + IncludedColumns: []string{}, + } + byName[name] = item + order = append(order, name) + } + item.Columns = append(item.Columns, column) + } + if err := rows.Err(); err != nil { + return nil, err + } + result := make([]indexInfo, 0, len(order)) + for _, name := range order { + result = append(result, *byName[name]) + } + return emptyIfNil(result), nil +} + +func (s *server) listForeignKeys(schema, table string) ([]foreignKeyInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT ac.CONSTRAINT_NAME, + acc.COLUMN_NAME, + rcc.TABLE_NAME AS REF_TABLE, + rcc.COLUMN_NAME AS REF_COLUMN +FROM ALL_CONSTRAINTS ac +JOIN ALL_CONS_COLUMNS acc ON acc.OWNER = ac.OWNER AND acc.CONSTRAINT_NAME = ac.CONSTRAINT_NAME +JOIN ALL_CONS_COLUMNS rcc ON rcc.OWNER = ac.R_OWNER AND rcc.CONSTRAINT_NAME = ac.R_CONSTRAINT_NAME + AND rcc.POSITION = acc.POSITION +WHERE ac.OWNER = :1 + AND ac.TABLE_NAME = :2 + AND ac.CONSTRAINT_TYPE = 'R' +ORDER BY ac.CONSTRAINT_NAME, acc.POSITION`, []any{schema, table}) + if err != nil { + return nil, err + } + defer rows.Close() + var result []foreignKeyInfo + for rows.Next() { + var item foreignKeyInfo + if err := rows.Scan(&item.Name, &item.Column, &item.RefTable, &item.RefColumn); err != nil { + return nil, err + } + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) listTriggers(schema, table string) ([]triggerInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT TRIGGER_NAME, TRIGGERING_EVENT, TRIGGER_TYPE +FROM ALL_TRIGGERS +WHERE OWNER = :1 AND TABLE_NAME = :2 +ORDER BY TRIGGER_NAME`, []any{schema, table}) + if err != nil { + return nil, err + } + defer rows.Close() + var result []triggerInfo + for rows.Next() { + var item triggerInfo + if err := rows.Scan(&item.Name, &item.Event, &item.Timing); err != nil { + return nil, err + } + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) getObjectSource(schema, name, objectType string) (map[string]any, error) { + var err error + schema, err = s.normalizeSchema(schema) + if err != nil { + return nil, err + } + upperType := strings.ToUpper(objectType) + if upperType == "VIEW" { + // ALL_VIEWS.TEXT for views — ALL_SOURCE doesn't contain views, and + // DBMS_METADATA.GET_DDL fails on XE editions. + var source string + err = s.db.QueryRow( + "SELECT TEXT FROM ALL_VIEWS WHERE OWNER = :1 AND VIEW_NAME = :2", + schema, strings.ToUpper(name), + ).Scan(&source) + if errors.Is(err, sql.ErrNoRows) { + return map[string]any{"name": name, "object_type": objectType, "schema": schema, "source": ""}, nil + } + if err != nil { + return nil, err + } + return map[string]any{"name": name, "object_type": objectType, "schema": schema, "source": source}, nil + } + + rows, err := s.queryRows(` +SELECT TEXT +FROM ALL_SOURCE +WHERE OWNER = :1 AND NAME = :2 AND TYPE = :3 +ORDER BY LINE`, []any{schema, strings.ToUpper(name), upperType}) + if err != nil { + return nil, err + } + defer rows.Close() + var builder strings.Builder + for rows.Next() { + var line string + if err := rows.Scan(&line); err != nil { + return nil, err + } + builder.WriteString(line) + } + return map[string]any{"name": name, "object_type": objectType, "schema": schema, "source": builder.String()}, rows.Err() +} + +func (s *server) getTableDDL(schema, table, objectType string) (string, error) { + var err error + schema, err = s.normalizeSchema(schema) + if err != nil { + return "", err + } + db, err := s.requireDB() + if err != nil { + return "", err + } + objectType, err = s.resolveDDLObjectType(schema, table, objectType) + if err != nil { + return "", err + } + if objectType == "VIEW" { + return s.buildViewDDL(schema, table) + } + var ddl string + err = db.QueryRow("SELECT DBMS_METADATA.GET_DDL(:1, :2, :3) FROM DUAL", objectType, strings.ToUpper(table), schema).Scan(&ddl) + if err == nil && strings.TrimSpace(ddl) != "" { + return ddl, nil + } + if objectType == "TABLE" { + return s.buildTableDDL(schema, table) + } + return "", err +} + +func (s *server) resolveDDLObjectType(schema, name, requested string) (string, error) { + objectType := normalizeDDLObjectType(requested) + if objectType != "" { + return objectType, nil + } + db, err := s.requireDB() + if err != nil { + return "", err + } + err = db.QueryRow(` +SELECT OBJECT_TYPE +FROM ( + SELECT OBJECT_TYPE + FROM ALL_OBJECTS + WHERE OWNER = :1 + AND OBJECT_NAME = :2 + AND OBJECT_TYPE IN ('TABLE', 'VIEW', 'MATERIALIZED VIEW') + ORDER BY CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 ELSE 2 END +) +WHERE ROWNUM = 1`, schema, strings.ToUpper(name)).Scan(&objectType) + if errors.Is(err, sql.ErrNoRows) { + return "", fmt.Errorf("object not found: %s.%s", schema, name) + } + if err != nil { + return "", err + } + return normalizeDDLObjectType(objectType), nil +} + +func normalizeDDLObjectType(value string) string { + switch strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(value), " ", "_")) { + case "TABLE": + return "TABLE" + case "VIEW": + return "VIEW" + case "MATERIALIZED_VIEW": + return "MATERIALIZED_VIEW" + default: + return "" + } +} + +func (s *server) buildViewDDL(schema, name string) (string, error) { + db, err := s.requireDB() + if err != nil { + return "", err + } + var source string + err = db.QueryRow( + "SELECT TEXT FROM ALL_VIEWS WHERE OWNER = :1 AND VIEW_NAME = :2", + schema, strings.ToUpper(name), + ).Scan(&source) + if errors.Is(err, sql.ErrNoRows) { + return "", fmt.Errorf("view not found: %s.%s", schema, name) + } + if err != nil { + return "", err + } + return fmt.Sprintf("CREATE OR REPLACE VIEW %s.%s AS\n%s", quoteIdentifier(schema), quoteIdentifier(name), strings.TrimSpace(source)), nil +} + +func (s *server) buildTableDDL(schema, table string) (string, error) { + columns, err := s.getColumns(schema, table) + if err != nil { + return "", err + } + if len(columns) == 0 { + return "", fmt.Errorf("table not found: %s.%s", schema, table) + } + var builder strings.Builder + builder.WriteString("CREATE TABLE ") + builder.WriteString(quoteIdentifier(schema)) + builder.WriteByte('.') + builder.WriteString(quoteIdentifier(table)) + builder.WriteString(" (\n") + for i, column := range columns { + if i > 0 { + builder.WriteString(",\n") + } + builder.WriteString(" ") + builder.WriteString(quoteIdentifier(column.Name)) + builder.WriteByte(' ') + builder.WriteString(oracleColumnTypeDDL(column)) + if column.ColumnDefault != nil && strings.TrimSpace(*column.ColumnDefault) != "" { + builder.WriteString(" DEFAULT ") + builder.WriteString(strings.TrimSpace(*column.ColumnDefault)) + } + if !column.IsNullable { + builder.WriteString(" NOT NULL") + } + } + primary := make([]string, 0) + for _, column := range columns { + if column.IsPrimaryKey { + primary = append(primary, quoteIdentifier(column.Name)) + } + } + if len(primary) > 0 { + builder.WriteString(",\n PRIMARY KEY (") + builder.WriteString(strings.Join(primary, ", ")) + builder.WriteByte(')') + } + builder.WriteString("\n)") + return builder.String(), nil +} + +func oracleColumnTypeDDL(column columnInfo) string { + dataType := strings.ToUpper(strings.TrimSpace(column.DataType)) + if dataType == "" { + return "VARCHAR2(4000)" + } + if strings.Contains(dataType, "(") { + return dataType + } + if isOracleCharacterType(dataType) && column.CharacterMaximumLength != nil && *column.CharacterMaximumLength > 0 { + return fmt.Sprintf("%s(%d)", dataType, *column.CharacterMaximumLength) + } + if dataType == "NUMBER" { + if column.NumericPrecision != nil && *column.NumericPrecision > 0 { + if column.NumericScale != nil && *column.NumericScale != 0 { + return fmt.Sprintf("NUMBER(%d,%d)", *column.NumericPrecision, *column.NumericScale) + } + return fmt.Sprintf("NUMBER(%d)", *column.NumericPrecision) + } + return "NUMBER" + } + if (dataType == "FLOAT" || dataType == "BINARY_FLOAT" || dataType == "BINARY_DOUBLE") && + column.NumericPrecision != nil && *column.NumericPrecision > 0 { + return fmt.Sprintf("%s(%d)", dataType, *column.NumericPrecision) + } + return dataType +} + +func isOracleCharacterType(dataType string) bool { + switch dataType { + case "CHAR", "VARCHAR2", "VARCHAR", "NCHAR", "NVARCHAR2", "RAW": + return true + default: + return false + } +} + +func (s *server) getExplainInfo(sqlText string) (string, error) { + if strings.TrimSpace(sqlText) == "" { + return "", errors.New("sql is required") + } + rows, err := s.queryRows("EXPLAIN PLAN FOR "+trimStatementSQL(sqlText), nil) + if err != nil { + return "", err + } + rows.Close() + planRows, err := s.queryRows("SELECT PLAN_TABLE_OUTPUT FROM TABLE(DBMS_XPLAN.DISPLAY())", nil) + if err != nil { + return "", err + } + defer planRows.Close() + var builder strings.Builder + for planRows.Next() { + var line string + if err := planRows.Scan(&line); err != nil { + return "", err + } + builder.WriteString(line) + builder.WriteByte('\n') + } + return strings.TrimSpace(builder.String()), planRows.Err() +} + +func (s *server) executeTransaction(params map[string]json.RawMessage) (queryResult, error) { + var payload struct { + Statements []string `json:"statements"` + Schema string `json:"schema"` + } + if err := decodeParams(params, &payload); err != nil { + return queryResult{}, err + } + db, err := s.requireDB() + if err != nil { + return queryResult{}, err + } + tx, err := db.Begin() + if err != nil { + return queryResult{}, err + } + if strings.TrimSpace(payload.Schema) != "" { + if _, err := tx.Exec("ALTER SESSION SET CURRENT_SCHEMA = " + quoteIdentifier(payload.Schema)); err != nil { + tx.Rollback() + return queryResult{}, err + } + } + var affected int64 + start := time.Now() + for _, statement := range payload.Statements { + statement = trimStatementSQL(statement) + if statement == "" { + continue + } + result, err := tx.Exec(statement) + if err != nil { + tx.Rollback() + return queryResult{}, err + } + count, _ := result.RowsAffected() + affected += count + } + if err := tx.Commit(); err != nil { + return queryResult{}, err + } + return queryResult{ + Columns: []string{}, + Rows: [][]any{}, + AffectedRows: affected, + ExecutionTimeMS: time.Since(start).Milliseconds(), + }, nil +} + +func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageResult, error) { + start := time.Now() + if strings.TrimSpace(opts.Schema) != "" { + if err := s.setSchema(opts.Schema); err != nil { + return queryPageResult{}, err + } + } + sqlText := trimStatementSQL(opts.SQL) + if !isQuerySQL(sqlText) { + result, err := s.executeQuery(opts) + return queryPageResult{ + Columns: result.Columns, + Rows: result.Rows, + AffectedRows: result.AffectedRows, + ExecutionTimeMS: result.ExecutionTimeMS, + Truncated: result.Truncated, + SessionID: nil, + HasMore: false, + }, err + } + rows, err := s.queryRows(sqlText, nil) + if err != nil { + return queryPageResult{}, err + } + columns, err := rows.Columns() + if err != nil { + rows.Close() + return queryPageResult{}, err + } + maxRows := opts.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + session := &querySession{rows: rows, columns: columns, remaining: maxRows} + result, err := readQuerySessionPage(session, pageSize) + result.ExecutionTimeMS = time.Since(start).Milliseconds() + if err != nil { + rows.Close() + return queryPageResult{}, err + } + if result.HasMore { + sessionID := s.storeQuerySession(session) + result.SessionID = &sessionID + } else { + rows.Close() + } + return result, nil +} + +func (s *server) fetchQueryPage(sessionID string, pageSize int) (queryPageResult, error) { + session := s.sessions[sessionID] + if session == nil { + return queryPageResult{Columns: []string{}, Rows: [][]any{}, SessionID: nil, HasMore: false}, nil + } + result, err := readQuerySessionPage(session, pageSize) + if err != nil { + s.closeQuerySession(sessionID) + return queryPageResult{}, err + } + if result.HasMore { + result.SessionID = &sessionID + } else { + s.closeQuerySession(sessionID) + } + return result, nil +} + +func (s *server) storeQuerySession(session *querySession) string { + s.nextSessionID++ + sessionID := fmt.Sprintf("oracle-go-%d", s.nextSessionID) + s.sessions[sessionID] = session + return sessionID +} + +func (s *server) closeQuerySession(sessionID string) bool { + session := s.sessions[sessionID] + if session == nil { + return false + } + session.rows.Close() + delete(s.sessions, sessionID) + return true +} + +func (s *server) closeAllQuerySessions() { + for sessionID := range s.sessions { + s.closeQuerySession(sessionID) + } +} + +func readQuerySessionPage(session *querySession, pageSize int) (queryPageResult, error) { + if pageSize <= 0 { + pageSize = defaultMaxRows + } + result := queryPageResult{Columns: session.columns, Rows: [][]any{}, SessionID: nil, HasMore: false} + for len(result.Rows) < pageSize && session.remaining > 0 { + if session.pending != nil { + result.Rows = append(result.Rows, session.pending) + session.pending = nil + session.remaining-- + continue + } + if !session.rows.Next() { + return result, session.rows.Err() + } + row, err := scanRow(session.rows, len(session.columns)) + if err != nil { + return queryPageResult{}, err + } + result.Rows = append(result.Rows, row) + session.remaining-- + } + if session.remaining <= 0 { + result.Truncated = true + return result, nil + } + if session.rows.Next() { + row, err := scanRow(session.rows, len(session.columns)) + if err != nil { + return queryPageResult{}, err + } + session.pending = row + result.HasMore = true + return result, nil + } + return result, session.rows.Err() +} + +func (s *server) executeQuery(opts queryOptions) (queryResult, error) { + start := time.Now() + if strings.TrimSpace(opts.Schema) != "" { + if err := s.setSchema(opts.Schema); err != nil { + return queryResult{}, err + } + } + sqlText := trimStatementSQL(opts.SQL) + maxRows := opts.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + if isQuerySQL(sqlText) { + result, err := s.executeSelect(sqlText, maxRows) + result.ExecutionTimeMS = time.Since(start).Milliseconds() + return result, err + } + db, err := s.requireDB() + if err != nil { + return queryResult{}, err + } + execResult, err := db.Exec(sqlText) + if err != nil { + return queryResult{}, err + } + affected, _ := execResult.RowsAffected() + return queryResult{Columns: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil +} + +func (s *server) executeSelect(sqlText string, maxRows int) (queryResult, error) { + rows, err := s.queryRows(sqlText, nil) + if err != nil { + return queryResult{}, err + } + defer rows.Close() + columns, err := rows.Columns() + if err != nil { + return queryResult{}, err + } + result := queryResult{Columns: columns, Rows: [][]any{}} + for rows.Next() { + if len(result.Rows) >= maxRows { + result.Truncated = true + break + } + values, err := scanRow(rows, len(columns)) + if err != nil { + return queryResult{}, err + } + result.Rows = append(result.Rows, values) + } + return result, rows.Err() +} + +func scanRow(rows *sql.Rows, columnCount int) ([]any, error) { + values := make([]any, columnCount) + scanTargets := make([]any, columnCount) + for i := range values { + scanTargets[i] = &values[i] + } + if err := rows.Scan(scanTargets...); err != nil { + return nil, err + } + for i, value := range values { + values[i] = normalizeValue(value) + } + return values, nil +} + +func (s *server) setSchema(schema string) error { + db, err := s.requireDB() + if err != nil { + return err + } + _, err = db.Exec("ALTER SESSION SET CURRENT_SCHEMA = " + quoteIdentifier(schema)) + return err +} + +func (s *server) queryRows(sqlText string, args []any) (*sql.Rows, error) { + db, err := s.requireDB() + if err != nil { + return nil, err + } + if len(args) == 0 { + return db.Query(sqlText) + } + return db.Query(sqlText, args...) +} + +func decodeParams(params map[string]json.RawMessage, target any) error { + if params == nil { + params = map[string]json.RawMessage{} + } + 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 params == nil || len(params[key]) == 0 { + return "" + } + var value string + _ = json.Unmarshal(params[key], &value) + return value +} + +func intParam(params map[string]json.RawMessage, key string) int { + if params == nil || len(params[key]) == 0 { + return 0 + } + var value int + _ = json.Unmarshal(params[key], &value) + return value +} + +func errorResponse(id json.RawMessage, err error) response { + return response{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: -1, Message: err.Error()}} +} + +func parseURLParams(raw string) map[string]string { + result := map[string]string{} + values, err := url.ParseQuery(raw) + if err != nil { + return result + } + for key, items := range values { + if len(items) > 0 { + result[key] = items[len(items)-1] + } + } + return result +} + +func trimStatementSQL(sqlText string) string { + return strings.TrimRight(strings.TrimSpace(sqlText), "; \t\r\n") +} + +func isQuerySQL(sqlText string) bool { + lower := strings.ToLower(strings.TrimSpace(sqlText)) + return strings.HasPrefix(lower, "select") || strings.HasPrefix(lower, "with") +} + +func quoteIdentifier(value string) string { + return `"` + strings.ReplaceAll(value, `"`, `""`) + `"` +} + +func normalizeValue(value any) any { + switch v := value.(type) { + case nil: + return nil + case []byte: + return string(v) + case time.Time: + return v.Format(time.RFC3339Nano) + case int64, float64, bool, string: + return v + case fmt.Stringer: + return v.String() + default: + return fmt.Sprint(v) + } +} + +func emptyIfNil[T any](values []T) []T { + if values == nil { + return []T{} + } + return values +} + +func intPtrFromString(value string) *int { + if value == "" { + return nil + } + parsed, err := strconv.Atoi(value) + if err != nil { + return nil + } + return &parsed +} diff --git a/agents/drivers/oracle-go/main_test.go b/agents/drivers/oracle-go/main_test.go new file mode 100644 index 0000000000..3830f92963 --- /dev/null +++ b/agents/drivers/oracle-go/main_test.go @@ -0,0 +1,284 @@ +package main + +import ( + "encoding/json" + "errors" + "os" + "strings" + "testing" +) + +func TestHandshakeResponse(t *testing.T) { + s := newServer() + resp, shutdown := s.handleLine(`{"jsonrpc":"2.0","id":7,"method":"handshake","params":{"appVersion":"dev"}}`) + if shutdown { + t.Fatal("handshake should not shut down the server") + } + if resp.Error != nil { + t.Fatalf("unexpected error: %v", resp.Error) + } + data, err := json.Marshal(resp.Result) + if err != nil { + t.Fatal(err) + } + var result struct { + ProtocolVersion int `json:"protocolVersion"` + AgentProtocolVersion int `json:"agentProtocolVersion"` + Capabilities []string `json:"capabilities"` + } + if err := json.Unmarshal(data, &result); err != nil { + t.Fatal(err) + } + if result.ProtocolVersion != 1 || result.AgentProtocolVersion != 1 { + t.Fatalf("unexpected protocol versions: %+v", result) + } + contract := protocolContract(t) + if result.ProtocolVersion != contract.ProtocolVersion || result.AgentProtocolVersion != contract.ProtocolVersion { + t.Fatalf("handshake protocol versions do not match contract: result=%+v contract=%+v", result, contract) + } + for _, capability := range result.Capabilities { + if !contains(contract.AllCapabilities, capability) { + t.Fatalf("handshake returned capability %q outside protocol contract %v", capability, contract.AllCapabilities) + } + } + if !contains(result.Capabilities, "query") || !contains(result.Capabilities, "metadata") { + t.Fatalf("expected query and metadata capabilities, got %v", result.Capabilities) + } +} + +func TestCloseMissingQuerySessionReturnsFalse(t *testing.T) { + s := newServer() + resp, shutdown := s.handleLine(`{"jsonrpc":"2.0","id":8,"method":"close_query_session","params":{"sessionId":"missing"}}`) + if shutdown { + t.Fatal("close_query_session should not shut down the server") + } + if resp.Error != nil { + t.Fatalf("unexpected error: %v", resp.Error) + } + if resp.Result != false { + t.Fatalf("expected false result, got %#v", resp.Result) + } +} + +func TestEmptyResultSlicesMarshalAsArrays(t *testing.T) { + data, err := json.Marshal(queryResult{}) + if err != nil { + t.Fatal(err) + } + text := string(data) + if strings.Contains(text, `"columns":null`) || strings.Contains(text, `"rows":null`) { + t.Fatalf("query result should marshal nil slices as arrays: %s", text) + } + + data, err = json.Marshal(indexInfo{}) + if err != nil { + t.Fatal(err) + } + text = string(data) + if strings.Contains(text, `"columns":null`) || strings.Contains(text, `"included_columns":null`) { + t.Fatalf("index info should marshal nil slices as arrays: %s", text) + } +} + +func TestGetTableDDLResultMarshalsAsString(t *testing.T) { + data, err := json.Marshal("CREATE TABLE HR.ORDERS (ID NUMBER)") + if err != nil { + t.Fatal(err) + } + var ddl string + if err := json.Unmarshal(data, &ddl); err != nil { + t.Fatalf("get_table_ddl result must deserialize as a string: %v", err) + } +} + +func TestNormalizeDDLObjectType(t *testing.T) { + tests := map[string]string{ + "": "", + "table": "TABLE", + "VIEW": "VIEW", + "materialized view": "MATERIALIZED_VIEW", + "MATERIALIZED_VIEW": "MATERIALIZED_VIEW", + "procedure": "", + } + for input, want := range tests { + if got := normalizeDDLObjectType(input); got != want { + t.Fatalf("normalizeDDLObjectType(%q) = %q, want %q", input, got, want) + } + } +} + +func protocolContract(t *testing.T) struct { + ProtocolVersion int `json:"protocolVersion"` + AllCapabilities []string `json:"allCapabilities"` +} { + t.Helper() + data, err := os.ReadFile("../../common/src/main/resources/agent-protocol-v1.json") + if err != nil { + t.Fatal(err) + } + var contract struct { + ProtocolVersion int `json:"protocolVersion"` + AllCapabilities []string `json:"allCapabilities"` + } + if err := json.Unmarshal(data, &contract); err != nil { + t.Fatal(err) + } + return contract +} + +func TestOracleColumnTypeDDL(t *testing.T) { + charLen := 64 + precision := 10 + scale := 2 + zeroScale := 0 + + tests := []struct { + name string + column columnInfo + want string + }{ + {name: "varchar", column: columnInfo{DataType: "VARCHAR2", CharacterMaximumLength: &charLen}, want: "VARCHAR2(64)"}, + {name: "number scale", column: columnInfo{DataType: "NUMBER", NumericPrecision: &precision, NumericScale: &scale}, want: "NUMBER(10,2)"}, + {name: "number zero scale", column: columnInfo{DataType: "NUMBER", NumericPrecision: &precision, NumericScale: &zeroScale}, want: "NUMBER(10)"}, + {name: "timestamp preserves precision", column: columnInfo{DataType: "TIMESTAMP(6)"}, want: "TIMESTAMP(6)"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := oracleColumnTypeDDL(tt.column); got != tt.want { + t.Fatalf("oracleColumnTypeDDL() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestBuildDSNUsesConnectionStringWhenProvided(t *testing.T) { + dsn := buildDSN(connectParams{ConnectionString: "oracle://scott:tiger@db.example.com:1521/ORCLPDB1"}) + + if dsn != "oracle://scott:tiger@db.example.com:1521/ORCLPDB1" { + t.Fatalf("unexpected dsn: %s", dsn) + } +} + +func TestBuildDSNUsesJdbcServiceHostAndPort(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "127.0.0.1", + Port: 11521, + Database: "ORCLPDB1", + Username: "scott", + Password: "tiger", + ConnectionString: "jdbc:oracle:thin:@//oracle.example.com:1521/ORCLPDB1", + }) + + if strings.Contains(strings.ToLower(dsn), "jdbc:") { + t.Fatalf("dsn should be go-ora format, got: %s", dsn) + } + if !strings.Contains(dsn, "oracle.example.com:1521") || !strings.Contains(dsn, "ORCLPDB1") { + t.Fatalf("dsn should use JDBC host/port/database fields, got: %s", dsn) + } +} + +func TestBuildDSNUsesRewrittenJdbcServiceHostAndPort(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "127.0.0.1", + Port: 11521, + Database: "ORCLPDB1", + Username: "scott", + Password: "tiger", + ConnectionString: "jdbc:oracle:thin:@//127.0.0.1:11521/ORCLPDB1", + }) + + if strings.Contains(strings.ToLower(dsn), "jdbc:") { + t.Fatalf("dsn should be go-ora format, got: %s", dsn) + } + if !strings.Contains(dsn, "127.0.0.1:11521") || !strings.Contains(dsn, "ORCLPDB1") { + t.Fatalf("dsn should use rewritten JDBC host/port/database fields, got: %s", dsn) + } +} + +func TestBuildDSNConvertsJdbcSID(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "127.0.0.1", + Port: 11521, + Database: "ORCL", + Username: "scott", + Password: "tiger", + ConnectionString: "jdbc:oracle:thin:@oracle.example.com:1521:ORCL", + }) + + if strings.Contains(strings.ToLower(dsn), "jdbc:") { + t.Fatalf("dsn should be go-ora format, got: %s", dsn) + } + upperDSN := strings.ToUpper(dsn) + if !strings.Contains(dsn, "oracle.example.com:1521") || !strings.Contains(upperDSN, "SID=ORCL") { + t.Fatalf("dsn should use JDBC host/port and SID option, got: %s", dsn) + } +} + +func TestBuildDSNConvertsJdbcDescriptor(t *testing.T) { + dsn := buildDSN(connectParams{ + Username: "scott", + Password: "tiger", + ConnectionString: "jdbc:oracle:thin:@(DESCRIPTION=(ADDRESS=(PROTOCOL=TCP)(HOST=db.example.com)(PORT=1521))(CONNECT_DATA=(SERVICE_NAME=ORCLPDB1)))", + }) + + if !strings.HasPrefix(dsn, "oracle://scott:tiger@") { + t.Fatalf("descriptor should become go-ora url, got: %s", dsn) + } + if !strings.Contains(dsn, "connStr=") { + t.Fatalf("descriptor should be passed via connStr option, got: %s", dsn) + } +} + +func TestBuildDSNAddsSysDbaOption(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "127.0.0.1", + Port: 1521, + Database: "SYSDBA:ORCLPDB1", + Username: "sys", + Password: "secret", + SysDBA: true, + URLParams: "TRACE FILE=trace.log", + }) + + if strings.Contains(dsn, "SYSDBA:") { + t.Fatalf("dsn should strip SYSDBA prefix: %s", dsn) + } + if !strings.Contains(dsn, "ORCLPDB1") { + t.Fatalf("dsn should include service name: %s", dsn) + } + upperDSN := strings.ToUpper(dsn) + if !strings.Contains(upperDSN, "AUTH TYPE=SYSDBA") && + !strings.Contains(upperDSN, "AUTH+TYPE=SYSDBA") && + !strings.Contains(upperDSN, "AUTH%20TYPE=SYSDBA") { + t.Fatalf("dsn should include SYSDBA auth option: %s", dsn) + } +} + +func TestListDatabasesSQLUsesUserDictionaryInsteadOfObjectDictionary(t *testing.T) { + sqlText := strings.ToUpper(oracleListDatabasesSQL) + + if !strings.Contains(sqlText, "ALL_USERS") { + t.Fatalf("schema listing should query ALL_USERS, got: %s", oracleListDatabasesSQL) + } + if strings.Contains(sqlText, "ALL_TABLES") || strings.Contains(sqlText, "ALL_VIEWS") { + t.Fatalf("schema listing should not scan object dictionaries, got: %s", oracleListDatabasesSQL) + } +} + +func TestIsOraclePGALimitError(t *testing.T) { + if !isOraclePGALimitError(errors.New("ORA-04036: PGA memory used by the instance exceeds PGA_AGGREGATE_LIMIT")) { + t.Fatal("expected ORA-04036 to be detected") + } + if isOraclePGALimitError(errors.New("ORA-00942: table or view does not exist")) { + t.Fatal("unexpected ORA-00942 match") + } +} + +func contains(values []string, target string) bool { + for _, value := range values { + if value == target { + return true + } + } + return false +} diff --git a/agents/drivers/oracle-legacy/build.gradle b/agents/drivers/oracle-legacy/build.gradle new file mode 100644 index 0000000000..f43773d5e8 --- /dev/null +++ b/agents/drivers/oracle-legacy/build.gradle @@ -0,0 +1,19 @@ +sourceSets { + main { + java { + srcDirs = ['src/main/java', '../oracle/src/main/java'] + } + } +} + +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.oracle.database.jdbc:ojdbc8:19.26.0.0' + implementation 'com.oracle.database.nls:orai18n:19.26.0.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Oracle Legacy', 'Main-Class': 'com.dbx.agent.oraclelegacy.OracleLegacyAgent') + } +} diff --git a/agents/drivers/oracle-legacy/src/main/java/com/dbx/agent/oraclelegacy/OracleLegacyAgent.java b/agents/drivers/oracle-legacy/src/main/java/com/dbx/agent/oraclelegacy/OracleLegacyAgent.java new file mode 100644 index 0000000000..d2fc003ea6 --- /dev/null +++ b/agents/drivers/oracle-legacy/src/main/java/com/dbx/agent/oraclelegacy/OracleLegacyAgent.java @@ -0,0 +1,10 @@ +package com.dbx.agent.oraclelegacy; + +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.oracle.OracleAgent; + +public final class OracleLegacyAgent extends OracleAgent { + public static void main(String[] args) { + new JsonRpcServer(new OracleLegacyAgent()).run(); + } +} diff --git a/agents/drivers/oracle/build.gradle b/agents/drivers/oracle/build.gradle new file mode 100644 index 0000000000..32455273d2 --- /dev/null +++ b/agents/drivers/oracle/build.gradle @@ -0,0 +1,11 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.oracle.database.jdbc:ojdbc17:23.26.2.0.0' + implementation 'com.oracle.database.nls:orai18n:23.26.2.0.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Oracle', 'Main-Class': 'com.dbx.agent.oracle.OracleAgent') + } +} diff --git a/agents/drivers/oracle/src/main/java/com/dbx/agent/oracle/OracleAgent.java b/agents/drivers/oracle/src/main/java/com/dbx/agent/oracle/OracleAgent.java new file mode 100644 index 0000000000..c495dcbcdc --- /dev/null +++ b/agents/drivers/oracle/src/main/java/com/dbx/agent/oracle/OracleAgent.java @@ -0,0 +1,616 @@ +package com.dbx.agent.oracle; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryPageOptions; +import com.dbx.agent.QueryPageResult; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Properties; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import java.util.stream.Collectors; + +public class OracleAgent extends BaseDatabaseAgent { + private static final Set SYSTEM_SCHEMAS = Set.of( + "SYS", "SYSTEM", "SYSMAN", "DBSNMP", "SYSBACKUP", "SYSDG", "SYSKM", "OUTLN", + "AUDSYS", "LBACSYS", "DVF", "DVSYS", "APPQOSSYS", "CTXSYS", "MDSYS", "MDDATA", + "ORDSYS", "ORDDATA", "ORDPLUGINS", "XDB", "ANONYMOUS", "DIP", "EXFSYS", + "GSMADMIN_INTERNAL", "GSMCATUSER", "GSMUSER", "OJVMSYS", "OLAPSYS", + "ORACLE_OCM", "SI_INFORMTN_SCHEMA", "WMSYS", "XS$NULL", "DBSFWUSER", + "REMOTE_SCHEDULER_AGENT", "PDBADMIN", "DGPDB_INT", "OPS$ORACLE", + "GGSYS", "FLOWS_FILES", "APEX_PUBLIC_USER", "GSMROOTUSER", "SYSRAC" + ); + + private static final Pattern OFFSET_FETCH_RE = Pattern.compile( + "(.+?)\\s+OFFSET\\s+(\\d+)\\s+ROWS?\\s+FETCH\\s+(FIRST|NEXT)\\s+(\\d+)\\s+ROWS?\\s+ONLY", + Pattern.CASE_INSENSITIVE + ); + private static final Pattern FETCH_ONLY_RE = Pattern.compile( + "(.+?)\\s+FETCH\\s+(FIRST|NEXT)\\s+(\\d+)\\s+ROWS?\\s+ONLY", + Pattern.CASE_INSENSITIVE + ); + private static final Pattern PLSQL_OBJECT_DDL_RE = Pattern.compile( + "^\\s*CREATE\\s+(?:OR\\s+REPLACE\\s+)?(?:(?:NON)?EDITIONABLE\\s+)?(?:PROCEDURE|FUNCTION|PACKAGE(?:\\s+BODY)?|TRIGGER|TYPE(?:\\s+BODY)?)\\b", + Pattern.CASE_INSENSITIVE + ); + + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + public static String rewriteFetchFirst(String sql) { + Matcher offsetMatcher = OFFSET_FETCH_RE.matcher(sql); + if (offsetMatcher.matches()) { + String innerSql = offsetMatcher.group(1); + long offset = Long.parseLong(offsetMatcher.group(2)); + long limit = Long.parseLong(offsetMatcher.group(4)); + long upper = offset + limit; + return "SELECT * FROM (SELECT a.*, ROWNUM rn__ FROM (" + innerSql + ") a WHERE ROWNUM <= " + upper + + ") WHERE rn__ > " + offset; + } + + Matcher fetchMatcher = FETCH_ONLY_RE.matcher(sql); + if (fetchMatcher.matches()) { + String innerSql = fetchMatcher.group(1); + long limit = Long.parseLong(fetchMatcher.group(3)); + return "SELECT * FROM (" + innerSql + ") WHERE ROWNUM <= " + limit; + } + + return sql; + } + + @Override + public String setSchemaSQL(String schema) { + return "ALTER SESSION SET CURRENT_SCHEMA = " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("oracle.jdbc.OracleDriver"); + String serviceName = params.getDatabase(); + Properties props = connectionProperties(params); + + if (serviceName.toUpperCase(Locale.ROOT).startsWith("SYSDBA:")) { + serviceName = serviceName.substring(7); + props.setProperty("internal_logon", "SYSDBA"); + } + + connection = DriverManager.getConnection(buildUrl(params, serviceName), props); + try (var stmt = connection.createStatement()) { + stmt.execute("ALTER SESSION SET NLS_LANGUAGE='AMERICAN'"); + } + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("oracle.jdbc.OracleDriver"); + String serviceName = params.getDatabase(); + Properties props = connectionProperties(params); + + if (serviceName.toUpperCase(Locale.ROOT).startsWith("SYSDBA:")) { + serviceName = serviceName.substring(7); + props.setProperty("internal_logon", "SYSDBA"); + } + + try (Connection conn = DriverManager.getConnection(buildUrl(params, serviceName), props)) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + try { + return queryDatabaseInfos(listDatabasesSql()); + } catch (SQLException e) { + if (isPgaLimitError(e)) { + return currentSchemaDatabase(); + } + throw e; + } + }); + } + + static String listDatabasesSql() { + String placeholders = SYSTEM_SCHEMAS.stream() + .map(schema -> "'" + schema + "'") + .collect(Collectors.joining(",")); + return """ + SELECT username AS owner + FROM all_users + WHERE username IS NOT NULL + AND username NOT IN (%s) + AND username NOT LIKE 'APEX_%%' + AND username NOT LIKE 'FLOWS_%%' + AND username NOT LIKE '%%$%%' + ORDER BY CASE + WHEN username = SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') THEN 0 + WHEN username = SYS_CONTEXT('USERENV', 'SESSION_USER') THEN 1 + ELSE 2 + END, username + """.formatted(placeholders).stripIndent().trim(); + } + + static boolean isPgaLimitError(SQLException error) { + for (Throwable current = error; current != null; current = current.getCause()) { + if (current instanceof SQLException sqlException) { + if (sqlException.getErrorCode() == 4036 || String.valueOf(sqlException.getMessage()).contains("ORA-04036")) { + return true; + } + SQLException next = sqlException.getNextException(); + while (next != null) { + if (next.getErrorCode() == 4036 || String.valueOf(next.getMessage()).contains("ORA-04036")) { + return true; + } + next = next.getNextException(); + } + } else if (String.valueOf(current.getMessage()).contains("ORA-04036")) { + return true; + } + } + return false; + } + + private List queryDatabaseInfos(String sql) throws SQLException { + List result = new ArrayList<>(); + try (var stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + } + + private List currentSchemaDatabase() throws SQLException { + try (var stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SELECT SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA') FROM DUAL")) { + if (rs.next()) { + String schema = rs.getString(1); + if (schema != null && !schema.isBlank()) { + return List.of(new DatabaseInfo(schema)); + } + } + } + return List.of(); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + try { + String sql = """ + SELECT /*+ NO_QUERY_TRANSFORMATION */ o.OBJECT_NAME, + CASE o.OBJECT_TYPE WHEN 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS TABLE_TYPE, + c.COMMENTS + FROM ALL_OBJECTS o + LEFT JOIN ALL_TAB_COMMENTS c ON c.OWNER = o.OWNER AND c.TABLE_NAME = o.OBJECT_NAME + WHERE o.OWNER = ? AND o.OBJECT_TYPE IN ('TABLE','VIEW') + ORDER BY o.OBJECT_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + } catch (SQLException e) { + if (isPgaLimitError(e)) { + return List.of(); + } + throw e; + } + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + try { + String sql = """ + SELECT /*+ NO_QUERY_TRANSFORMATION */ OBJECT_NAME, OBJECT_TYPE FROM ALL_OBJECTS + WHERE OWNER = ? AND OBJECT_TYPE IN ('TABLE', 'VIEW', 'PROCEDURE', 'FUNCTION') + ORDER BY CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 ELSE 3 END, OBJECT_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + } catch (SQLException e) { + if (isPgaLimitError(e)) { + return List.of(); + } + throw e; + } + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String upperType = objectType.toUpperCase(Locale.ROOT); + String source = ""; + if ("VIEW".equals(upperType)) { + // ALL_VIEWS is more reliable than DBMS_METADATA.GET_DDL, + // which fails on XE editions where XSL stylesheets are missing. + try (var stmt = requireConnected().prepareStatement( + "SELECT TEXT FROM ALL_VIEWS WHERE OWNER = ? AND VIEW_NAME = ?")) { + stmt.setString(1, schema); + stmt.setString(2, name); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + String value = rs.getString(1); + source = value == null ? "" : value; + } + } + } + } else { + String dbmsType = switch (upperType) { + case "PROCEDURE" -> "PROCEDURE"; + case "FUNCTION" -> "FUNCTION"; + default -> throw new IllegalArgumentException("Unsupported object type: " + objectType); + }; + try (var stmt = requireConnected().prepareStatement( + "SELECT DBMS_METADATA.GET_DDL(?, ?, ?) FROM DUAL")) { + stmt.setString(1, dbmsType); + stmt.setString(2, name); + stmt.setString(3, schema); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + String value = rs.getString(1); + source = value == null ? "" : value; + } + } + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public String getTableDdl(String schema, String table) { + return unchecked(() -> { + try (var stmt = requireConnected().prepareStatement("SELECT DBMS_METADATA.GET_DDL(?, ?, ?) FROM DUAL")) { + stmt.setString(1, "TABLE"); + stmt.setString(2, table); + stmt.setString(3, schema); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + String value = rs.getString(1); + return value == null ? "" : value; + } + } + } + throw new IllegalArgumentException("Table not found: " + schema + "." + table); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + String sql = """ + SELECT c.COLUMN_NAME, c.DATA_TYPE, c.NULLABLE, c.DATA_PRECISION, c.DATA_SCALE, + c.DATA_LENGTH, c.CHAR_LENGTH, cc.COMMENTS, + CASE WHEN pk.COLUMN_NAME IS NULL THEN 0 ELSE 1 END AS IS_PK + FROM ALL_TAB_COLUMNS c + LEFT JOIN ALL_COL_COMMENTS cc + ON cc.OWNER = c.OWNER AND cc.TABLE_NAME = c.TABLE_NAME AND cc.COLUMN_NAME = c.COLUMN_NAME + LEFT JOIN ( + SELECT cols.COLUMN_NAME FROM ALL_CONS_COLUMNS cols + JOIN ALL_CONSTRAINTS cons + ON cols.CONSTRAINT_NAME = cons.CONSTRAINT_NAME AND cols.OWNER = cons.OWNER + WHERE cons.CONSTRAINT_TYPE = 'P' AND cons.OWNER = ? AND cons.TABLE_NAME = ? + ) pk ON pk.COLUMN_NAME = c.COLUMN_NAME + WHERE c.OWNER = ? AND c.TABLE_NAME = ? + ORDER BY c.COLUMN_ID + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + stmt.setString(3, schema); + stmt.setString(4, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String name = rs.getString("COLUMN_NAME"); + String baseType = rs.getString("DATA_TYPE"); + Integer numPrec = intOrNull(rs, "DATA_PRECISION"); + Integer numScale = intOrNull(rs, "DATA_SCALE"); + Integer dataLen = intOrNull(rs, "DATA_LENGTH"); + Integer charLen = intOrNull(rs, "CHAR_LENGTH"); + String dataType = formatDataType(baseType, numPrec, numScale, dataLen, charLen); + + result.add(new ColumnInfo( + name, + dataType, + "Y".equals(rs.getString("NULLABLE")), + null, + rs.getInt("IS_PK") == 1, + null, + rs.getString("COMMENTS"), + numPrec, + numScale, + charLen + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + String sql = """ + SELECT i.INDEX_NAME, + LISTAGG(ic.COLUMN_NAME, ',') WITHIN GROUP (ORDER BY ic.COLUMN_POSITION) AS COLUMNS, + i.UNIQUENESS, + CASE WHEN c.CONSTRAINT_TYPE = 'P' THEN 1 ELSE 0 END AS IS_PK, + i.INDEX_TYPE + FROM ALL_INDEXES i + JOIN ALL_IND_COLUMNS ic ON i.INDEX_NAME = ic.INDEX_NAME AND i.OWNER = ic.INDEX_OWNER AND i.TABLE_OWNER = ic.TABLE_OWNER + LEFT JOIN ALL_CONSTRAINTS c ON i.INDEX_NAME = c.INDEX_NAME AND i.TABLE_OWNER = c.OWNER + AND c.CONSTRAINT_TYPE = 'P' + WHERE i.TABLE_OWNER = ? AND i.TABLE_NAME = ? + GROUP BY i.INDEX_NAME, i.UNIQUENESS, c.CONSTRAINT_TYPE, i.INDEX_TYPE + ORDER BY i.INDEX_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String colsStr = rs.getString(2); + List columns = new ArrayList<>(); + if (colsStr != null && !colsStr.isEmpty()) { + columns.addAll(List.of(colsStr.split(","))); + } + result.add(new IndexInfo( + rs.getString(1), + columns, + "UNIQUE".equals(rs.getString(3)), + "1".equals(rs.getString(4)), + null, + rs.getString(5), + null, + null + )); + } + } + } + return result; + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + String sql = """ + SELECT c.CONSTRAINT_NAME, cc.COLUMN_NAME, rc.TABLE_NAME, rcc.COLUMN_NAME + FROM ALL_CONSTRAINTS c + JOIN ALL_CONS_COLUMNS cc ON c.CONSTRAINT_NAME = cc.CONSTRAINT_NAME AND c.OWNER = cc.OWNER + JOIN ALL_CONSTRAINTS rc ON c.R_CONSTRAINT_NAME = rc.CONSTRAINT_NAME AND c.R_OWNER = rc.OWNER + JOIN ALL_CONS_COLUMNS rcc ON rc.CONSTRAINT_NAME = rcc.CONSTRAINT_NAME AND rc.OWNER = rcc.OWNER + WHERE c.CONSTRAINT_TYPE = 'R' AND c.OWNER = ? AND c.TABLE_NAME = ? + ORDER BY c.CONSTRAINT_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString(1), + rs.getString(2), + rs.getString(3), + rs.getString(4) + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + String sql = """ + SELECT TRIGGER_NAME, TRIGGERING_EVENT, TRIGGER_TYPE + FROM ALL_TRIGGERS + WHERE OWNER = ? AND TABLE_NAME = ? + ORDER BY TRIGGER_NAME + """.stripIndent().trim(); + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo(rs.getString(1), rs.getString(2), rs.getString(3))); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + prepareExecutableSql(sql), + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + OracleAgent::stringResultValue + ); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + return JdbcExecutor.INSTANCE.executePage( + requireConnected(), + prepareExecutableSql(sql), + schema, + this::setSchemaSQL, + options, + OracleAgent::stringResultValue + ); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private static Properties connectionProperties(ConnectParams params) { + Properties props = new Properties(); + props.setProperty("user", params.getUsername()); + props.setProperty("password", params.getPassword()); + props.setProperty("oracle.jdbc.defaultNChar", "true"); + return props; + } + + static String buildUrl(ConnectParams params) { + return buildUrl(params, serviceName(params)); + } + + private static String buildUrl(ConnectParams params, String serviceName) { + String connectionString = params.getConnection_string(); + if (connectionString != null && !connectionString.trim().isEmpty()) { + return connectionString; + } + return "jdbc:oracle:thin:@" + params.getHost() + ":" + params.getPort() + "/" + serviceName; + } + + private static String serviceName(ConnectParams params) { + String serviceName = params.getDatabase(); + if (serviceName.toUpperCase(Locale.ROOT).startsWith("SYSDBA:")) { + return serviceName.substring(7); + } + return serviceName; + } + + private static String formatDataType( + String base, + Integer numPrec, + Integer numScale, + Integer dataLen, + Integer charLen + ) { + return switch (base.toUpperCase(Locale.ROOT)) { + case "VARCHAR2", "NVARCHAR2", "CHAR", "NCHAR" -> { + Integer len = charLen == null ? dataLen : charLen; + yield len == null ? base : base + "(" + len + ")"; + } + case "NUMBER" -> { + if (numPrec != null && numScale != null && numScale > 0) { + yield base + "(" + numPrec + "," + numScale + ")"; + } + if (numPrec != null && numPrec > 0) { + yield base + "(" + numPrec + ")"; + } + yield base; + } + case "RAW" -> dataLen == null ? "RAW" : "RAW(" + dataLen + ")"; + default -> base; + }; + } + + private static Object stringResultValue(ResultSet rs, int index, int sqlType) { + try { + return JdbcExecutor.stringResultValue(rs, index, sqlType); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static String trimEndSemicolons(String sql) { + String trimmed = sql; + while (trimmed.endsWith(";")) { + trimmed = trimmed.substring(0, trimmed.length() - 1); + } + return trimmed; + } + + static String prepareExecutableSql(String sql) { + String trimmed = sql.trim(); + if (PLSQL_OBJECT_DDL_RE.matcher(trimmed).find()) { + return trimmed; + } + return rewriteFetchFirst(trimEndSemicolons(trimmed)); + } + + public static void main(String[] args) { + new JsonRpcServer(new OracleAgent()).run(); + } +} diff --git a/agents/drivers/oracle/src/test/java/com/dbx/agent/oracle/OracleAgentTest.java b/agents/drivers/oracle/src/test/java/com/dbx/agent/oracle/OracleAgentTest.java new file mode 100644 index 0000000000..a612e65f8c --- /dev/null +++ b/agents/drivers/oracle/src/test/java/com/dbx/agent/oracle/OracleAgentTest.java @@ -0,0 +1,76 @@ +package com.dbx.agent.oracle; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.sql.SQLException; +import java.util.Locale; + +class OracleAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new OracleAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL DBMS_XPLAN.DISPLAY_CURSOR()"; + } + + @Test + void buildUrlUsesExplicitConnectionString() { + ConnectParams params = new ConnectParams( + "oracle.example.com", + 1521, + "ORCL", + "scott", + "tiger", + "", + "jdbc:oracle:thin:@oracle.example.com:1521:ORCL", + false + ); + + Assertions.assertEquals("jdbc:oracle:thin:@oracle.example.com:1521:ORCL", OracleAgent.buildUrl(params)); + } + + @Test + void prepareExecutableSqlKeepsPlsqlObjectTerminator() { + String sql = "CREATE OR REPLACE PROCEDURE APP_PROC AS BEGIN NULL; END;"; + + Assertions.assertEquals(sql, OracleAgent.prepareExecutableSql(sql)); + } + + @Test + void prepareExecutableSqlTrimsPlainStatementTerminator() { + Assertions.assertEquals("SELECT 1 FROM DUAL", OracleAgent.prepareExecutableSql("SELECT 1 FROM DUAL;")); + } + + @Test + void prepareExecutableSqlStillRewritesFetchFirst() { + Assertions.assertEquals( + "SELECT * FROM (SELECT * FROM EMP) WHERE ROWNUM <= 10", + OracleAgent.prepareExecutableSql("SELECT * FROM EMP FETCH FIRST 10 ROWS ONLY;") + ); + } + + @Test + void listDatabasesSqlUsesUserDictionaryInsteadOfObjectDictionary() { + String sql = OracleAgent.listDatabasesSql().toUpperCase(Locale.ROOT); + + Assertions.assertTrue(sql.contains("ALL_USERS")); + Assertions.assertFalse(sql.contains("ALL_TABLES")); + Assertions.assertFalse(sql.contains("ALL_VIEWS")); + } + + @Test + void detectsPgaLimitError() { + Assertions.assertTrue(OracleAgent.isPgaLimitError(new SQLException( + "ORA-04036: PGA memory used by the instance exceeds PGA_AGGREGATE_LIMIT", + "72000", + 4036 + ))); + } +} diff --git a/agents/drivers/saphana/build.gradle b/agents/drivers/saphana/build.gradle new file mode 100644 index 0000000000..112aff5227 --- /dev/null +++ b/agents/drivers/saphana/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.sap.cloud.db.jdbc:ngdbc:2.28.7' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'SAP HANA', 'Main-Class': 'com.dbx.agent.saphana.SapHanaAgent') + } +} diff --git a/agents/drivers/saphana/src/main/java/com/dbx/agent/saphana/SapHanaAgent.java b/agents/drivers/saphana/src/main/java/com/dbx/agent/saphana/SapHanaAgent.java new file mode 100644 index 0000000000..269b2fd24b --- /dev/null +++ b/agents/drivers/saphana/src/main/java/com/dbx/agent/saphana/SapHanaAgent.java @@ -0,0 +1,26 @@ +package com.dbx.agent.saphana; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; +import java.util.Arrays; +import java.util.Collections; + +public final class SapHanaAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile SAPHANA_PROFILE = new JdbcAgentProfile( + "com.sap.db.jdbc.Driver", + "jdbc:sap://{host}:{port}/?databaseName={database}", + 30015, + false, + Collections.emptySet(), + Arrays.asList("COLUMN TABLE", "ROW TABLE", "TABLE", "VIEW") + ); + + public SapHanaAgent() { + super(SAPHANA_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new SapHanaAgent()).run(); + } +} diff --git a/agents/drivers/snowflake/build.gradle b/agents/drivers/snowflake/build.gradle new file mode 100644 index 0000000000..7ba1819842 --- /dev/null +++ b/agents/drivers/snowflake/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'net.snowflake:snowflake-jdbc:3.22.0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Snowflake', 'Main-Class': 'com.dbx.agent.snowflake.SnowflakeAgent') + } +} diff --git a/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java b/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java new file mode 100644 index 0000000000..da11603db8 --- /dev/null +++ b/agents/drivers/snowflake/src/main/java/com/dbx/agent/snowflake/SnowflakeAgent.java @@ -0,0 +1,243 @@ +package com.dbx.agent.snowflake; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Locale; + +public final class SnowflakeAgent extends AbstractJdbcAgent { + + @Override + protected String driverClass() { + return "net.snowflake.client.jdbc.SnowflakeDriver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString("name"))); + } + } + result.sort(Comparator.comparing(DatabaseInfo::getName)); + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW SCHEMAS")) { + while (rs.next()) { + result.add(rs.getString("name")); + } + } + Collections.sort(result); + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT TABLE_NAME, TABLE_TYPE " + + "FROM INFORMATION_SCHEMA.TABLES " + + "WHERE TABLE_SCHEMA = ? " + + "ORDER BY TABLE_NAME"; + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("TABLE_NAME"), + normalizeTableType(rs.getString("TABLE_TYPE")), + null + )); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + + String sql = "SELECT PROCEDURE_NAME, 'PROCEDURE' " + + "FROM INFORMATION_SCHEMA.PROCEDURES " + + "WHERE PROCEDURE_SCHEMA = ? " + + "ORDER BY PROCEDURE_NAME"; + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String ddlType = objectType.toUpperCase(Locale.ROOT); + String sql = "SELECT GET_DDL('" + ddlType + "', '\"" + schema + "\".\"" + name + "\"')"; + String source = ""; + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + if (rs.next()) { + String value = rs.getString(1); + source = value == null ? "" : value; + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Connection conn = requireConnected(); + List primaryKeys = new ArrayList<>(); + try (java.sql.Statement stmt = conn.createStatement()) { + String qualifiedTable = JdbcIdentifiers.INSTANCE.doubleQuote(schema) + "." + JdbcIdentifiers.INSTANCE.doubleQuote(table); + try (ResultSet rs = stmt.executeQuery("SHOW PRIMARY KEYS IN TABLE " + qualifiedTable)) { + while (rs.next()) { + primaryKeys.add(rs.getString("column_name")); + } + } + } catch (Exception ignored) { + // Snowflake may not return primary key metadata for every object. + } + + List result = new ArrayList<>(); + String sql = "SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT, " + + "NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH, COMMENT " + + "FROM INFORMATION_SCHEMA.COLUMNS " + + "WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? " + + "ORDER BY ORDINAL_POSITION"; + try (java.sql.PreparedStatement stmt = conn.prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String colName = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + colName, + rs.getString("DATA_TYPE"), + "YES".equals(rs.getString("IS_NULLABLE")), + rs.getString("COLUMN_DEFAULT"), + primaryKeys.contains(colName), + null, + blankToNull(rs.getString("COMMENT")), + intOrNull(rs, "NUMERIC_PRECISION"), + intOrNull(rs, "NUMERIC_SCALE"), + intOrNull(rs, "CHARACTER_MAXIMUM_LENGTH") + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + try (java.sql.Statement stmt = requireConnected().createStatement()) { + String qualifiedTable = JdbcIdentifiers.INSTANCE.doubleQuote(schema) + "." + JdbcIdentifiers.INSTANCE.doubleQuote(table); + try (ResultSet rs = stmt.executeQuery("SHOW IMPORTED KEYS IN TABLE " + qualifiedTable)) { + List result = new ArrayList<>(); + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("fk_name"), + rs.getString("fk_column_name"), + rs.getString("pk_table_name"), + rs.getString("pk_column_name") + )); + } + return result; + } + } catch (Exception ignored) { + return Collections.emptyList(); + } + }); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:snowflake://" + params.getHost() + ":" + params.getPort() + "/?db=" + params.getDatabase(); + } + + private static String normalizeTableType(String tableType) { + return "BASE TABLE".equals(tableType) ? "TABLE" : tableType; + } + + private static String blankToNull(String value) { + return value == null || value.isEmpty() ? null : value; + } + + private static Integer intOrNull(ResultSet rs, String column) throws java.sql.SQLException { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + public static void main(String[] args) { + new JsonRpcServer(new SnowflakeAgent()).run(); + } +} diff --git a/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java b/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java new file mode 100644 index 0000000000..847d48878d --- /dev/null +++ b/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentMetadataTest.java @@ -0,0 +1,29 @@ +package com.dbx.agent.snowflake; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.test.JdbcAgentFake; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class SnowflakeAgentMetadataTest { + @Test + void quotesSchemaAndTableIdentifiersInKeyMetadataSql() { + SnowflakeAgent agent = new SnowflakeAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.getColumns("bad\"schema", "bad\"table"); + agent.listForeignKeys("bad\"schema", "bad\"table"); + + Assertions.assertTrue( + JdbcMetadataSqlFake.statements.contains( + "SHOW PRIMARY KEYS IN TABLE \"bad\"\"schema\".\"bad\"\"table\"" + ) + ); + Assertions.assertTrue( + JdbcMetadataSqlFake.statements.contains( + "SHOW IMPORTED KEYS IN TABLE \"bad\"\"schema\".\"bad\"\"table\"" + ) + ); + } +} diff --git a/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentTest.java b/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentTest.java new file mode 100644 index 0000000000..11af6e4638 --- /dev/null +++ b/agents/drivers/snowflake/src/test/java/com/dbx/agent/snowflake/SnowflakeAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.snowflake; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class SnowflakeAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new SnowflakeAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } +} diff --git a/agents/drivers/sundb/build.gradle b/agents/drivers/sundb/build.gradle new file mode 100644 index 0000000000..ef14a7312d --- /dev/null +++ b/agents/drivers/sundb/build.gradle @@ -0,0 +1,13 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) +} + +tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'SunDB', + 'Agent-External-Driver': 'true', + 'Main-Class': 'com.dbx.agent.sundb.SundbAgent' + ) + } +} diff --git a/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java b/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java new file mode 100644 index 0000000000..9a23eb897b --- /dev/null +++ b/agents/drivers/sundb/src/main/java/com/dbx/agent/sundb/SundbAgent.java @@ -0,0 +1,384 @@ +package com.dbx.agent.sundb; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +public final class SundbAgent extends BaseDatabaseAgent { + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + Class.forName("com.sundb.jdbc.SundbDriver"); + String url = buildUrl(params); + connection = DriverManager.getConnection(url, params.getUsername(), params.getPassword()); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + Class.forName("com.sundb.jdbc.SundbDriver"); + String url = buildUrl(params); + try (Connection conn = DriverManager.getConnection(url, params.getUsername(), params.getPassword())) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String tableType = rs.getString("TABLE_TYPE"); + if ("BASE TABLE".equals(tableType)) { + tableType = "TABLE"; + } + result.add(new TableInfo(rs.getString("TABLE_NAME"), tableType, null)); + } + } + } + return result; + }); + } + + @Override + public List listObjects(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + + try (var stmt = requireConnected().prepareStatement( + "SELECT ROUTINE_NAME, ROUTINE_TYPE FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = ? ORDER BY ROUTINE_NAME" + )) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo(rs.getString(1), rs.getString(2), schema, null)); + } + } + } + return result; + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + return unchecked(() -> { + String sql = switch (objectType.toUpperCase(Locale.ROOT)) { + case "VIEW" -> "SHOW CREATE VIEW `" + name.replace("`", "``") + "`"; + case "PROCEDURE" -> "SHOW CREATE PROCEDURE `" + name.replace("`", "``") + "`"; + case "FUNCTION" -> "SHOW CREATE FUNCTION `" + name.replace("`", "``") + "`"; + default -> throw new IllegalArgumentException("Unsupported object type: " + objectType); + }; + + String source = ""; + try (var stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + if (rs.next()) { + int index = "VIEW".equals(objectType.toUpperCase(Locale.ROOT)) ? 1 : 2; + String value = rs.getString(index + 1); + source = value == null ? "" : value; + } + } + return new ObjectSource(name, objectType, schema, source); + }); + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + Set primaryKeys = new HashSet<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT COLUMN_NAME + FROM information_schema.KEY_COLUMN_USAGE + WHERE TABLE_SCHEMA = ? + AND TABLE_NAME = ? + AND CONSTRAINT_NAME = 'PRIMARY' + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + primaryKeys.add(rs.getString("COLUMN_NAME")); + } + } + } + + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_DEFAULT, EXTRA, COLUMN_COMMENT, + NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH + FROM information_schema.COLUMNS + WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? + ORDER BY ORDINAL_POSITION + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + String columnName = rs.getString("COLUMN_NAME"); + result.add(new ColumnInfo( + columnName, + rs.getString("COLUMN_TYPE"), + "YES".equals(rs.getString("IS_NULLABLE")), + rs.getString("COLUMN_DEFAULT"), + primaryKeys.contains(columnName), + rs.getString("EXTRA"), + emptyToNull(rs.getString("COLUMN_COMMENT")), + integerOrNull(rs, "NUMERIC_PRECISION"), + integerOrNull(rs, "NUMERIC_SCALE"), + numberToIntOrNull(rs, "CHARACTER_MAXIMUM_LENGTH") + )); + } + } + } + return result; + }); + } + + @Override + public List listIndexes(String schema, String table) { + return unchecked(() -> { + try (var stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery( + "SHOW INDEX FROM " + JdbcIdentifiers.INSTANCE.backtick(table) + + " FROM " + JdbcIdentifiers.INSTANCE.backtick(schema) + )) { + Map> indexMap = new LinkedHashMap<>(); + Map uniqueMap = new HashMap<>(); + Map typeMap = new HashMap<>(); + + while (rs.next()) { + String indexName = rs.getString("Key_name"); + String columnName = rs.getString("Column_name"); + int sequence = rs.getInt("Seq_in_index"); + int nonUnique = rs.getInt("Non_unique"); + String indexType = rs.getString("Index_type"); + + indexMap.computeIfAbsent(indexName, ignored -> new ArrayList<>()) + .add(new IndexColumn(sequence, columnName)); + uniqueMap.put(indexName, nonUnique == 0); + typeMap.put(indexName, indexType); + } + + List result = new ArrayList<>(); + for (Map.Entry> entry : indexMap.entrySet()) { + String name = entry.getKey(); + List columns = entry.getValue().stream() + .sorted(Comparator.comparingInt(IndexColumn::sequence)) + .map(IndexColumn::name) + .toList(); + result.add(new IndexInfo( + name, + columns, + uniqueMap.getOrDefault(name, false), + "PRIMARY".equals(name), + null, + typeMap.get(name), + null, + null + )); + } + return result; + } + }); + } + + @Override + public List listForeignKeys(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT CONSTRAINT_NAME, COLUMN_NAME, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME + FROM information_schema.KEY_COLUMN_USAGE + WHERE TABLE_SCHEMA = ? + AND TABLE_NAME = ? + AND REFERENCED_TABLE_NAME IS NOT NULL + ORDER BY CONSTRAINT_NAME + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ForeignKeyInfo( + rs.getString("CONSTRAINT_NAME"), + rs.getString("COLUMN_NAME"), + rs.getString("REFERENCED_TABLE_NAME"), + rs.getString("REFERENCED_COLUMN_NAME") + )); + } + } + } + return result; + }); + } + + @Override + public List listTriggers(String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + try (var stmt = requireConnected().prepareStatement( + """ + SELECT TRIGGER_NAME, EVENT_MANIPULATION, ACTION_TIMING + FROM information_schema.TRIGGERS + WHERE TRIGGER_SCHEMA = ? AND EVENT_OBJECT_TABLE = ? + ORDER BY TRIGGER_NAME + """ + )) { + stmt.setString(1, schema); + stmt.setString(2, table); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TriggerInfo( + rs.getString("TRIGGER_NAME"), + rs.getString("EVENT_MANIPULATION"), + rs.getString("ACTION_TIMING") + )); + } + } + } + return result; + }); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::getResultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + JdbcIdentifiers.INSTANCE.backtick(schema); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private static String buildUrl(ConnectParams params) { + return "jdbc:sundb://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private static Integer integerOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Integer ? (Integer) value : null; + } + + private static Integer numberToIntOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static String emptyToNull(String value) { + return value == null || value.isEmpty() ? null : value; + } + + private Object getResultValue(ResultSet rs, int index, int sqlType) { + try { + Object value = switch (sqlType) { + case Types.BIGINT -> rs.getLong(index); + case Types.INTEGER, Types.SMALLINT, Types.TINYINT -> rs.getInt(index); + case Types.FLOAT, Types.REAL -> rs.getFloat(index); + case Types.DOUBLE -> rs.getDouble(index); + case Types.DECIMAL, Types.NUMERIC -> rs.getBigDecimal(index); + case Types.BOOLEAN, Types.BIT -> rs.getBoolean(index); + default -> rs.getString(index); + }; + return rs.wasNull() ? null : value; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private record IndexColumn(int sequence, String name) { + } + + public static void main(String[] args) { + new JsonRpcServer(new SundbAgent()).run(); + } +} diff --git a/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java b/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java new file mode 100644 index 0000000000..943e7953f2 --- /dev/null +++ b/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentMetadataTest.java @@ -0,0 +1,23 @@ +package com.dbx.agent.sundb; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.test.JdbcAgentFake; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import java.util.Collections; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class SundbAgentMetadataTest { + @Test + void quotesSchemaAndTableIdentifiersInIndexMetadataSql() { + SundbAgent agent = new SundbAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listIndexes("bad`schema", "bad`table"); + + Assertions.assertEquals( + Collections.singletonList("SHOW INDEX FROM `bad``table` FROM `bad``schema`"), + JdbcMetadataSqlFake.statements + ); + } +} diff --git a/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentTest.java b/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentTest.java new file mode 100644 index 0000000000..1a563804c7 --- /dev/null +++ b/agents/drivers/sundb/src/test/java/com/dbx/agent/sundb/SundbAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.sundb; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class SundbAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new SundbAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } +} diff --git a/agents/drivers/tdengine/build.gradle b/agents/drivers/tdengine/build.gradle new file mode 100644 index 0000000000..3d537919af --- /dev/null +++ b/agents/drivers/tdengine/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.taosdata.jdbc:taos-jdbcdriver:3.6.3' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'TDengine', 'Main-Class': 'com.dbx.agent.tdengine.TDengineAgent') + } +} diff --git a/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java b/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java new file mode 100644 index 0000000000..3f3b9cadeb --- /dev/null +++ b/agents/drivers/tdengine/src/main/java/com/dbx/agent/tdengine/TDengineAgent.java @@ -0,0 +1,598 @@ +package com.dbx.agent.tdengine; + +import com.dbx.agent.BaseDatabaseAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.QueryPageOptions; +import com.dbx.agent.QueryPageResult; +import com.dbx.agent.QueryResult; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Timestamp; +import java.sql.Types; +import java.time.format.DateTimeFormatter; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +public final class TDengineAgent extends BaseDatabaseAgent { + private static final DateTimeFormatter TDENGINE_TIMESTAMP_FORMAT = + DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS"); + private static final Pattern NUMERIC_PRECISION_PATTERN = + Pattern.compile("(?i)^(decimal|numeric)\\((\\d+)(?:,\\s*\\d+)?\\)"); + private static final Pattern NUMERIC_SCALE_PATTERN = + Pattern.compile("(?i)^(decimal|numeric)\\(\\d+,\\s*(\\d+)\\)"); + private static final Pattern CHARACTER_LENGTH_PATTERN = + Pattern.compile("(?i)^(binary|nchar|varchar|varbinary)\\((\\d+)\\)"); + + private Connection connection; + + @Override + public Connection getConnection() { + return connection; + } + + @Override + public void connect(ConnectParams params) { + uncheckedVoid(() -> { + connection = TDengineConnectionFactory.open(params); + }); + } + + @Override + public boolean testConnection(ConnectParams params) { + return unchecked(() -> { + try (Connection conn = TDengineConnectionFactory.open(params)) { + return conn.isValid(5); + } + }); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW DATABASES")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + List result = new ArrayList<>(); + for (DatabaseInfo database : listDatabases()) { + result.add(database.getName()); + } + return result; + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "STABLES", "STABLE")); + result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "TABLES", "TABLE")); + + Map distinct = new LinkedHashMap<>(); + for (TableInfo table : result) { + distinct.putIfAbsent(table.getTable_type() + ":" + table.getName(), table); + } + List sorted = new ArrayList<>(distinct.values()); + sorted.sort(Comparator.comparing(table -> table.getName().toLowerCase(Locale.ROOT))); + return sorted; + }); + } + + @Override + public List listObjects(String schema) { + List result = new ArrayList<>(); + for (TableInfo table : listTables(schema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment())); + } + return result; + } + + @Override + public List getColumns(String schema, String table) { + return unchecked(() -> { + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("DESCRIBE " + qualifiedName(schema, table))) { + return readDescribeColumns(rs); + } + }); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + String source = getCreateSql(schema, name, objectType); + if (source.isBlank()) { + source = getTableDdl(schema, name); + } + return new ObjectSource(name, objectType, schema, source); + } + + @Override + public String getTableDdl(String schema, String table) { + String source = getCreateSql(schema, table, "STABLE"); + if (source.isBlank()) { + source = getCreateSql(schema, table, "TABLE"); + } + if (source.isBlank()) { + return super.getTableDdl(schema, table); + } + return source; + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { + return JdbcExecutor.INSTANCE.execute( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options.getMaxRows(), + options.getFetchSize(), + options.getTimeoutSecs(), + this::tdengineResultValue + ); + } + + @Override + public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { + return JdbcExecutor.INSTANCE.executePage( + requireConnected(), + sql, + schema, + this::setSchemaSQL, + options, + this::tdengineResultValue + ); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + quoteIdentifier(schema); + } + + @Override + public void disconnect() { + uncheckedVoid(() -> { + if (connection != null) { + connection.close(); + } + connection = null; + }); + } + + private List queryTables(String sql, String tableType) throws Exception { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + result.add(new TableInfo(rs.getString(1), tableType, null)); + } + } + return result; + } + + private static List readDescribeColumns(ResultSet rs) throws Exception { + List result = new ArrayList<>(); + int ordinal = 0; + while (rs.next()) { + ordinal += 1; + String name = rs.getString(1); + String dataType = coalesce(rs.getString(2)); + String note = optionalString(rs, 4); + boolean isTag = note != null && note.toUpperCase(Locale.ROOT).contains("TAG"); + result.add(new ColumnInfo( + name, + dataType, + ordinal != 1, + null, + ordinal == 1 && !isTag, + note, + isTag ? "TAG" : null, + parseNumericPrecision(dataType), + parseNumericScale(dataType), + parseCharacterMaximumLength(dataType) + )); + } + return result; + } + + private String getCreateSql(String schema, String name, String objectType) { + String showType = switch (objectType.toUpperCase(Locale.ROOT)) { + case "STABLE", "SUPER TABLE", "SUPERTABLE" -> "STABLE"; + case "TABLE", "BASE TABLE", "CHILD TABLE" -> "TABLE"; + default -> null; + }; + if (showType == null) { + return ""; + } + + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW CREATE " + showType + " " + qualifiedName(schema, name))) { + if (!rs.next()) { + return ""; + } + int columnCount = rs.getMetaData().getColumnCount(); + for (int i = 2; i <= columnCount; i++) { + String value = rs.getString(i); + if (value != null && value.toUpperCase(Locale.ROOT).contains("CREATE")) { + return value; + } + } + String value = rs.getString(columnCount); + return value == null ? "" : value; + } catch (Exception e) { + return ""; + } + } + + private Object tdengineResultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = switch (sqlType) { + case Types.BIGINT -> rs.getLong(index); + case Types.INTEGER, Types.SMALLINT, Types.TINYINT -> rs.getInt(index); + case Types.FLOAT, Types.REAL -> rs.getFloat(index); + case Types.DOUBLE -> rs.getDouble(index); + case Types.DECIMAL, Types.NUMERIC -> rs.getBigDecimal(index); + case Types.BOOLEAN, Types.BIT -> rs.getBoolean(index); + default -> rs.getObject(index); + }; + return rs.wasNull() ? null : decodeTdengineValue(value); + }); + } + + public static Object decodeTdengineValue(Object value) { + if (value instanceof byte[] bytes) { + return new String(bytes, StandardCharsets.UTF_8); + } + if (value instanceof Timestamp timestamp) { + return TDENGINE_TIMESTAMP_FORMAT.format(timestamp.toLocalDateTime()); + } + return value; + } + + private static String quoteQualifiedPrefix(String schema) { + String trimmed = schema.trim(); + return trimmed.isEmpty() ? "" : quoteIdentifier(trimmed) + "."; + } + + private static String qualifiedName(String schema, String name) { + String table = quoteIdentifier(name); + String trimmed = schema.trim(); + return trimmed.isEmpty() ? table : quoteIdentifier(trimmed) + "." + table; + } + + private static String quoteIdentifier(String identifier) { + return "`" + identifier.replace("`", "``") + "`"; + } + + private static Integer parseNumericPrecision(String dataType) { + return parseIntGroup(NUMERIC_PRECISION_PATTERN, dataType, 2); + } + + private static Integer parseNumericScale(String dataType) { + return parseIntGroup(NUMERIC_SCALE_PATTERN, dataType, 2); + } + + private static Integer parseCharacterMaximumLength(String dataType) { + return parseIntGroup(CHARACTER_LENGTH_PATTERN, dataType, 2); + } + + private static Integer parseIntGroup(Pattern pattern, String value, int group) { + Matcher matcher = pattern.matcher(value); + if (!matcher.find()) { + return null; + } + try { + return Integer.valueOf(matcher.group(group)); + } catch (NumberFormatException e) { + return null; + } + } + + private static String optionalString(ResultSet rs, int index) { + try { + return rs.getString(index); + } catch (Exception e) { + return null; + } + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + public static void main(String[] args) { + new JsonRpcServer(new TDengineAgent()).run(); + } +} + +final class TDengineJdbcUrl { + enum TransportPreference { + AUTO, + WEBSOCKET, + REST + } + + private static final String PARAM_TRANSPORT = "transport"; + private static final String PARAM_DBX_TRANSPORT = "dbx.transport"; + + private TDengineJdbcUrl() { + } + + static String from(ConnectParams params) { + TransportPreference preference = transportPreference(params.getUrl_params()); + if (preference == TransportPreference.REST) { + return from(params, TDengineTransport.REST); + } + String host = params.getHost().isBlank() ? "localhost" : params.getHost(); + int port = params.getPort() > 0 ? params.getPort() : 6041; + String database = params.getDatabase().trim(); + String path = database.isBlank() ? "/" : "/" + database; + String query = normalizeQueryString(params.getUrl_params()); + String suffix = query.isBlank() ? "" : "?" + query; + return "jdbc:TAOS-WS://" + host + ":" + port + path + suffix; + } + + static String from(ConnectParams params, TDengineTransport transport) { + String host = params.getHost().isBlank() ? "localhost" : params.getHost(); + int port = params.getPort() > 0 ? params.getPort() : 6041; + String database = params.getDatabase().trim(); + String path = database.isBlank() ? "/" : "/" + database; + String query = normalizeQueryString(params.getUrl_params()); + String suffix = query.isBlank() ? "" : "?" + query; + return transport.urlPrefix() + host + ":" + port + path + suffix; + } + + static String sanitizeConnectionString(String connectionString) { + if (connectionString == null || connectionString.isBlank()) { + return ""; + } + int queryIndex = connectionString.indexOf('?'); + if (queryIndex < 0) { + return connectionString; + } + String base = connectionString.substring(0, queryIndex); + String rawQuery = connectionString.substring(queryIndex + 1); + int fragmentIndex = rawQuery.indexOf('#'); + String fragment = ""; + if (fragmentIndex >= 0) { + fragment = rawQuery.substring(fragmentIndex); + rawQuery = rawQuery.substring(0, fragmentIndex); + } + String query = normalizeQueryString(rawQuery); + if (query.isBlank()) { + return base + fragment; + } + return base + "?" + query + fragment; + } + + static TransportPreference transportPreference(String rawQuery) { + String query = rawQuery == null ? "" : rawQuery.trim(); + if (query.startsWith("?")) { + query = query.substring(1); + } + if (query.isBlank()) { + return TransportPreference.AUTO; + } + String[] parts = query.split("&"); + for (String part : parts) { + if (part == null || part.isBlank()) { + continue; + } + int index = part.indexOf('='); + String key = index < 0 ? part.trim() : part.substring(0, index).trim(); + if (!PARAM_TRANSPORT.equalsIgnoreCase(key) && !PARAM_DBX_TRANSPORT.equalsIgnoreCase(key)) { + continue; + } + String value = index < 0 ? "" : part.substring(index + 1).trim(); + if ("rest".equalsIgnoreCase(value) || "rs".equalsIgnoreCase(value)) { + return TransportPreference.REST; + } + if ("ws".equalsIgnoreCase(value) || "websocket".equalsIgnoreCase(value)) { + return TransportPreference.WEBSOCKET; + } + } + return TransportPreference.AUTO; + } + + private static String normalizeQueryString(String rawQuery) { + String query = rawQuery == null ? "" : rawQuery.trim(); + if (query.startsWith("?")) { + query = query.substring(1); + } + if (query.isBlank()) { + return ""; + } + String[] parts = query.split("&"); + List kept = new ArrayList<>(); + for (String part : parts) { + if (part == null || part.isBlank()) { + continue; + } + int index = part.indexOf('='); + String key = index < 0 ? part.trim() : part.substring(0, index).trim(); + if (PARAM_TRANSPORT.equalsIgnoreCase(key) || PARAM_DBX_TRANSPORT.equalsIgnoreCase(key)) { + continue; + } + kept.add(part); + } + return String.join("&", kept); + } +} + +enum TDengineTransport { + WEBSOCKET("jdbc:TAOS-WS://", "com.taosdata.jdbc.ws.WebSocketDriver"), + REST("jdbc:TAOS-RS://", "com.taosdata.jdbc.rs.RestfulDriver"); + + private final String urlPrefix; + private final String driverClass; + + TDengineTransport(String urlPrefix, String driverClass) { + this.urlPrefix = urlPrefix; + this.driverClass = driverClass; + } + + String urlPrefix() { + return urlPrefix; + } + + String driverClass() { + return driverClass; + } +} + +final class TDengineConnectionFactory { + private static final int AUTO_TRANSPORT_ATTEMPT_TIMEOUT_SECONDS = 8; + + private TDengineConnectionFactory() { + } + + static Connection open(ConnectParams params) throws Exception { + String explicitConnectionString = params.getConnection_string() == null ? "" : params.getConnection_string().trim(); + if (!explicitConnectionString.isBlank()) { + String sanitizedConnectionString = TDengineJdbcUrl.sanitizeConnectionString(explicitConnectionString); + Class.forName(driverClassFromConnectionString(sanitizedConnectionString)); + return DriverManager.getConnection(sanitizedConnectionString, params.getUsername(), params.getPassword()); + } + + TDengineJdbcUrl.TransportPreference preference = TDengineJdbcUrl.transportPreference(params.getUrl_params()); + TDengineTransport[] candidates = transportsFor(preference); + Exception firstError = null; + for (TDengineTransport transport : candidates) { + try { + Connection connection = preference == TDengineJdbcUrl.TransportPreference.AUTO + ? openWithAttemptTimeout(params, transport) + : openDirect(params, transport); + return connection; + } catch (Exception e) { + if (firstError == null) { + firstError = e; + } else { + firstError.addSuppressed(e); + } + } + } + if (firstError != null) { + throw new RuntimeException("Failed to connect TDengine using transports: " + candidateNames(candidates), firstError); + } + throw new IllegalStateException("No TDengine transport available."); + } + + private static Connection openDirect(ConnectParams params, TDengineTransport transport) throws Exception { + Class.forName(transport.driverClass()); + String url = TDengineJdbcUrl.from(params, transport); + return DriverManager.getConnection(url, params.getUsername(), params.getPassword()); + } + + private static Connection openWithAttemptTimeout(ConnectParams params, TDengineTransport transport) throws Exception { + String threadName = "tdengine-connect-" + transport.name().toLowerCase(Locale.ROOT); + ThreadFactory factory = runnable -> { + Thread thread = new Thread(runnable, threadName); + thread.setDaemon(true); + return thread; + }; + ExecutorService executor = Executors.newSingleThreadExecutor(factory); + try { + Future future = executor.submit(new Callable() { + @Override + public Connection call() throws Exception { + return openDirect(params, transport); + } + }); + return future.get(AUTO_TRANSPORT_ATTEMPT_TIMEOUT_SECONDS, TimeUnit.SECONDS); + } catch (TimeoutException e) { + throw new RuntimeException( + "TDengine " + transport.name().toLowerCase(Locale.ROOT) + + " transport connect timed out after " + AUTO_TRANSPORT_ATTEMPT_TIMEOUT_SECONDS + "s", + e + ); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while connecting TDengine via " + transport.name().toLowerCase(Locale.ROOT), e); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception exception) { + throw exception; + } + throw new RuntimeException(cause); + } finally { + executor.shutdownNow(); + } + } + + private static TDengineTransport[] transportsFor(TDengineJdbcUrl.TransportPreference preference) { + return switch (preference) { + case REST -> new TDengineTransport[] {TDengineTransport.REST}; + case WEBSOCKET -> new TDengineTransport[] {TDengineTransport.WEBSOCKET}; + case AUTO -> new TDengineTransport[] {TDengineTransport.WEBSOCKET, TDengineTransport.REST}; + }; + } + + private static String candidateNames(TDengineTransport[] candidates) { + List names = new ArrayList<>(); + for (TDengineTransport candidate : candidates) { + names.add(candidate.name().toLowerCase(Locale.ROOT)); + } + return String.join(" -> ", names); + } + + private static String driverClassFromConnectionString(String connectionString) { + if (connectionString.regionMatches(true, 0, TDengineTransport.REST.urlPrefix(), 0, TDengineTransport.REST.urlPrefix().length())) { + return TDengineTransport.REST.driverClass(); + } + if (connectionString.regionMatches(true, 0, TDengineTransport.WEBSOCKET.urlPrefix(), 0, TDengineTransport.WEBSOCKET.urlPrefix().length())) { + return TDengineTransport.WEBSOCKET.driverClass(); + } + return TDengineTransport.WEBSOCKET.driverClass(); + } +} diff --git a/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentPagingTest.java b/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentPagingTest.java new file mode 100644 index 0000000000..91ef0b9174 --- /dev/null +++ b/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentPagingTest.java @@ -0,0 +1,178 @@ +package com.dbx.agent.tdengine; + +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.test.TestSupport; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.sql.Types; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class TDengineAgentPagingTest { + @Test + void executeQueryPageDecodesBinaryVarcharValuesAsStrings() { + TDengineAgent agent = new TDengineAgent(); + TestSupport.setPrivateConnection(agent, binaryVarcharConnection()); + JsonRpcServer server = new JsonRpcServer(agent); + + String response = handleRequest(server, """ + { + "jsonrpc": "2.0", + "id": 1, + "method": "execute_query_page", + "params": { + "sql": "SELECT v FROM sample", + "schema": "power", + "pageSize": 100 + } + } + """); + + assertFalse(response.contains("\"error\""), response); + assertTrue(response.contains("\"hello\""), response); + assertFalse(response.contains("[104,101,108,108,111]"), response); + } + + private static Connection binaryVarcharConnection() { + Statement statement = binaryVarcharStatement(); + return proxy(Connection.class, (method, args) -> { + String name = method.getName(); + if ("createStatement".equals(name)) { + return statement; + } + if ("getAutoCommit".equals(name)) { + return true; + } + if ("setAutoCommit".equals(name) || "commit".equals(name) || "rollback".equals(name) || "close".equals(name)) { + return null; + } + if ("isClosed".equals(name)) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static Statement binaryVarcharStatement() { + ResultSet resultSet = binaryVarcharResultSet(); + return proxy(Statement.class, (method, args) -> { + String name = method.getName(); + if ("execute".equals(name)) { + return true; + } + if ("getResultSet".equals(name)) { + return resultSet; + } + if ("getUpdateCount".equals(name)) { + return 0; + } + if ("setMaxRows".equals(name) || "setFetchSize".equals(name) || "close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet binaryVarcharResultSet() { + final int[] index = {-1}; + ResultSetMetaData metadata = binaryVarcharMetadata(); + byte[] value = "hello".getBytes(StandardCharsets.UTF_8); + return proxy(ResultSet.class, (method, args) -> { + String name = method.getName(); + if ("next".equals(name)) { + index[0] += 1; + return index[0] == 0; + } + if ("getMetaData".equals(name)) { + return metadata; + } + if ("getObject".equals(name)) { + return value; + } + if ("wasNull".equals(name)) { + return false; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSetMetaData binaryVarcharMetadata() { + return proxy(ResultSetMetaData.class, (method, args) -> { + String name = method.getName(); + if ("getColumnCount".equals(name)) { + return 1; + } + if ("getColumnLabel".equals(name)) { + return "v"; + } + if ("getColumnType".equals(name)) { + return Types.VARCHAR; + } + return defaultValue(method.getReturnType()); + }); + } + + private static String handleRequest(JsonRpcServer server, String request) { + try { + Method method = JsonRpcServer.class.getDeclaredMethod("handleRequest", String.class); + method.setAccessible(true); + return (String) method.invoke(server, request); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java b/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java new file mode 100644 index 0000000000..51bdc114b3 --- /dev/null +++ b/agents/drivers/tdengine/src/test/java/com/dbx/agent/tdengine/TDengineAgentTest.java @@ -0,0 +1,140 @@ +package com.dbx.agent.tdengine; + +import com.dbx.agent.test.TestSupport; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.test.JdbcAgentFake; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import com.dbx.agent.test.JdbcMetadataSqlFake; +import java.nio.charset.StandardCharsets; +import java.sql.Timestamp; +import java.util.Arrays; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TDengineAgentExecutionTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new TDengineAgent(); + } + + @Override + protected String resultSetSql() { + return "SHOW DATABASES"; + } +} + +class TDengineAgentMetadataTest { + @Test + void buildsWebsocketJdbcUrlWithDefaultPortAndDatabase() { + String url = TDengineJdbcUrl.from( + new ConnectParams("127.0.0.1", 0, "meters", "root", "taosdata", "", "", false) + ); + + Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/meters", url); + } + + @Test + void preservesCustomWebsocketPortAndUrlParams() { + String url = TDengineJdbcUrl.from( + new ConnectParams("td.local", 6042, "", "", "", "timezone=UTC&charset=UTF-8", "", false) + ); + + Assertions.assertEquals("jdbc:TAOS-WS://td.local:6042/?timezone=UTC&charset=UTF-8", url); + } + + @Test + void supportsRestTransportViaCompatibilityUrlParam() { + String url = TDengineJdbcUrl.from( + new ConnectParams("127.0.0.1", 0, "testdb", "", "", "dbx.transport=rest&charset=UTF-8", "", false) + ); + + Assertions.assertEquals("jdbc:TAOS-RS://127.0.0.1:6041/testdb?charset=UTF-8", url); + } + + @Test + void stripsTransportControlParamFromJdbcQuery() { + String url = TDengineJdbcUrl.from( + new ConnectParams("127.0.0.1", 6041, "", "", "", "transport=ws&timezone=UTC", "", false) + ); + + Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/?timezone=UTC", url); + } + + @Test + void usesTdengineMetadataStatements() { + TDengineAgent agent = new TDengineAgent(); + TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection()); + + agent.listDatabases(); + agent.listTables("power"); + agent.getColumns("power", "meters"); + + Assertions.assertEquals( + Arrays.asList( + "SHOW DATABASES", + "SHOW `power`.STABLES", + "SHOW `power`.TABLES", + "DESCRIBE `power`.`meters`" + ), + JdbcMetadataSqlFake.statements + ); + } + + @Test + void setsDatabaseBeforeExecutionWhenSchemaIsProvided() { + TDengineAgent agent = new TDengineAgent(); + TestSupport.setPrivateConnection(agent, JdbcAgentFake.connection()); + + agent.executeQuery("SELECT 1", "power", new ExecuteQueryOptions()); + + Assertions.assertEquals( + Arrays.asList("setMaxRows:10001", "execute"), + JdbcAgentFake.calls + ); + } + + @Test + void decodesTdengineByteArrayTextValues() { + Assertions.assertEquals( + "d1001", + TDengineAgent.decodeTdengineValue("d1001".getBytes(StandardCharsets.UTF_8)) + ); + } + + @Test + void formatsTdengineTimestampsAsSqlLiterals() { + Assertions.assertEquals( + "2026-05-16 09:35:58.123", + TDengineAgent.decodeTdengineValue(Timestamp.valueOf("2026-05-16 09:35:58.123")) + ); + } + + @Test + void unknownTransportValueDefaultsToWebsocketAndKeepsOtherParams() { + String url = TDengineJdbcUrl.from( + new ConnectParams("127.0.0.1", 6041, "", "", "", "transport=foo&timezone=UTC", "", false) + ); + + Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/?timezone=UTC", url); + } + + @Test + void sanitizeConnectionStringStripsTransportControlParams() { + String sanitized = TDengineJdbcUrl.sanitizeConnectionString( + "jdbc:TAOS-WS://127.0.0.1:6041/db?dbx.transport=rest&charset=UTF-8&transport=ws" + ); + + Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/db?charset=UTF-8", sanitized); + } + + @Test + void sanitizeConnectionStringKeepsFragmentAndNonControlParams() { + String sanitized = TDengineJdbcUrl.sanitizeConnectionString( + "jdbc:TAOS-RS://127.0.0.1:6041/db?timezone=UTC&dbx.transport=rest#anchor" + ); + + Assertions.assertEquals("jdbc:TAOS-RS://127.0.0.1:6041/db?timezone=UTC#anchor", sanitized); + } +} diff --git a/agents/drivers/teradata/build.gradle b/agents/drivers/teradata/build.gradle new file mode 100644 index 0000000000..73c6bad937 --- /dev/null +++ b/agents/drivers/teradata/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.teradata.jdbc:terajdbc:20.00.00.56' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Teradata', 'Main-Class': 'com.dbx.agent.teradata.TeradataAgent') + } +} diff --git a/agents/drivers/teradata/src/main/java/com/dbx/agent/teradata/TeradataAgent.java b/agents/drivers/teradata/src/main/java/com/dbx/agent/teradata/TeradataAgent.java new file mode 100644 index 0000000000..f24aafcff5 --- /dev/null +++ b/agents/drivers/teradata/src/main/java/com/dbx/agent/teradata/TeradataAgent.java @@ -0,0 +1,21 @@ +package com.dbx.agent.teradata; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +public final class TeradataAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile TERADATA_PROFILE = new JdbcAgentProfile( + "com.teradata.jdbc.TeraDriver", + "jdbc:teradata://{host}/DBS_PORT={port},DATABASE={database}", + 1025 + ); + + public TeradataAgent() { + super(TERADATA_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new TeradataAgent()).run(); + } +} diff --git a/agents/drivers/trino/build.gradle b/agents/drivers/trino/build.gradle new file mode 100644 index 0000000000..28b6cd8146 --- /dev/null +++ b/agents/drivers/trino/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'io.trino:trino-jdbc:471' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Trino (Presto)', 'Main-Class': 'com.dbx.agent.trino.TrinoAgent') + } +} diff --git a/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java b/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java new file mode 100644 index 0000000000..6ffdad8ace --- /dev/null +++ b/agents/drivers/trino/src/main/java/com/dbx/agent/trino/TrinoAgent.java @@ -0,0 +1,272 @@ +package com.dbx.agent.trino; + +import com.dbx.agent.AbstractJdbcAgent; +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.ConnectParams; +import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ForeignKeyInfo; +import com.dbx.agent.IndexInfo; +import com.dbx.agent.JdbcIdentifiers; +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.TableInfo; +import com.dbx.agent.TriggerInfo; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSet; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +public final class TrinoAgent extends AbstractJdbcAgent { + private static final Pattern NUMERIC_PRECISION_PATTERN = Pattern.compile("(?i)^(decimal|numeric)\\((\\d+)(?:,\\s*\\d+)?\\)"); + private static final Pattern NUMERIC_SCALE_PATTERN = Pattern.compile("(?i)^(decimal|numeric)\\(\\d+,\\s*(\\d+)\\)"); + private static final Pattern CHARACTER_LENGTH_PATTERN = Pattern.compile("(?i)^(char|varchar)\\((\\d+)\\)"); + + @Override + protected String driverClass() { + return "io.trino.jdbc.TrinoDriver"; + } + + @Override + protected String buildJdbcUrl(ConnectParams params) { + return buildUrl(params); + } + + @Override + public List listDatabases() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW CATALOGS")) { + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString(1))); + } + } + return result; + }); + } + + @Override + public List listSchemas() { + return unchecked(() -> { + List result = new ArrayList<>(); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW SCHEMAS")) { + while (rs.next()) { + result.add(rs.getString(1)); + } + } + return result; + }); + } + + @Override + public List listTables(String schema) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT table_name, table_type " + + "FROM information_schema.tables " + + "WHERE table_schema = ? " + + "ORDER BY table_name"; + try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, schema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("table_name"), + normalizeTableType(rs.getString("table_type")), + null + )); + } + } + } + return result; + }); + } + + @Override + public List getColumns(String schema, String table) { + Connection conn = requireConnected(); + List metadataColumns; + try { + metadataColumns = getColumnsFromMetadata(conn, schema, table); + } catch (RuntimeException e) { + metadataColumns = Collections.emptyList(); + } + return metadataColumns.isEmpty() ? getColumnsFromDescribe(conn, schema, table) : metadataColumns; + } + + @Override + public List listIndexes(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listForeignKeys(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public List listTriggers(String schema, String table) { + return Collections.emptyList(); + } + + @Override + public String setSchemaSQL(String schema) { + return "USE " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); + } + + private List getColumnsFromMetadata(Connection conn, String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String catalog = blankToNull(conn.getCatalog()); + try (ResultSet rs = conn.getMetaData().getColumns(catalog, schema, table, null)) { + while (rs.next()) { + int sqlType = rs.getInt("DATA_TYPE"); + Integer size = intOrNull(rs, "COLUMN_SIZE"); + Integer scale = intOrNull(rs, "DECIMAL_DIGITS"); + result.add(new ColumnInfo( + rs.getString("COLUMN_NAME"), + coalesce(rs.getString("TYPE_NAME")), + rs.getInt("NULLABLE") != DatabaseMetaData.columnNoNulls, + rs.getString("COLUMN_DEF"), + false, + null, + rs.getString("REMARKS"), + isNumericType(sqlType) ? size : null, + isNumericType(sqlType) ? scale : null, + isCharacterType(sqlType) ? size : null + )); + } + } + return result; + }); + } + + private List getColumnsFromDescribe(Connection conn, String schema, String table) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "DESCRIBE " + quoteIdentifier(schema) + "." + quoteIdentifier(table); + try (java.sql.Statement stmt = conn.createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + String dataType = rs.getString(2); + String extra = rs.getString(3); + result.add(new ColumnInfo( + rs.getString(1), + dataType, + !coalesce(extra).toLowerCase(java.util.Locale.ROOT).contains("not null"), + null, + false, + extra, + rs.getString(4), + parseNumericPrecision(dataType), + parseNumericScale(dataType), + parseCharacterMaximumLength(dataType) + )); + } + } + return result; + }); + } + + @Override + protected Object resultValue(ResultSet rs, int index, int sqlType) { + return unchecked(() -> { + Object value = rs.getObject(index); + return rs.wasNull() ? null : value == null ? null : value.toString(); + }); + } + + private static String buildUrl(ConnectParams params) { + if (params.getConnection_string() != null && !params.getConnection_string().isBlank()) { + return params.getConnection_string(); + } + return "jdbc:trino://" + params.getHost() + ":" + params.getPort() + "/" + params.getDatabase(); + } + + private static String quoteIdentifier(String identifier) { + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } + + private static boolean isNumericType(int sqlType) { + switch (sqlType) { + case Types.BIGINT: + case Types.DECIMAL: + case Types.DOUBLE: + case Types.FLOAT: + case Types.INTEGER: + case Types.NUMERIC: + case Types.REAL: + case Types.SMALLINT: + case Types.TINYINT: + return true; + default: + return false; + } + } + + private static boolean isCharacterType(int sqlType) { + switch (sqlType) { + case Types.CHAR: + case Types.LONGNVARCHAR: + case Types.LONGVARCHAR: + case Types.NCHAR: + case Types.NVARCHAR: + case Types.VARCHAR: + return true; + default: + return false; + } + } + + private static Integer parseNumericPrecision(String dataType) { + return firstIntGroup(NUMERIC_PRECISION_PATTERN, dataType, 2); + } + + private static Integer parseNumericScale(String dataType) { + return firstIntGroup(NUMERIC_SCALE_PATTERN, dataType, 1); + } + + private static Integer parseCharacterMaximumLength(String dataType) { + return firstIntGroup(CHARACTER_LENGTH_PATTERN, dataType, 2); + } + + private static Integer firstIntGroup(Pattern pattern, String value, int group) { + if (value == null) { + return null; + } + Matcher matcher = pattern.matcher(value); + if (!matcher.find()) { + return null; + } + return Integer.valueOf(matcher.group(group)); + } + + private static Integer intOrNull(ResultSet rs, String column) throws Exception { + Object value = rs.getObject(column); + return value instanceof Number ? ((Number) value).intValue() : null; + } + + private static String normalizeTableType(String type) { + if ("BASE TABLE".equals(type)) { + return "TABLE"; + } + return type; + } + + private static String blankToNull(String value) { + return value == null || value.trim().isEmpty() ? null : value; + } + + private static String coalesce(String value) { + return value == null ? "" : value; + } + + public static void main(String[] args) { + new JsonRpcServer(new TrinoAgent()).run(); + } +} diff --git a/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java b/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java new file mode 100644 index 0000000000..f1d94db0f6 --- /dev/null +++ b/agents/drivers/trino/src/test/java/com/dbx/agent/trino/TrinoAgentTest.java @@ -0,0 +1,16 @@ +package com.dbx.agent.trino; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; + +class TrinoAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new TrinoAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL system.runtime.nodes()"; + } +} diff --git a/agents/drivers/vastbase/build.gradle b/agents/drivers/vastbase/build.gradle new file mode 100644 index 0000000000..11ece4ac12 --- /dev/null +++ b/agents/drivers/vastbase/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'cn.com.vastdata:vastbase-jdbc:2.11v' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Vastbase', 'Main-Class': 'com.dbx.agent.vastbase.VastbaseAgent') + } +} diff --git a/agents/drivers/vastbase/libs/.gitkeep b/agents/drivers/vastbase/libs/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/agents/drivers/vastbase/src/main/java/com/dbx/agent/vastbase/VastbaseAgent.java b/agents/drivers/vastbase/src/main/java/com/dbx/agent/vastbase/VastbaseAgent.java new file mode 100644 index 0000000000..920db9b16a --- /dev/null +++ b/agents/drivers/vastbase/src/main/java/com/dbx/agent/vastbase/VastbaseAgent.java @@ -0,0 +1,61 @@ +package com.dbx.agent.vastbase; + +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.PostgresLikeAgent; +import com.dbx.agent.PostgresLikeAgentProfile; + +public final class VastbaseAgent extends PostgresLikeAgent { + public static final PostgresLikeAgentProfile VASTBASE_PROFILE = new PostgresLikeAgentProfile( + "cn.com.vastbase.Driver", + "jdbc:vastbase://{host}:{port}/{database}" + ); + + public VastbaseAgent() { + super(VASTBASE_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new VastbaseAgent()).run(); + } + + @Override + public ObjectSource getObjectSource(String schema, String name, String objectType) { + ObjectSource result = super.getObjectSource(schema, name, objectType); + String source = unwrapRecordText(result.getSource()); + return new ObjectSource(result.getName(), result.getObject_type(), result.getSchema(), source); + } + + /** + * Vastbase 的 pg_get_functiondef 在 SQL Server 兼容模式下返回的是 + * PostgreSQL 行记录文本格式 {@code (1,"source text")} 而非纯文本。 + * 此方法剥离行记录包装，提取实际的源码内容。 + */ + static String unwrapRecordText(String source) { + if (source == null || source.isEmpty()) { + return source; + } + String trimmed = source.trim(); + if (!trimmed.startsWith("(") || !trimmed.endsWith(")")) { + return source; + } + // 格式: (field1,field2,...)，第一个字段通常是整数，第二个是源码文本 + String inner = trimmed.substring(1, trimmed.length() - 1); + int commaIdx = inner.indexOf(','); + if (commaIdx <= 0) { + // 单字段行记录: (source_text) + return unquoteRecordField(inner); + } + // 多字段行记录: 跳过第一个字段，提取第二个字段及之后的内容 + String rest = inner.substring(commaIdx + 1); + return unquoteRecordField(rest); + } + + private static String unquoteRecordField(String value) { + String trimmed = value.trim(); + if (trimmed.startsWith("\"") && trimmed.endsWith("\"")) { + return trimmed.substring(1, trimmed.length() - 1).replace("\"\"", "\""); + } + return trimmed; + } +} diff --git a/agents/drivers/vastbase/src/test/java/com/dbx/agent/vastbase/VastbaseAgentTest.java b/agents/drivers/vastbase/src/test/java/com/dbx/agent/vastbase/VastbaseAgentTest.java new file mode 100644 index 0000000000..ab8cb9a55e --- /dev/null +++ b/agents/drivers/vastbase/src/test/java/com/dbx/agent/vastbase/VastbaseAgentTest.java @@ -0,0 +1,26 @@ +package com.dbx.agent.vastbase; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class VastbaseAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new VastbaseAgent(); + } + + @Override + protected String resultSetSql() { + return "CALL sample_proc()"; + } + + @Test + void declaresVastbasePostgresLikeProfile() { + VastbaseAgent agent = new VastbaseAgent(); + + Assertions.assertEquals("cn.com.vastbase.Driver", agent.getProfile().getDriverClass()); + Assertions.assertEquals("jdbc:vastbase://{host}:{port}/{database}", agent.getProfile().getUrlTemplate()); + } +} diff --git a/agents/drivers/vertica/build.gradle b/agents/drivers/vertica/build.gradle new file mode 100644 index 0000000000..fc1e0cf03e --- /dev/null +++ b/agents/drivers/vertica/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'com.vertica.jdbc:vertica-jdbc:25.3.0-0' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'Vertica', 'Main-Class': 'com.dbx.agent.vertica.VerticaAgent') + } +} diff --git a/agents/drivers/vertica/src/main/java/com/dbx/agent/vertica/VerticaAgent.java b/agents/drivers/vertica/src/main/java/com/dbx/agent/vertica/VerticaAgent.java new file mode 100644 index 0000000000..05cef7a2b5 --- /dev/null +++ b/agents/drivers/vertica/src/main/java/com/dbx/agent/vertica/VerticaAgent.java @@ -0,0 +1,21 @@ +package com.dbx.agent.vertica; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +public final class VerticaAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile VERTICA_PROFILE = new JdbcAgentProfile( + "com.vertica.jdbc.Driver", + "jdbc:vertica://{host}:{port}/{database}", + 5433 + ); + + public VerticaAgent() { + super(VERTICA_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new VerticaAgent()).run(); + } +} diff --git a/agents/drivers/xugu/README.md b/agents/drivers/xugu/README.md new file mode 100644 index 0000000000..09631c301a --- /dev/null +++ b/agents/drivers/xugu/README.md @@ -0,0 +1,36 @@ +# DBX Xugu Agent + +Native XuguDB agent for DBX using `gitee.com/XuguDB/go-xugu-driver`. + +## Build + +```bash +go build -o agent . +``` + +Cross-compile release builds use pure Go output: + +```bash +CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o dbx-agent-xugu-linux-x64 . +CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 go build -trimpath -ldflags="-s -w" -o dbx-agent-xugu-macos-aarch64 . +CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o dbx-agent-xugu-windows-x64.exe . +``` + +## Local DBX Test + +Build the binary, then copy it into DBX's installed XuguDB driver directory: + +```bash +mkdir -p ~/.dbx/agents/drivers/xugu +cp agent ~/.dbx/agents/drivers/xugu/agent +chmod +x ~/.dbx/agents/drivers/xugu/agent +``` + +DBX prefers `agent` over `agent.jar`, so XuguDB connections will use this Go +agent until the file is removed. + +To restore the Java/JDBC agent: + +```bash +rm ~/.dbx/agents/drivers/xugu/agent +``` diff --git a/agents/drivers/xugu/go.mod b/agents/drivers/xugu/go.mod new file mode 100644 index 0000000000..acc70480f7 --- /dev/null +++ b/agents/drivers/xugu/go.mod @@ -0,0 +1,5 @@ +module github.com/t8y2/dbx/agents/drivers/xugu + +go 1.22 + +require gitee.com/XuguDB/go-xugu-driver v1.0.12 diff --git a/agents/drivers/xugu/go.sum b/agents/drivers/xugu/go.sum new file mode 100644 index 0000000000..7b5f9a2f7c --- /dev/null +++ b/agents/drivers/xugu/go.sum @@ -0,0 +1,2 @@ +gitee.com/XuguDB/go-xugu-driver v1.0.12 h1:Le49iCI5YJBmQDJUUWx2SHHR3aYio0CTFjNx5O8m2Ew= +gitee.com/XuguDB/go-xugu-driver v1.0.12/go.mod h1:VJlSyp+6BjYL1US00cVb4traVISlm8iGywYe3viLPbE= diff --git a/agents/drivers/xugu/main.go b/agents/drivers/xugu/main.go new file mode 100644 index 0000000000..f3af1f39ec --- /dev/null +++ b/agents/drivers/xugu/main.go @@ -0,0 +1,1714 @@ +package main + +import ( + "bufio" + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "regexp" + "strconv" + "strings" + "time" + + _ "gitee.com/XuguDB/go-xugu-driver" +) + +const protocolVersion = 1 +const defaultMaxRows = 1000 +const defaultXuguPort = 5138 +const xuguListDatabasesSQL = ` +SELECT DB_NAME +FROM ALL_DATABASES +ORDER BY DB_NAME` +const xuguListSchemasSQL = ` +SELECT SCHEMA_NAME +FROM ALL_SCHEMAS +ORDER BY SCHEMA_NAME` +const xuguPrimaryKeyColumnsSQL = ` +SELECT c.DEFINE +FROM ALL_CONSTRAINTS c +JOIN ALL_TABLES t ON t.DB_ID = c.DB_ID AND t.TABLE_ID = c.TABLE_ID +JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) + AND UPPER(t.TABLE_NAME) = UPPER(?) + AND c.CONS_TYPE = 'P'` + +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 rpcError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +type connectParams struct { + Host string `json:"host"` + Port int `json:"port"` + Database string `json:"database"` + Username string `json:"username"` + Password string `json:"password"` + SysDBA bool `json:"sysdba"` + URLParams string `json:"url_params"` + ConnectionString string `json:"connection_string"` +} + +type queryOptions struct { + SQL string `json:"sql"` + Database string `json:"database"` + Schema string `json:"schema"` + MaxRows int `json:"maxRows"` + FetchSize int `json:"fetchSize"` +} + +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"` +} + +func (r queryResult) MarshalJSON() ([]byte, error) { + type alias queryResult + value := alias(r) + if value.Columns == nil { + value.Columns = []string{} + } + if value.ColumnTypes == nil { + value.ColumnTypes = []string{} + } + if value.Rows == nil { + value.Rows = [][]any{} + } + return json.Marshal(value) +} + +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"` +} + +func (r queryPageResult) MarshalJSON() ([]byte, error) { + type alias queryPageResult + value := alias(r) + if value.Columns == nil { + value.Columns = []string{} + } + if value.ColumnTypes == nil { + value.ColumnTypes = []string{} + } + if value.Rows == nil { + value.Rows = [][]any{} + } + return json.Marshal(value) +} + +type querySession struct { + rows *sql.Rows + columns []string + columnTypes []string + pending []any + remaining int +} + +type databaseInfo struct { + Name string `json:"name"` +} + +type tableInfo struct { + Name string `json:"name"` + TableType string `json:"table_type"` + Comment *string `json:"comment"` +} + +type objectInfo struct { + Name string `json:"name"` + ObjectType string `json:"object_type"` + Schema string `json:"schema"` + Comment *string `json:"comment"` +} + +type columnInfo struct { + Name string `json:"name"` + DataType string `json:"data_type"` + IsNullable bool `json:"is_nullable"` + ColumnDefault *string `json:"column_default"` + IsPrimaryKey bool `json:"is_primary_key"` + Extra *string `json:"extra"` + Comment *string `json:"comment"` + NumericPrecision *int `json:"numeric_precision"` + NumericScale *int `json:"numeric_scale"` + CharacterMaximumLength *int `json:"character_maximum_length"` +} + +type indexInfo struct { + Name string `json:"name"` + Columns []string `json:"columns"` + IsUnique bool `json:"is_unique"` + IsPrimary bool `json:"is_primary"` + Filter *string `json:"filter"` + IndexType *string `json:"index_type"` + IncludedColumns []string `json:"included_columns"` + Comment *string `json:"comment"` +} + +func (i indexInfo) MarshalJSON() ([]byte, error) { + type alias indexInfo + value := alias(i) + if value.Columns == nil { + value.Columns = []string{} + } + if value.IncludedColumns == nil { + value.IncludedColumns = []string{} + } + return json.Marshal(value) +} + +type foreignKeyInfo struct { + Name string `json:"name"` + Column string `json:"column"` + RefTable string `json:"ref_table"` + RefColumn string `json:"ref_column"` +} + +type triggerInfo struct { + Name string `json:"name"` + Event string `json:"event"` + Timing string `json:"timing"` +} + +type server struct { + db *sql.DB + params connectParams + sessions map[string]*querySession + nextSessionID int64 +} + +func main() { + s := newServer() + encoder := json.NewEncoder(os.Stdout) + 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 + } + resp, shutdown := s.handleLine(line) + if err := encoder.Encode(resp); err != nil { + fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err) + return + } + if shutdown { + return + } + } + if err := scanner.Err(); err != nil && !errors.Is(err, io.EOF) { + fmt.Fprintf(os.Stderr, "failed to read stdin: %v\n", err) + } +} + +func newServer() *server { + return &server{sessions: map[string]*querySession{}} +} + +func (s *server) 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 := s.dispatch(req.Method, req.Params) + if err != nil { + return errorResponse(req.ID, err), false + } + return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown +} + +func (s *server) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) { + switch method { + case "handshake": + return map[string]any{ + "protocolVersion": protocolVersion, + "agentProtocolVersion": protocolVersion, + "capabilities": []string{"connect", "test_connection", "metadata", "query", "ddl"}, + }, false, nil + case "connect": + var cp connectParams + if err := decodeParams(params, &cp); err != nil { + return nil, false, err + } + return map[string]bool{"ok": true}, false, s.connect(cp) + case "test_connection": + var cp connectParams + if err := decodeParams(params, &cp); err != nil { + return nil, false, err + } + db, err := openDB(cp) + if err != nil { + return nil, false, err + } + defer db.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := db.PingContext(ctx); err != nil { + return nil, false, err + } + return map[string]bool{"ok": true}, false, nil + case "list_databases": + result, err := s.listDatabases() + return result, false, err + case "list_schemas": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + result, err := s.listSchemas() + return result, false, err + case "list_tables": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + result, err := s.listTables(schema) + return result, false, err + case "list_objects": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + result, err := s.listObjects(schema) + return result, false, err + case "get_columns": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.getColumns(schema, table) + return result, false, err + case "get_object_source": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + name := stringParam(params, "name") + objectType := stringParam(params, "object_type") + source, err := s.getObjectSource(schema, name, objectType) + return source, false, err + case "get_table_ddl": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + table := stringParam(params, "table") + ddl, err := s.getTableDDL(schema, table) + return ddl, false, err + case "execute_query": + var opts queryOptions + if err := decodeParams(params, &opts); err != nil { + return nil, false, err + } + result, err := s.executeQuery(opts) + return result, false, err + case "execute_query_page": + var opts queryOptions + if err := decodeParams(params, &opts); err != nil { + return nil, false, err + } + result, err := s.executeQueryPage(opts, intParam(params, "pageSize")) + return result, false, err + case "fetch_query_page": + result, err := s.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize")) + return result, false, err + case "close_query_session": + return s.closeQuerySession(stringParam(params, "sessionId")), false, nil + case "list_indexes": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.listIndexes(schema, table) + return result, false, err + case "list_foreign_keys": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.listForeignKeys(schema, table) + return result, false, err + case "list_triggers": + if err := s.useDatabase(stringParam(params, "database")); err != nil { + return nil, false, err + } + schema := stringParam(params, "schema") + table := stringParam(params, "table") + result, err := s.listTriggers(schema, table) + return result, false, err + case "get_explain_info": + sqlText := stringParam(params, "sql") + plan, err := s.getExplainInfo(sqlText) + return map[string]any{"plan": plan, "has_actual_stats": false}, false, err + case "execute_transaction": + result, err := s.executeTransaction(params) + return result, false, err + case "disconnect": + return map[string]bool{"ok": true}, false, s.disconnect() + case "shutdown": + _ = s.disconnect() + return map[string]bool{"ok": true}, true, nil + default: + return nil, false, fmt.Errorf("unknown method: %s", method) + } +} + +func (s *server) connect(params connectParams) error { + _ = s.disconnect() + db, err := openDB(params) + if err != nil { + return err + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := db.PingContext(ctx); err != nil { + db.Close() + return err + } + s.db = db + s.params = params + return nil +} + +func (s *server) disconnect() error { + s.closeAllQuerySessions() + if s.db == nil { + return nil + } + err := s.db.Close() + s.db = nil + return err +} + +func openDB(params connectParams) (*sql.DB, error) { + dsn := buildDSN(params) + db, err := sql.Open("xugu", dsn) + if err != nil { + return nil, err + } + db.SetMaxOpenConns(4) + db.SetMaxIdleConns(1) + db.SetConnMaxLifetime(30 * time.Minute) + return db, nil +} + +func buildDSN(params connectParams) string { + connectionString := strings.TrimSpace(params.ConnectionString) + if looksLikeXuguDSN(connectionString) { + return connectionString + } + if parsed := parseXuguURL(connectionString); parsed.Host != "" { + if parsed.Database == "" { + parsed.Database = params.Database + } + if parsed.Username == "" { + parsed.Username = params.Username + } + if parsed.Password == "" { + parsed.Password = params.Password + } + return buildXuguDSN(parsed.Host, parsed.Port, parsed.Database, parsed.Username, parsed.Password, params.URLParams) + } + + if jdbc := parseXuguJDBCURL(connectionString); jdbc.Host != "" { + return buildXuguDSN(jdbc.Host, jdbc.Port, jdbc.Database, params.Username, params.Password, params.URLParams) + } + + return buildXuguDSN(params.Host, params.Port, params.Database, params.Username, params.Password, params.URLParams) +} + +func looksLikeXuguDSN(value string) bool { + upper := strings.ToUpper(value) + return strings.Contains(upper, "IP=") && strings.Contains(upper, "DB=") && strings.Contains(upper, "USER=") +} + +type xuguURLInfo struct { + Host string + Port int + Database string + Username string + Password string +} + +var xuguJDBCRegexp = regexp.MustCompile(`(?i)^jdbc:xugu://([^/:]+)(?::([0-9]+))?/([^?;]+)`) + +func parseXuguJDBCURL(value string) xuguURLInfo { + value = strings.TrimSpace(value) + match := xuguJDBCRegexp.FindStringSubmatch(value) + if len(match) != 4 { + return xuguURLInfo{} + } + return xuguURLInfo{Host: match[1], Port: parsePort(match[2]), Database: match[3]} +} + +func parseXuguURL(value string) xuguURLInfo { + value = strings.TrimSpace(value) + if !strings.HasPrefix(strings.ToLower(value), "xugu://") { + return xuguURLInfo{} + } + withoutScheme := value[len("xugu://"):] + var userInfo string + hostPart := withoutScheme + if at := strings.LastIndex(hostPart, "@"); at >= 0 { + userInfo = hostPart[:at] + hostPart = hostPart[at+1:] + } + if slash := strings.IndexAny(hostPart, "/?"); slash >= 0 { + databasePart := strings.TrimLeft(hostPart[slash:], "/") + hostPart = hostPart[:slash] + if q := strings.Index(databasePart, "?"); q >= 0 { + databasePart = databasePart[:q] + } + info := parseHostPort(hostPart) + info.Database = databasePart + if userInfo != "" { + info.Username, info.Password = splitUserInfo(userInfo) + } + return info + } + info := parseHostPort(hostPart) + if userInfo != "" { + info.Username, info.Password = splitUserInfo(userInfo) + } + return info +} + +func parseHostPort(value string) xuguURLInfo { + host := strings.TrimSpace(value) + port := 0 + if idx := strings.LastIndex(host, ":"); idx > 0 { + port = parsePort(host[idx+1:]) + host = host[:idx] + } + return xuguURLInfo{Host: host, Port: port} +} + +func splitUserInfo(value string) (string, string) { + if idx := strings.Index(value, ":"); idx >= 0 { + return value[:idx], value[idx+1:] + } + return value, "" +} + +func buildXuguDSN(host string, port int, database, username, password, urlParams string) string { + if port <= 0 { + port = defaultXuguPort + } + parts := []string{ + "IP=" + host, + "DB=" + strings.TrimSpace(database), + "User=" + strings.TrimSpace(username), + "PWD=" + strings.TrimSpace(password), + "Port=" + strconv.Itoa(port), + } + if !hasDSNParam(urlParams, "CHAR_SET") && !hasDSNParam(urlParams, "CHARSET") { + parts = append(parts, "CHAR_SET=UTF8") + } + for _, param := range splitDSNParams(urlParams) { + parts = append(parts, param) + } + return strings.Join(parts, ";") +} + +func hasDSNParam(raw, key string) bool { + key = strings.ToUpper(strings.TrimSpace(key)) + for _, param := range splitDSNParams(raw) { + name, _, _ := strings.Cut(param, "=") + if strings.ToUpper(strings.TrimSpace(name)) == key { + return true + } + } + return false +} + +func splitDSNParams(raw string) []string { + raw = strings.TrimSpace(strings.Trim(raw, ";")) + if raw == "" { + return nil + } + if strings.Contains(raw, ";") { + items := strings.Split(raw, ";") + result := make([]string, 0, len(items)) + for _, item := range items { + item = strings.TrimSpace(item) + if item != "" { + result = append(result, item) + } + } + return result + } + if strings.Contains(raw, "&") { + items := strings.Split(raw, "&") + result := make([]string, 0, len(items)) + for _, item := range items { + item = strings.TrimSpace(item) + if item != "" { + result = append(result, item) + } + } + return result + } + return []string{raw} +} + +func parsePort(value string) int { + port, _ := strconv.Atoi(value) + return port +} + +func (s *server) requireDB() (*sql.DB, error) { + if s.db == nil { + return nil, errors.New("agent is not connected") + } + return s.db, nil +} + +func (s *server) useDatabase(database string) error { + database = strings.TrimSpace(database) + if database == "" { + return nil + } + if strings.EqualFold(database, configuredDatabaseName(s.params)) { + return nil + } + db, err := s.requireDB() + if err != nil { + return err + } + _, err = db.Exec("USE " + quoteIdentifier(database)) + return err +} + +func (s *server) listDatabases() ([]databaseInfo, error) { + rows, err := s.queryRows(xuguListDatabasesSQL, nil) + if err != nil { + if fallback := fallbackDatabasesFromParams(s.params); len(fallback) > 0 && isXuguMetadataAccessError(err) { + return fallback, nil + } + return nil, err + } + defer rows.Close() + var result []databaseInfo + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return nil, err + } + result = append(result, databaseInfo{Name: name}) + } + return emptyIfNil(result), rows.Err() +} + +func fallbackDatabasesFromParams(params connectParams) []databaseInfo { + if name := configuredDatabaseName(params); name != "" { + return []databaseInfo{{Name: name}} + } + return nil +} + +func configuredDatabaseName(params connectParams) string { + if name := strings.TrimSpace(params.Database); name != "" { + return name + } + connectionString := strings.TrimSpace(params.ConnectionString) + if parsed := parseXuguURL(connectionString); parsed.Database != "" { + return parsed.Database + } + if jdbc := parseXuguJDBCURL(connectionString); jdbc.Database != "" { + return jdbc.Database + } + if value := xuguDSNValue(connectionString, "DB"); value != "" { + return value + } + return "" +} + +func isXuguMetadataAccessError(err error) bool { + message := strings.ToUpper(err.Error()) + return strings.Contains(message, "E18012") || + strings.Contains(message, "权限不够") || + strings.Contains(message, "ALL_DATABASES") || + strings.Contains(message, "SYS_DATABASES") || + strings.Contains(message, "ALL_SCHEMAS") || + strings.Contains(message, "SYS_SCHEMAS") || + strings.Contains(message, "ALL_TABLES") || + strings.Contains(message, "SYS_TABLES") || + strings.Contains(message, "ALL_VIEWS") || + strings.Contains(message, "SYS_VIEWS") || + strings.Contains(message, "ALL_CONSTRAINTS") || + strings.Contains(message, "SYS_COLUMNS") || + strings.Contains(message, "SYS_CONSTRAINTS") || + strings.Contains(message, "SYS_INDEXES") || + strings.Contains(message, "SYS_TRIGGERS") +} + +func xuguDSNValue(dsn string, key string) string { + for _, part := range strings.Split(dsn, ";") { + name, value, ok := strings.Cut(part, "=") + if ok && strings.EqualFold(strings.TrimSpace(name), key) { + return strings.TrimSpace(value) + } + } + return "" +} + +func (s *server) listSchemas() ([]string, error) { + rows, err := s.queryRows(xuguListSchemasSQL, nil) + if err != nil { + if fallback := strings.ToUpper(strings.TrimSpace(s.params.Username)); fallback != "" && isXuguMetadataAccessError(err) { + return []string{fallback}, nil + } + return nil, err + } + defer rows.Close() + var result []string + for rows.Next() { + var schema string + if err := rows.Scan(&schema); err != nil { + return nil, err + } + result = append(result, schema) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) currentSchema() (string, error) { + rows, err := s.queryRows(` +SELECT s.SCHEMA_NAME +FROM SYS_SCHEMAS s +JOIN SYS_USERS u ON u.DB_ID = s.DB_ID AND u.USER_ID = s.USER_ID +WHERE UPPER(u.USER_NAME) = UPPER(?) +ORDER BY CASE WHEN UPPER(s.SCHEMA_NAME) = UPPER(?) THEN 0 ELSE 1 END, s.SCHEMA_NAME`, []any{s.params.Username, s.params.Username}) + if err != nil { + if fallback := strings.ToUpper(strings.TrimSpace(s.params.Username)); fallback != "" && isXuguMetadataAccessError(err) { + return fallback, nil + } + return "", err + } + defer rows.Close() + if rows.Next() { + var schema string + if err := rows.Scan(&schema); err != nil { + return "", err + } + return strings.ToUpper(schema), nil + } + if err := rows.Err(); err != nil { + return "", err + } + return strings.ToUpper(strings.TrimSpace(s.params.Username)), nil +} + +func (s *server) normalizeSchema(schema string) (string, error) { + schema = strings.TrimSpace(schema) + if schema == "" { + return s.currentSchema() + } + return strings.ToUpper(schema), nil +} + +func (s *server) listTables(schema string) ([]tableInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + rows, err := s.queryRows(` +SELECT t.TABLE_NAME, 'TABLE' AS TABLE_TYPE, t.COMMENTS +FROM ALL_TABLES t +JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) +UNION ALL +SELECT v.VIEW_NAME, 'VIEW' AS TABLE_TYPE, v.COMMENTS +FROM ALL_VIEWS v +JOIN ALL_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) +ORDER BY 2, 1`, []any{schema, schema}) + if err != nil { + return nil, err + } + defer rows.Close() + var result []tableInfo + for rows.Next() { + var item tableInfo + if err := rows.Scan(&item.Name, &item.TableType, &item.Comment); err != nil { + return nil, err + } + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) listObjects(schema string) ([]objectInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + rows, err := s.queryRows(` +SELECT t.TABLE_NAME, 'TABLE' AS OBJECT_TYPE, t.COMMENTS +FROM ALL_TABLES t +JOIN ALL_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) +UNION ALL +SELECT v.VIEW_NAME, 'VIEW' AS OBJECT_TYPE, v.COMMENTS +FROM ALL_VIEWS v +JOIN ALL_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) +ORDER BY 2, 1`, []any{schema, schema}) + if err != nil { + return nil, err + } + defer rows.Close() + var result []objectInfo + for rows.Next() { + var item objectInfo + item.Schema = schema + if err := rows.Scan(&item.Name, &item.ObjectType, &item.Comment); err != nil { + return nil, err + } + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) getColumns(schema, table string) ([]columnInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + primaryKeys, err := s.primaryKeyColumns(schema, table) + if err != nil { + return nil, err + } + rows, err := s.queryRows(` +SELECT c.COL_NAME, c.TYPE_NAME, c.NOT_NULL, c.DEF_VAL, c.COMMENTS, c.SCALE +FROM SYS_COLUMNS c +JOIN SYS_TABLES t ON t.DB_ID = c.DB_ID AND t.TABLE_ID = c.TABLE_ID +JOIN SYS_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) + AND UPPER(t.TABLE_NAME) = UPPER(?) + AND (c.IS_HIDE IS NULL OR c.IS_HIDE = FALSE) +ORDER BY c.COL_NO`, []any{schema, table}) + if err != nil { + if isXuguMetadataAccessError(err) { + return s.columnsFromSelect(schema, table, primaryKeys) + } + return nil, err + } + defer rows.Close() + var result []columnInfo + for rows.Next() { + var item columnInfo + var notNull any + var scale *int + if err := rows.Scan( + &item.Name, + &item.DataType, + ¬Null, + &item.ColumnDefault, + &item.Comment, + &scale, + ); err != nil { + return nil, err + } + item.IsNullable = !truthy(notNull) + item.IsPrimaryKey = primaryKeys[strings.ToUpper(item.Name)] + item.NumericPrecision, item.NumericScale, item.CharacterMaximumLength = decodeXuguScale(item.DataType, scale) + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) columnsFromSelect(schema, table string, primaryKeys map[string]bool) ([]columnInfo, error) { + rows, err := s.queryRows( + "SELECT * FROM "+quoteIdentifier(schema)+"."+quoteIdentifier(table)+" WHERE 1 = 0", + nil, + ) + if err != nil { + return nil, err + } + defer rows.Close() + types, err := rows.ColumnTypes() + if err != nil { + return nil, err + } + result := make([]columnInfo, 0, len(types)) + for _, columnType := range types { + item := columnInfo{ + Name: columnType.Name(), + DataType: columnType.DatabaseTypeName(), + IsPrimaryKey: primaryKeys[strings.ToUpper(columnType.Name())], + } + if nullable, ok := columnType.Nullable(); ok { + item.IsNullable = nullable + } else { + item.IsNullable = true + } + if length, ok := columnType.Length(); ok { + value := int(length) + item.CharacterMaximumLength = &value + } + result = append(result, item) + } + return emptyIfNil(result), nil +} + +func (s *server) primaryKeyColumns(schema, table string) (map[string]bool, error) { + rows, err := s.queryRows(xuguPrimaryKeyColumnsSQL, []any{schema, table}) + if err != nil { + if isXuguMetadataAccessError(err) { + return map[string]bool{}, nil + } + return nil, err + } + defer rows.Close() + result := map[string]bool{} + for rows.Next() { + var define string + if err := rows.Scan(&define); err != nil { + return nil, err + } + for _, column := range parseQuotedIdentifiers(define) { + result[strings.ToUpper(column)] = true + } + } + return result, rows.Err() +} + +func (s *server) listIndexes(schema, table string) ([]indexInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT i.INDEX_NAME, i.KEYS, i.IS_UNIQUE, i.IS_PRIMARY, i.INDEX_TYPE, i.FILTER +FROM SYS_INDEXES i +JOIN SYS_TABLES t ON t.DB_ID = i.DB_ID AND t.TABLE_ID = i.TABLE_ID +JOIN SYS_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) + AND UPPER(t.TABLE_NAME) = UPPER(?) +ORDER BY i.INDEX_NAME`, []any{schema, table}) + if err != nil { + if isXuguMetadataAccessError(err) { + return []indexInfo{}, nil + } + return nil, err + } + defer rows.Close() + + var result []indexInfo + for rows.Next() { + var item indexInfo + var keys string + var unique, primary any + var indexType any + if err := rows.Scan(&item.Name, &keys, &unique, &primary, &indexType, &item.Filter); err != nil { + return nil, err + } + item.Columns = parseIndexKeys(keys) + item.IsUnique = truthy(unique) + item.IsPrimary = truthy(primary) + item.IndexType = stringPtr(indexTypeName(indexType)) + item.IncludedColumns = []string{} + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) listForeignKeys(schema, table string) ([]foreignKeyInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT c.CONS_NAME, c.DEFINE, rt.TABLE_NAME +FROM SYS_CONSTRAINTS c +JOIN SYS_TABLES t ON t.DB_ID = c.DB_ID AND t.TABLE_ID = c.TABLE_ID +JOIN SYS_TABLES rt ON rt.DB_ID = c.DB_ID AND rt.TABLE_ID = c.REF_TABLE_ID +JOIN SYS_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) + AND UPPER(t.TABLE_NAME) = UPPER(?) + AND c.CONS_TYPE = 'F' +ORDER BY c.CONS_NAME`, []any{schema, table}) + if err != nil { + if isXuguMetadataAccessError(err) { + return []foreignKeyInfo{}, nil + } + return nil, err + } + defer rows.Close() + var result []foreignKeyInfo + for rows.Next() { + var name, define, refTable string + if err := rows.Scan(&name, &define, &refTable); err != nil { + return nil, err + } + local, ref := parseForeignKeyColumns(define) + for i, column := range local { + item := foreignKeyInfo{Name: name, Column: column, RefTable: refTable} + if i < len(ref) { + item.RefColumn = ref[i] + } + result = append(result, item) + } + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) listTriggers(schema, table string) ([]triggerInfo, error) { + schema, err := s.normalizeSchema(schema) + if err != nil { + return nil, err + } + table = strings.ToUpper(strings.TrimSpace(table)) + rows, err := s.queryRows(` +SELECT tr.TRIG_NAME, tr.TRIG_EVENT, tr.TRIG_TIME +FROM SYS_TRIGGERS tr +JOIN SYS_TABLES t ON t.DB_ID = tr.DB_ID AND t.TABLE_ID = tr.OBJ_ID +JOIN SYS_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) + AND UPPER(t.TABLE_NAME) = UPPER(?) +ORDER BY tr.TRIG_NAME`, []any{schema, table}) + if err != nil { + if isXuguMetadataAccessError(err) { + return []triggerInfo{}, nil + } + return nil, err + } + defer rows.Close() + var result []triggerInfo + for rows.Next() { + var item triggerInfo + var event, timing any + if err := rows.Scan(&item.Name, &event, &timing); err != nil { + return nil, err + } + item.Event = triggerEventName(event) + item.Timing = triggerTimingName(timing) + result = append(result, item) + } + return emptyIfNil(result), rows.Err() +} + +func (s *server) getObjectSource(schema, name, objectType string) (map[string]any, error) { + var err error + schema, err = s.normalizeSchema(schema) + if err != nil { + return nil, err + } + sourceSQL, args, err := objectSourceQuery(schema, name, objectType) + if err != nil { + return nil, err + } + rows, err := s.queryRows(sourceSQL, args) + if err != nil { + return nil, err + } + defer rows.Close() + var builder strings.Builder + for rows.Next() { + var line string + if err := rows.Scan(&line); err != nil { + return nil, err + } + builder.WriteString(line) + } + return map[string]any{"name": name, "object_type": objectType, "schema": schema, "source": builder.String()}, rows.Err() +} + +func (s *server) getTableDDL(schema, table string) (string, error) { + var err error + schema, err = s.normalizeSchema(schema) + if err != nil { + return "", err + } + var ddl string + rows, err := s.queryRows("SELECT TO_CHAR(DBMS_METADATA.GET_DDL('TABLE', ?, ?)) FROM DUAL", []any{strings.ToUpper(table), schema}) + if err == nil { + defer rows.Close() + if rows.Next() { + if scanErr := rows.Scan(&ddl); scanErr == nil && strings.TrimSpace(ddl) != "" { + return ddl, rows.Err() + } + } + } + ddl, err = s.buildTableDDL(schema, table) + if err != nil { + return "", err + } + return ddl, nil +} + +func (s *server) getExplainInfo(sqlText string) (string, error) { + if strings.TrimSpace(sqlText) == "" { + return "", errors.New("sql is required") + } + rows, err := s.queryRows("EXPLAIN "+trimStatementSQL(sqlText), nil) + if err != nil { + return "", err + } + defer rows.Close() + columns, err := rows.Columns() + if err != nil { + return "", err + } + var builder strings.Builder + for rows.Next() { + values, err := scanRow(rows, len(columns)) + if err != nil { + return "", err + } + builder.WriteString(joinValues(values, "\t")) + builder.WriteByte('\n') + } + return strings.TrimSpace(builder.String()), rows.Err() +} + +func (s *server) executeTransaction(params map[string]json.RawMessage) (queryResult, error) { + var payload struct { + Statements []string `json:"statements"` + Schema string `json:"schema"` + } + if err := decodeParams(params, &payload); err != nil { + return queryResult{}, err + } + db, err := s.requireDB() + if err != nil { + return queryResult{}, err + } + tx, err := db.Begin() + if err != nil { + return queryResult{}, err + } + if strings.TrimSpace(payload.Schema) != "" { + if _, err := tx.Exec("SET SCHEMA " + quoteIdentifier(payload.Schema)); err != nil { + tx.Rollback() + return queryResult{}, err + } + } + var affected int64 + start := time.Now() + for _, statement := range payload.Statements { + statement = trimStatementSQL(statement) + if statement == "" { + continue + } + result, err := tx.Exec(statement) + if err != nil { + tx.Rollback() + return queryResult{}, err + } + count, _ := result.RowsAffected() + affected += count + } + if err := tx.Commit(); err != nil { + return queryResult{}, err + } + return queryResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + AffectedRows: affected, + ExecutionTimeMS: time.Since(start).Milliseconds(), + }, nil +} + +func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageResult, error) { + start := time.Now() + if err := s.useDatabase(opts.Database); err != nil { + return queryPageResult{}, err + } + if strings.TrimSpace(opts.Schema) != "" { + if err := s.setSchema(opts.Schema); err != nil { + return queryPageResult{}, err + } + } + sqlText := trimStatementSQL(opts.SQL) + if !isQuerySQL(sqlText) { + result, err := s.executeQuery(opts) + return queryPageResult{ + Columns: result.Columns, + ColumnTypes: result.ColumnTypes, + Rows: result.Rows, + AffectedRows: result.AffectedRows, + ExecutionTimeMS: result.ExecutionTimeMS, + Truncated: result.Truncated, + SessionID: nil, + HasMore: false, + }, err + } + rows, err := s.queryRows(sqlText, nil) + if err != nil { + return queryPageResult{}, err + } + columns, err := rows.Columns() + if err != nil { + rows.Close() + return queryPageResult{}, err + } + columnTypes := columnTypeNames(rows) + maxRows := opts.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + session := &querySession{rows: rows, columns: columns, columnTypes: columnTypes, remaining: maxRows} + result, err := readQuerySessionPage(session, pageSize) + result.ExecutionTimeMS = time.Since(start).Milliseconds() + if err != nil { + rows.Close() + return queryPageResult{}, err + } + if result.HasMore { + sessionID := s.storeQuerySession(session) + result.SessionID = &sessionID + } else { + rows.Close() + } + return result, nil +} + +func (s *server) fetchQueryPage(sessionID string, pageSize int) (queryPageResult, error) { + session := s.sessions[sessionID] + if session == nil { + return queryPageResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, SessionID: nil, HasMore: false}, nil + } + result, err := readQuerySessionPage(session, pageSize) + if err != nil { + s.closeQuerySession(sessionID) + return queryPageResult{}, err + } + if result.HasMore { + result.SessionID = &sessionID + } else { + s.closeQuerySession(sessionID) + } + return result, nil +} + +func (s *server) storeQuerySession(session *querySession) string { + s.nextSessionID++ + sessionID := fmt.Sprintf("xugu-%d", s.nextSessionID) + s.sessions[sessionID] = session + return sessionID +} + +func (s *server) closeQuerySession(sessionID string) bool { + session := s.sessions[sessionID] + if session == nil { + return false + } + session.rows.Close() + delete(s.sessions, sessionID) + return true +} + +func (s *server) closeAllQuerySessions() { + for sessionID := range s.sessions { + s.closeQuerySession(sessionID) + } +} + +func readQuerySessionPage(session *querySession, pageSize int) (queryPageResult, error) { + if pageSize <= 0 { + pageSize = defaultMaxRows + } + result := queryPageResult{Columns: session.columns, ColumnTypes: session.columnTypes, Rows: [][]any{}, SessionID: nil, HasMore: false} + for len(result.Rows) < pageSize && session.remaining > 0 { + if session.pending != nil { + result.Rows = append(result.Rows, session.pending) + session.pending = nil + session.remaining-- + continue + } + if !session.rows.Next() { + return result, session.rows.Err() + } + row, err := scanRow(session.rows, len(session.columns)) + if err != nil { + return queryPageResult{}, err + } + result.Rows = append(result.Rows, row) + session.remaining-- + } + if session.remaining <= 0 { + result.Truncated = true + return result, nil + } + if session.rows.Next() { + row, err := scanRow(session.rows, len(session.columns)) + if err != nil { + return queryPageResult{}, err + } + session.pending = row + result.HasMore = true + return result, nil + } + return result, session.rows.Err() +} + +func (s *server) executeQuery(opts queryOptions) (queryResult, error) { + start := time.Now() + if err := s.useDatabase(opts.Database); err != nil { + return queryResult{}, err + } + if strings.TrimSpace(opts.Schema) != "" { + if err := s.setSchema(opts.Schema); err != nil { + return queryResult{}, err + } + } + sqlText := trimStatementSQL(opts.SQL) + maxRows := opts.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + if isQuerySQL(sqlText) { + result, err := s.executeSelect(sqlText, maxRows) + result.ExecutionTimeMS = time.Since(start).Milliseconds() + return result, err + } + db, err := s.requireDB() + if err != nil { + return queryResult{}, err + } + execResult, err := db.Exec(sqlText) + if err != nil { + return queryResult{}, err + } + affected, _ := execResult.RowsAffected() + return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil +} + +func (s *server) executeSelect(sqlText string, maxRows int) (queryResult, error) { + rows, err := s.queryRows(sqlText, nil) + if err != nil { + return queryResult{}, err + } + defer rows.Close() + columns, err := rows.Columns() + if err != nil { + return queryResult{}, err + } + result := queryResult{Columns: columns, ColumnTypes: columnTypeNames(rows), Rows: [][]any{}} + for rows.Next() { + if len(result.Rows) >= maxRows { + result.Truncated = true + break + } + values, err := scanRow(rows, len(columns)) + if err != nil { + return queryResult{}, err + } + result.Rows = append(result.Rows, values) + } + return result, rows.Err() +} + +func scanRow(rows *sql.Rows, columnCount int) ([]any, error) { + values := make([]any, columnCount) + scanTargets := make([]any, columnCount) + for i := range values { + scanTargets[i] = &values[i] + } + if err := rows.Scan(scanTargets...); err != nil { + return nil, err + } + for i, value := range values { + values[i] = normalizeValue(value) + } + return values, nil +} + +func columnTypeNames(rows *sql.Rows) []string { + types, err := rows.ColumnTypes() + if err != nil { + return []string{} + } + result := make([]string, 0, len(types)) + for _, columnType := range types { + result = append(result, columnType.DatabaseTypeName()) + } + return result +} + +func (s *server) setSchema(schema string) error { + db, err := s.requireDB() + if err != nil { + return err + } + _, err = db.Exec("SET SCHEMA " + quoteIdentifier(schema)) + return err +} + +func (s *server) queryRows(sqlText string, args []any) (*sql.Rows, error) { + db, err := s.requireDB() + if err != nil { + return nil, err + } + if len(args) == 0 { + return db.Query(sqlText) + } + return db.Query(sqlText, args...) +} + +func objectSourceQuery(schema, name, objectType string) (string, []any, error) { + objectType = strings.ToUpper(strings.TrimSpace(objectType)) + name = strings.ToUpper(strings.TrimSpace(name)) + switch objectType { + case "VIEW": + return ` +SELECT TO_CHAR(v.DEFINE) +FROM SYS_VIEWS v +JOIN SYS_SCHEMAS s ON s.DB_ID = v.DB_ID AND s.SCHEMA_ID = v.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) AND UPPER(v.VIEW_NAME) = UPPER(?)`, []any{schema, name}, nil + case "TRIGGER": + return ` +SELECT TO_CHAR(t.DEFINE) +FROM SYS_TRIGGERS t +JOIN SYS_SCHEMAS s ON s.DB_ID = t.DB_ID AND s.SCHEMA_ID = t.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) AND UPPER(t.TRIG_NAME) = UPPER(?)`, []any{schema, name}, nil + case "PROCEDURE", "FUNCTION": + return ` +SELECT TO_CHAR(p.DEFINE) +FROM SYS_PROCEDURES p +JOIN SYS_SCHEMAS s ON s.DB_ID = p.DB_ID AND s.SCHEMA_ID = p.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) AND UPPER(p.PROC_NAME) = UPPER(?)`, []any{schema, name}, nil + case "PACKAGE", "PACKAGE BODY": + return ` +SELECT COALESCE(TO_CHAR(k.SPEC), '') || COALESCE(TO_CHAR(k.BODY), '') +FROM SYS_PACKAGES k +JOIN SYS_SCHEMAS s ON s.DB_ID = k.DB_ID AND s.SCHEMA_ID = k.SCHEMA_ID +WHERE UPPER(s.SCHEMA_NAME) = UPPER(?) AND UPPER(k.PACK_NAME) = UPPER(?)`, []any{schema, name}, nil + default: + return "", nil, fmt.Errorf("object source is not supported for %s", objectType) + } +} + +func (s *server) buildTableDDL(schema, table string) (string, error) { + columns, err := s.getColumns(schema, table) + if err != nil { + return "", err + } + if len(columns) == 0 { + return "", fmt.Errorf("table not found: %s.%s", schema, table) + } + var builder strings.Builder + builder.WriteString("CREATE TABLE ") + builder.WriteString(quoteIdentifier(schema)) + builder.WriteByte('.') + builder.WriteString(quoteIdentifier(table)) + builder.WriteString(" (\n") + for i, column := range columns { + if i > 0 { + builder.WriteString(",\n") + } + builder.WriteString(" ") + builder.WriteString(quoteIdentifier(column.Name)) + builder.WriteByte(' ') + builder.WriteString(columnTypeDDL(column)) + if column.ColumnDefault != nil && strings.TrimSpace(*column.ColumnDefault) != "" { + builder.WriteString(" DEFAULT ") + builder.WriteString(strings.TrimSpace(*column.ColumnDefault)) + } + if !column.IsNullable { + builder.WriteString(" NOT NULL") + } + } + primary := make([]string, 0) + for _, column := range columns { + if column.IsPrimaryKey { + primary = append(primary, quoteIdentifier(column.Name)) + } + } + if len(primary) > 0 { + builder.WriteString(",\n PRIMARY KEY (") + builder.WriteString(strings.Join(primary, ", ")) + builder.WriteByte(')') + } + builder.WriteString("\n)") + return builder.String(), nil +} + +func columnTypeDDL(column columnInfo) string { + dataType := strings.ToUpper(strings.TrimSpace(column.DataType)) + if column.CharacterMaximumLength != nil { + return fmt.Sprintf("%s(%d)", dataType, *column.CharacterMaximumLength) + } + if column.NumericPrecision != nil && column.NumericScale != nil { + return fmt.Sprintf("%s(%d,%d)", dataType, *column.NumericPrecision, *column.NumericScale) + } + return dataType +} + +func decodeXuguScale(dataType string, scale *int) (*int, *int, *int) { + if scale == nil || *scale < 0 { + return nil, nil, nil + } + upper := strings.ToUpper(dataType) + if strings.Contains(upper, "CHAR") || strings.Contains(upper, "BINARY") { + length := *scale + return nil, nil, &length + } + if strings.Contains(upper, "NUM") || strings.Contains(upper, "DECIMAL") { + precision := *scale / 65536 + numericScale := *scale % 65536 + return &precision, &numericScale, nil + } + return nil, nil, nil +} + +var quotedIdentifierRegexp = regexp.MustCompile(`"([^"]+)"`) + +func parseQuotedIdentifiers(value string) []string { + matches := quotedIdentifierRegexp.FindAllStringSubmatch(value, -1) + result := make([]string, 0, len(matches)) + for _, match := range matches { + if len(match) == 2 { + result = append(result, match[1]) + } + } + return result +} + +func parseIndexKeys(value string) []string { + quoted := parseQuotedIdentifiers(value) + if len(quoted) > 0 { + return quoted + } + parts := strings.Split(value, ",") + result := make([]string, 0, len(parts)) + for _, part := range parts { + part = strings.Trim(strings.TrimSpace(part), `"`) + if part != "" { + result = append(result, part) + } + } + return result +} + +func parseForeignKeyColumns(define string) ([]string, []string) { + groups := regexp.MustCompile(`\(([^()]*)\)`).FindAllStringSubmatch(define, -1) + if len(groups) < 2 { + columns := parseQuotedIdentifiers(define) + if len(columns)%2 == 0 { + mid := len(columns) / 2 + return columns[:mid], columns[mid:] + } + return columns, nil + } + return parseQuotedIdentifiers(groups[0][1]), parseQuotedIdentifiers(groups[1][1]) +} + +func truthy(value any) bool { + switch v := normalizeValue(value).(type) { + case bool: + return v + case int64: + return v != 0 + case float64: + return v != 0 + case string: + upper := strings.ToUpper(strings.TrimSpace(v)) + return upper == "T" || upper == "TRUE" || upper == "1" || upper == "Y" || upper == "YES" + default: + return false + } +} + +func stringPtr(value string) *string { + if strings.TrimSpace(value) == "" { + return nil + } + return &value +} + +func indexTypeName(value any) string { + switch fmt.Sprint(normalizeValue(value)) { + case "0": + return "BTREE" + case "1": + return "RTREE" + case "2": + return "FULLTEXT" + case "3": + return "BITMAP" + default: + return fmt.Sprint(normalizeValue(value)) + } +} + +func triggerEventName(value any) string { + switch fmt.Sprint(normalizeValue(value)) { + case "1": + return "INSERT" + case "2": + return "UPDATE" + case "3": + return "INSERT OR UPDATE" + case "4": + return "DELETE" + case "5": + return "INSERT OR DELETE" + case "6": + return "UPDATE OR DELETE" + case "7": + return "INSERT OR UPDATE OR DELETE" + case "8": + return "LOGON" + default: + return fmt.Sprint(normalizeValue(value)) + } +} + +func triggerTimingName(value any) string { + switch fmt.Sprint(normalizeValue(value)) { + case "1": + return "BEFORE" + case "2": + return "INSTEAD" + case "4": + return "AFTER" + default: + return fmt.Sprint(normalizeValue(value)) + } +} + +func joinValues(values []any, sep string) string { + parts := make([]string, len(values)) + for i, value := range values { + if value == nil { + parts[i] = "" + } else { + parts[i] = fmt.Sprint(value) + } + } + return strings.Join(parts, sep) +} + +func decodeParams(params map[string]json.RawMessage, target any) error { + if params == nil { + params = map[string]json.RawMessage{} + } + 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 params == nil || len(params[key]) == 0 { + return "" + } + var value string + _ = json.Unmarshal(params[key], &value) + return value +} + +func intParam(params map[string]json.RawMessage, key string) int { + if params == nil || len(params[key]) == 0 { + return 0 + } + var value int + _ = json.Unmarshal(params[key], &value) + return value +} + +func errorResponse(id json.RawMessage, err error) response { + return response{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: -1, Message: err.Error()}} +} + +func trimStatementSQL(sqlText string) string { + return strings.TrimRight(strings.TrimSpace(sqlText), "; \t\r\n") +} + +func isQuerySQL(sqlText string) bool { + lower := strings.ToLower(strings.TrimSpace(sqlText)) + return strings.HasPrefix(lower, "select") || strings.HasPrefix(lower, "with") +} + +func quoteIdentifier(value string) string { + return `"` + strings.ReplaceAll(value, `"`, `""`) + `"` +} + +func normalizeValue(value any) any { + switch v := value.(type) { + case nil: + return nil + case []byte: + return string(v) + case time.Time: + return v.Format(time.RFC3339Nano) + case int: + return int64(v) + case int8: + return int64(v) + case int16: + return int64(v) + case int32: + return int64(v) + case int64: + return v + case uint: + return uint64(v) + case uint8: + return uint64(v) + case uint16: + return uint64(v) + case uint32: + return uint64(v) + case uint64: + return v + case float32: + return float64(v) + case float64, bool, string: + return v + case fmt.Stringer: + return v.String() + default: + return fmt.Sprint(v) + } +} + +func emptyIfNil[T any](values []T) []T { + if values == nil { + return []T{} + } + return values +} diff --git a/agents/drivers/xugu/main_test.go b/agents/drivers/xugu/main_test.go new file mode 100644 index 0000000000..8b4346742c --- /dev/null +++ b/agents/drivers/xugu/main_test.go @@ -0,0 +1,370 @@ +package main + +import ( + "encoding/json" + "errors" + "os" + "strings" + "testing" +) + +func TestHandshakeResponse(t *testing.T) { + s := newServer() + resp, shutdown := s.handleLine(`{"jsonrpc":"2.0","id":7,"method":"handshake","params":{"appVersion":"dev"}}`) + if shutdown { + t.Fatal("handshake should not shut down the server") + } + if resp.Error != nil { + t.Fatalf("unexpected error: %v", resp.Error) + } + data, err := json.Marshal(resp.Result) + if err != nil { + t.Fatal(err) + } + var result struct { + ProtocolVersion int `json:"protocolVersion"` + AgentProtocolVersion int `json:"agentProtocolVersion"` + Capabilities []string `json:"capabilities"` + } + if err := json.Unmarshal(data, &result); err != nil { + t.Fatal(err) + } + if result.ProtocolVersion != 1 || result.AgentProtocolVersion != 1 { + t.Fatalf("unexpected protocol versions: %+v", result) + } + contract := protocolContract(t) + if result.ProtocolVersion != contract.ProtocolVersion || result.AgentProtocolVersion != contract.ProtocolVersion { + t.Fatalf("handshake protocol versions do not match contract: result=%+v contract=%+v", result, contract) + } + for _, capability := range result.Capabilities { + if !contains(contract.AllCapabilities, capability) { + t.Fatalf("handshake returned capability %q outside protocol contract %v", capability, contract.AllCapabilities) + } + } + if !contains(result.Capabilities, "query") || !contains(result.Capabilities, "metadata") { + t.Fatalf("expected query and metadata capabilities, got %v", result.Capabilities) + } +} + +func TestCloseMissingQuerySessionReturnsFalse(t *testing.T) { + s := newServer() + resp, shutdown := s.handleLine(`{"jsonrpc":"2.0","id":8,"method":"close_query_session","params":{"sessionId":"missing"}}`) + if shutdown { + t.Fatal("close_query_session should not shut down the server") + } + if resp.Error != nil { + t.Fatalf("unexpected error: %v", resp.Error) + } + if resp.Result != false { + t.Fatalf("expected false result, got %#v", resp.Result) + } +} + +func TestEmptyResultSlicesMarshalAsArrays(t *testing.T) { + data, err := json.Marshal(queryResult{}) + if err != nil { + t.Fatal(err) + } + text := string(data) + if strings.Contains(text, `"columns":null`) || strings.Contains(text, `"column_types":null`) || strings.Contains(text, `"rows":null`) { + t.Fatalf("query result should marshal nil slices as arrays: %s", text) + } + + data, err = json.Marshal(indexInfo{}) + if err != nil { + t.Fatal(err) + } + text = string(data) + if strings.Contains(text, `"columns":null`) || strings.Contains(text, `"included_columns":null`) { + t.Fatalf("index info should marshal nil slices as arrays: %s", text) + } +} + +func TestGetTableDDLResultMarshalsAsString(t *testing.T) { + data, err := json.Marshal("CREATE TABLE SYSDBA.ORDERS (ID INT)") + if err != nil { + t.Fatal(err) + } + var ddl string + if err := json.Unmarshal(data, &ddl); err != nil { + t.Fatalf("get_table_ddl result must deserialize as a string: %v", err) + } +} + +func TestBuildDSNUsesConnectionStringWhenProvided(t *testing.T) { + dsn := buildDSN(connectParams{ConnectionString: "IP=db.example.com;DB=SYSTEM;User=SYSDBA;PWD=secret;Port=5138"}) + + if dsn != "IP=db.example.com;DB=SYSTEM;User=SYSDBA;PWD=secret;Port=5138" { + t.Fatalf("unexpected dsn: %s", dsn) + } +} + +func protocolContract(t *testing.T) struct { + ProtocolVersion int `json:"protocolVersion"` + AllCapabilities []string `json:"allCapabilities"` +} { + t.Helper() + data, err := os.ReadFile("../../common/src/main/resources/agent-protocol-v1.json") + if err != nil { + t.Fatal(err) + } + var contract struct { + ProtocolVersion int `json:"protocolVersion"` + AllCapabilities []string `json:"allCapabilities"` + } + if err := json.Unmarshal(data, &contract); err != nil { + t.Fatal(err) + } + return contract +} + +func TestBuildDSNUsesConnectionFields(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "db.example.com", + Port: 15138, + Database: "demo", + Username: "sysdba", + Password: "secret", + }) + + for _, part := range []string{"IP=db.example.com", "DB=demo", "User=sysdba", "PWD=secret", "Port=15138"} { + if !strings.Contains(dsn, part) { + t.Fatalf("dsn should contain %s, got: %s", part, dsn) + } + } +} + +func TestBuildDSNUsesDefaultPort(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "db.example.com", + Database: "demo", + Username: "sysdba", + Password: "secret", + }) + + if !strings.Contains(dsn, "Port=5138") { + t.Fatalf("dsn should default to Xugu port, got: %s", dsn) + } +} + +func TestBuildDSNParsesJdbcURL(t *testing.T) { + dsn := buildDSN(connectParams{ + Username: "sysdba", + Password: "secret", + ConnectionString: "jdbc:xugu://db.example.com:15138/demo", + }) + + for _, part := range []string{"IP=db.example.com", "DB=demo", "User=sysdba", "PWD=secret", "Port=15138"} { + if !strings.Contains(dsn, part) { + t.Fatalf("dsn should contain %s, got: %s", part, dsn) + } + } +} + +func TestBuildDSNParsesDBXURL(t *testing.T) { + dsn := buildDSN(connectParams{ + ConnectionString: "xugu://sysdba:secret@db.example.com:15138/demo", + }) + + for _, part := range []string{"IP=db.example.com", "DB=demo", "User=sysdba", "PWD=secret", "Port=15138"} { + if !strings.Contains(dsn, part) { + t.Fatalf("dsn should contain %s, got: %s", part, dsn) + } + } +} + +func TestBuildDSNAppendsURLParams(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "db.example.com", + Database: "demo", + Username: "sysdba", + Password: "secret", + URLParams: "AUTO_COMMIT=on;CHAR_SET=UTF8", + }) + + for _, part := range []string{"AUTO_COMMIT=on", "CHAR_SET=UTF8"} { + if !strings.Contains(dsn, part) { + t.Fatalf("dsn should contain %s, got: %s", part, dsn) + } + } +} + +func TestBuildDSNDefaultsToUTF8(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "db.example.com", + Database: "demo", + Username: "sysdba", + Password: "secret", + }) + + if !strings.Contains(dsn, "CHAR_SET=UTF8") { + t.Fatalf("dsn should default to UTF8, got: %s", dsn) + } +} + +func TestBuildDSNRespectsExplicitCharset(t *testing.T) { + dsn := buildDSN(connectParams{ + Host: "db.example.com", + Database: "demo", + Username: "sysdba", + Password: "secret", + URLParams: "CHAR_SET=GBK", + }) + + if strings.Contains(dsn, "CHAR_SET=UTF8") || !strings.Contains(dsn, "CHAR_SET=GBK") { + t.Fatalf("dsn should respect explicit charset, got: %s", dsn) + } +} + +func TestListDatabasesSQLUsesXuguDictionary(t *testing.T) { + sqlText := strings.ToUpper(xuguListDatabasesSQL) + + if !strings.Contains(sqlText, "ALL_DATABASES") || strings.Contains(sqlText, "SYS_DATABASES") { + t.Fatalf("database listing should query low-privilege ALL_DATABASES, got: %s", xuguListDatabasesSQL) + } +} + +func TestFallbackDatabasesFromParams(t *testing.T) { + cases := []struct { + name string + params connectParams + want string + }{ + { + name: "database field", + params: connectParams{ + Database: "LOWPRIV", + }, + want: "LOWPRIV", + }, + { + name: "dbx url", + params: connectParams{ + ConnectionString: "xugu://user:secret@db.example.com:5138/demo", + }, + want: "demo", + }, + { + name: "jdbc url", + params: connectParams{ + ConnectionString: "jdbc:xugu://db.example.com:5138/reporting", + }, + want: "reporting", + }, + { + name: "native dsn", + params: connectParams{ + ConnectionString: "IP=db.example.com;DB=SYSTEM;User=SYSDBA;PWD=secret;Port=5138", + }, + want: "SYSTEM", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := fallbackDatabasesFromParams(tc.params) + if len(got) != 1 || got[0].Name != tc.want { + t.Fatalf("unexpected fallback databases: got=%v want=%s", got, tc.want) + } + }) + } +} + +func TestUseDatabaseSkipsConfiguredDatabase(t *testing.T) { + s := newServer() + s.params = connectParams{Database: "SYSTEM"} + + if err := s.useDatabase("system"); err != nil { + t.Fatalf("expected configured database USE to be skipped, got: %v", err) + } +} + +func TestConfiguredDatabaseName(t *testing.T) { + cases := []struct { + params connectParams + want string + }{ + {params: connectParams{Database: "SYSTEM"}, want: "SYSTEM"}, + {params: connectParams{ConnectionString: "xugu://user:secret@db.example.com:5138/demo"}, want: "demo"}, + {params: connectParams{ConnectionString: "jdbc:xugu://db.example.com:5138/reporting"}, want: "reporting"}, + {params: connectParams{ConnectionString: "IP=db.example.com;DB=SYSTEM;User=SYSDBA;PWD=secret"}, want: "SYSTEM"}, + } + + for _, tc := range cases { + if got := configuredDatabaseName(tc.params); got != tc.want { + t.Fatalf("configuredDatabaseName(%+v) = %q, want %q", tc.params, got, tc.want) + } + } +} + +func TestSchemaListingSQLUsesLowPrivilegeDictionary(t *testing.T) { + sqlText := strings.ToUpper(xuguListSchemasSQL) + + if !strings.Contains(sqlText, "ALL_SCHEMAS") || strings.Contains(sqlText, "SYS_SCHEMAS") { + t.Fatalf("schema listing should query low-privilege ALL_SCHEMAS, got: %s", xuguListSchemasSQL) + } +} + +func TestPrimaryKeySQLUsesLowPrivilegeDictionary(t *testing.T) { + sqlText := strings.ToUpper(xuguPrimaryKeyColumnsSQL) + + for _, want := range []string{"ALL_CONSTRAINTS", "ALL_TABLES", "ALL_SCHEMAS"} { + if !strings.Contains(sqlText, want) { + t.Fatalf("primary key listing should query %s, got: %s", want, xuguPrimaryKeyColumnsSQL) + } + } + for _, forbidden := range []string{"SYS_CONSTRAINTS", "SYS_TABLES", "SYS_SCHEMAS"} { + if strings.Contains(sqlText, forbidden) { + t.Fatalf("primary key listing should not query %s, got: %s", forbidden, xuguPrimaryKeyColumnsSQL) + } + } +} + +func TestXuguMetadataAccessErrorDetection(t *testing.T) { + if !isXuguMetadataAccessError(errors.New("[E18012] 权限不够")) { + t.Fatal("expected E18012 permission error to be treated as metadata access error") + } + if isXuguMetadataAccessError(errors.New("network timeout")) { + t.Fatal("network errors should not trigger database-list fallback") + } +} + +func TestParseForeignKeyColumns(t *testing.T) { + local, ref := parseForeignKeyColumns(`("C1","C2")("ID1","ID2")`) + + if strings.Join(local, ",") != "C1,C2" || strings.Join(ref, ",") != "ID1,ID2" { + t.Fatalf("unexpected foreign key columns: local=%v ref=%v", local, ref) + } +} + +func TestDecodeXuguScale(t *testing.T) { + numericScale := 32*65536 + 6 + precision, scale, length := decodeXuguScale("NUMERIC", &numericScale) + if precision == nil || *precision != 32 || scale == nil || *scale != 6 || length != nil { + t.Fatalf("unexpected numeric scale decode: precision=%v scale=%v length=%v", precision, scale, length) + } + + charScale := 128 + precision, scale, length = decodeXuguScale("VARCHAR", &charScale) + if precision != nil || scale != nil || length == nil || *length != 128 { + t.Fatalf("unexpected char scale decode: precision=%v scale=%v length=%v", precision, scale, length) + } +} + +func TestNormalizeValuePreservesDriverNumericTypes(t *testing.T) { + if value := normalizeValue(int32(7)); value != int64(7) { + t.Fatalf("expected int32 to normalize to int64, got %#v", value) + } + if value := normalizeValue(float32(1.25)); value != float64(float32(1.25)) { + t.Fatalf("expected float32 to normalize to float64, got %#v", value) + } +} + +func contains(values []string, target string) bool { + for _, value := range values { + if value == target { + return true + } + } + return false +} diff --git a/agents/drivers/yashandb/build.gradle b/agents/drivers/yashandb/build.gradle new file mode 100644 index 0000000000..5a8041288e --- /dev/null +++ b/agents/drivers/yashandb/build.gradle @@ -0,0 +1,9 @@ +dependencies { + implementation 'com.yashandb:yashandb-jdbc:1.9.27' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': '崖山 YashanDB', 'Main-Class': 'com.dbx.agent.yashandb.YashandbAgent') + } +} diff --git a/agents/drivers/yashandb/src/main/java/com/dbx/agent/yashandb/YashandbAgent.java b/agents/drivers/yashandb/src/main/java/com/dbx/agent/yashandb/YashandbAgent.java new file mode 100644 index 0000000000..cf11cc080e --- /dev/null +++ b/agents/drivers/yashandb/src/main/java/com/dbx/agent/yashandb/YashandbAgent.java @@ -0,0 +1,22 @@ +package com.dbx.agent.yashandb; + +import com.dbx.agent.ConfiguredJdbcAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.JsonRpcServer; + +public final class YashandbAgent extends ConfiguredJdbcAgent { + public static final JdbcAgentProfile YASHANDB_PROFILE = new JdbcAgentProfile( + "com.yashandb.jdbc.Driver", + "jdbc:yasdb://{host}:{port}/{database}", + 1688, + true + ); + + public YashandbAgent() { + super(YASHANDB_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new YashandbAgent()).run(); + } +} diff --git a/agents/drivers/yashandb/src/test/java/com/dbx/agent/yashandb/YashandbAgentTest.java b/agents/drivers/yashandb/src/test/java/com/dbx/agent/yashandb/YashandbAgentTest.java new file mode 100644 index 0000000000..0ae5c64091 --- /dev/null +++ b/agents/drivers/yashandb/src/test/java/com/dbx/agent/yashandb/YashandbAgentTest.java @@ -0,0 +1,30 @@ +package com.dbx.agent.yashandb; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.JdbcAgentProfile; +import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; +import java.util.Collections; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class YashandbAgentTest extends JdbcFakeExecutionBehaviorTest { + @Override + protected DatabaseAgent createAgent() { + return new YashandbAgent(); + } + + @Override + protected String resultSetSql() { + return "SELECT 1 FROM DUAL"; + } + + @Test + void declaresYashandbJdbcProfile() { + JdbcAgentProfile profile = YashandbAgent.YASHANDB_PROFILE; + + Assertions.assertEquals("com.yashandb.jdbc.Driver", profile.getDriverClass()); + Assertions.assertEquals("jdbc:yasdb://{host}:{port}/{database}", profile.getUrlTemplate()); + Assertions.assertTrue(profile.getSkipExecutionContext()); + Assertions.assertEquals(Collections.emptySet(), profile.getExcludedSchemas()); + } +} diff --git a/agents/drivers/zookeeper/build.gradle b/agents/drivers/zookeeper/build.gradle new file mode 100644 index 0000000000..f679c11a84 --- /dev/null +++ b/agents/drivers/zookeeper/build.gradle @@ -0,0 +1,13 @@ +dependencies { + implementation 'com.google.code.gson:gson:2.12.1' + implementation 'org.apache.curator:curator-framework:5.9.0' + runtimeOnly 'org.slf4j:slf4j-simple:1.7.36' + testImplementation 'org.apache.curator:curator-test:5.9.0' +} + +tasks.named('shadowJar') { + mergeServiceFiles() + manifest { + attributes('Agent-Label': 'Apache ZooKeeper', 'Main-Class': 'com.dbx.agent.zookeeper.ZooKeeperAgent') + } +} diff --git a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java new file mode 100644 index 0000000000..45e7534a14 --- /dev/null +++ b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java @@ -0,0 +1,549 @@ +package com.dbx.agent.zookeeper; + +import com.dbx.agent.AgentProtocol; +import com.google.gson.Gson; +import com.google.gson.GsonBuilder; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CharsetDecoder; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Base64; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import org.apache.curator.framework.CuratorFramework; +import org.apache.curator.framework.CuratorFrameworkFactory; +import org.apache.curator.retry.ExponentialBackoffRetry; +import org.apache.zookeeper.KeeperException; +import org.apache.zookeeper.data.Stat; + +public final class ZooKeeperAgent { + private static final Gson GSON = new GsonBuilder().serializeNulls().create(); + private static final int DEFAULT_LIMIT = 100; + private static final int DEFAULT_SESSION_TIMEOUT_MS = 30000; + private static final int DEFAULT_CONNECTION_TIMEOUT_MS = 15000; + private static final int DEFAULT_BASE_SLEEP_TIME_MS = 1000; + private static final int DEFAULT_MAX_RETRIES = 3; + private static final List CAPABILITIES = Collections.unmodifiableList(Arrays.asList( + AgentProtocol.CAPABILITY_CONNECT, + AgentProtocol.CAPABILITY_TEST_CONNECTION, + AgentProtocol.CAPABILITY_KV + )); + private static CuratorFramework client; + + private ZooKeeperAgent() { + } + + private static Object handshakeResult() { + return new HandshakeResult(AgentProtocol.PROTOCOL_VERSION, AgentProtocol.PROTOCOL_VERSION, CAPABILITIES); + } + + private static Object connect(JsonObject params) throws Exception { + JsonObject connection = connectionObject(params); + CuratorFramework nextClient = buildClient(connection); + try { + startAndVerify(nextClient, connection); + closeClient(); + client = nextClient; + return Collections.singletonMap("ok", true); + } catch (Exception e) { + nextClient.close(); + throw e; + } + } + + private static Object testConnection(JsonObject params) throws Exception { + JsonObject connection = connectionObject(params); + CuratorFramework probe = buildClient(connection); + try { + startAndVerify(probe, connection); + return Collections.singletonMap("ok", true); + } finally { + probe.close(); + } + } + + static CuratorFramework buildClient(JsonObject connection) { + if (hasTlsOptions(connection)) { + throw new IllegalArgumentException("ZooKeeper TLS is not supported"); + } + + int baseSleepTimeMs = intOrDefault(connection, "base_sleep_time_ms", DEFAULT_BASE_SLEEP_TIME_MS); + int maxRetries = intOrDefault(connection, "max_retries", DEFAULT_MAX_RETRIES); + CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder() + .connectString(connectString(connection)) + .sessionTimeoutMs(intOrDefault(connection, "session_timeout_ms", DEFAULT_SESSION_TIMEOUT_MS)) + .connectionTimeoutMs(intOrDefault(connection, "connection_timeout_ms", DEFAULT_CONNECTION_TIMEOUT_MS)) + .retryPolicy(new ExponentialBackoffRetry(baseSleepTimeMs, maxRetries)); + + String namespace = trimSlashes(stringOrEmpty(connection, "namespace")); + if (!namespace.isBlank()) { + builder.namespace(namespace); + } + + String username = stringOrEmpty(connection, "username"); + String password = stringOrEmpty(connection, "password"); + if (!username.isBlank()) { + builder.authorization("digest", (username + ":" + password).getBytes(StandardCharsets.UTF_8)); + } + + return builder.build(); + } + + static String connectString(JsonObject connection) { + String configured = firstNonBlank( + stringOrNull(connection, "zookeeper_connect_string"), + stringOrNull(connection, "connect_string"), + stringOrNull(connection, "connection_string") + ); + if (configured != null) { + return configured; + } + return stringOrDefault(connection, "host", "127.0.0.1") + ":" + + intOrDefault(connection, "port", 2181); + } + + private static void startAndVerify(CuratorFramework active, JsonObject connection) throws Exception { + active.start(); + int timeoutMs = intOrDefault(connection, "connection_timeout_ms", DEFAULT_CONNECTION_TIMEOUT_MS); + if (!active.blockUntilConnected(timeoutMs, TimeUnit.MILLISECONDS)) { + throw new IllegalStateException("Connection timed out"); + } + if (active.checkExists().forPath("/") == null) { + throw new IllegalStateException("Root znode is not readable"); + } + } + + private static Object listPrefix(JsonObject params) throws Exception { + CuratorFramework active = requireClient(); + String root = normalizePath(stringOrDefault(params, "prefix", "/")); + boolean recursive = boolOrDefault(params, "recursive", true); + int limit = Math.max(1, intOrDefault(params, "limit", DEFAULT_LIMIT)); + Cursor cursor = null; + String continuation = stringOrNull(params, "continuation"); + if (continuation != null && !continuation.isBlank()) { + cursor = decodeCursor(continuation); + if (!root.equals(cursor.root) || recursive != cursor.recursive) { + throw new IllegalArgumentException("Continuation does not match request"); + } + } + + Map result = new LinkedHashMap<>(); + List> keys = new ArrayList<>(); + Stat rootStat = active.checkExists().forPath(root); + if (rootStat == null) { + result.put("keys", keys); + result.put("continuation", null); + return result; + } + + List paths = recursive ? listRecursive(active, root) : listDirectChildren(active, root); + paths.removeIf(path -> shouldHideFromRootListing(root, path)); + Collections.sort(paths); + int offset = cursor == null ? 0 : Math.max(0, cursor.offset); + int end = Math.min(paths.size(), offset + limit); + for (int index = offset; index < end; index++) { + Stat stat = active.checkExists().forPath(paths.get(index)); + if (stat == null) { + continue; + } + Map row = new LinkedHashMap<>(); + row.put("key", paths.get(index)); + row.putAll(statMetadata(stat)); + keys.add(row); + } + + result.put("keys", keys); + result.put("continuation", end < paths.size() ? encodeCursor(new Cursor(root, recursive, end)) : null); + return result; + } + + private static Object get(JsonObject params) throws Exception { + CuratorFramework active = requireClient(); + String path = normalizePath(stringOrDefault(params, "key", "")); + Map result = new LinkedHashMap<>(); + Stat exists = active.checkExists().forPath(path); + if (exists == null) { + result.put("found", false); + result.put("key", path); + result.put("value", null); + result.put("metadata", null); + return result; + } + + Stat stat = new Stat(); + byte[] bytes = active.getData().storingStatIn(stat).forPath(path); + result.put("found", true); + result.put("key", path); + result.put("value", valueObject(bytes)); + result.put("metadata", statMetadata(stat)); + return result; + } + + private static Object put(JsonObject params) throws Exception { + CuratorFramework active = requireClient(); + String path = normalizePath(stringOrDefault(params, "key", "")); + requireMutablePath(path, "modified"); + byte[] bytes = parseValue(params.getAsJsonObject("value")); + + Stat stat = active.checkExists().forPath(path); + if (stat == null) { + active.create().creatingParentsIfNeeded().forPath(path, bytes); + stat = active.checkExists().forPath(path); + } else { + stat = active.setData().forPath(path, bytes); + } + + Map result = new LinkedHashMap<>(); + result.put("version", stat.getVersion()); + result.put("mtime", stat.getMtime()); + return result; + } + + private static Object delete(JsonObject params) throws Exception { + CuratorFramework active = requireClient(); + String path = normalizePath(stringOrDefault(params, "key", "")); + requireMutablePath(path, "deleted"); + boolean recursive = boolOrDefault(params, "recursive", false); + int deleted; + try { + if (active.checkExists().forPath(path) == null) { + deleted = 0; + } else if (recursive) { + deleted = countSubtree(active, path); + active.delete().deletingChildrenIfNeeded().forPath(path); + } else { + active.delete().forPath(path); + deleted = 1; + } + } catch (KeeperException.NoNodeException e) { + deleted = 0; + } + return Collections.singletonMap("deleted", deleted); + } + + private static Object dispatch(String method, JsonObject params) throws Exception { + return switch (method) { + case AgentProtocol.METHOD_HANDSHAKE -> handshakeResult(); + case AgentProtocol.METHOD_CONNECT -> connect(params); + case AgentProtocol.METHOD_TEST_CONNECTION -> testConnection(params); + case AgentProtocol.KV_METHOD_LIST_PREFIX -> listPrefix(params); + case AgentProtocol.KV_METHOD_GET -> get(params); + case AgentProtocol.KV_METHOD_PUT -> put(params); + case AgentProtocol.KV_METHOD_DELETE -> delete(params); + case AgentProtocol.METHOD_DISCONNECT -> { + closeClient(); + yield Collections.singletonMap("ok", true); + } + case AgentProtocol.METHOD_SHUTDOWN -> { + closeClient(); + System.exit(0); + yield Collections.singletonMap("ok", true); + } + default -> throw new IllegalArgumentException("Unknown method: " + method); + }; + } + + static String handleRequest(String line) { + JsonObject req = JsonParser.parseString(line).getAsJsonObject(); + JsonElement id = req.get("id"); + String method = req.get("method").getAsString(); + JsonObject params = req.has("params") && req.get("params").isJsonObject() + ? req.getAsJsonObject("params") + : new JsonObject(); + + JsonObject response = new JsonObject(); + response.addProperty("jsonrpc", "2.0"); + response.add("id", id); + + try { + Object result = dispatch(method, params); + response.add("result", GSON.toJsonTree(result)); + } catch (Exception e) { + JsonObject error = new JsonObject(); + error.addProperty("code", -1); + error.addProperty("message", e.getMessage() == null ? "Unknown error" : e.getMessage()); + response.add("error", error); + } + + return GSON.toJson(response); + } + + private static JsonObject connectionObject(JsonObject params) { + JsonElement connection = params.get("connection"); + return connection != null && connection.isJsonObject() ? connection.getAsJsonObject() : params; + } + + private static CuratorFramework requireClient() { + if (client == null) { + throw new IllegalStateException("Not connected"); + } + return client; + } + + private static void closeClient() { + if (client != null) { + client.close(); + client = null; + } + } + + static String normalizePath(String path) { + String normalized = path == null ? "" : path; + if (normalized.isEmpty() || "/".equals(normalized)) { + return "/"; + } + if (!normalized.startsWith("/")) { + normalized = "/" + normalized; + } + while (normalized.length() > 1 && normalized.endsWith("/")) { + normalized = normalized.substring(0, normalized.length() - 1); + } + return normalized; + } + + private static String childPath(String parent, String child) { + return "/".equals(parent) ? "/" + child : parent + "/" + child; + } + + private static void requireMutablePath(String path, String action) { + if ("/".equals(path)) { + throw new IllegalArgumentException("Root znode cannot be " + action); + } + } + + private static Map statMetadata(Stat stat) { + Map metadata = new LinkedHashMap<>(); + long czxid = stat.getCzxid(); + long mzxid = stat.getMzxid(); + int dataLength = stat.getDataLength(); + metadata.put("czxid", czxid); + metadata.put("mzxid", mzxid); + metadata.put("ctime", stat.getCtime()); + metadata.put("mtime", stat.getMtime()); + metadata.put("version", stat.getVersion()); + metadata.put("cversion", stat.getCversion()); + metadata.put("aversion", stat.getAversion()); + metadata.put("ephemeralOwner", stat.getEphemeralOwner()); + metadata.put("dataLength", dataLength); + metadata.put("numChildren", stat.getNumChildren()); + metadata.put("createRevision", czxid); + metadata.put("modRevision", mzxid); + metadata.put("valueSize", dataLength); + return metadata; + } + + private static Map valueObject(byte[] bytes) { + Map value = new LinkedHashMap<>(); + String utf8 = strictUtf8(bytes); + if (utf8 != null) { + value.put("encoding", "utf8"); + value.put("data", utf8); + } else { + value.put("encoding", "base64"); + value.put("data", Base64.getEncoder().encodeToString(bytes)); + } + return value; + } + + private static byte[] parseValue(JsonObject value) { + if (value == null) { + return new byte[0]; + } + String encoding = stringOrDefault(value, "encoding", "utf8"); + String data = stringOrDefault(value, "data", ""); + if ("base64".equals(encoding)) { + return Base64.getDecoder().decode(data); + } + if (!"utf8".equals(encoding)) { + throw new IllegalArgumentException("Unsupported value encoding: " + encoding); + } + return data.getBytes(StandardCharsets.UTF_8); + } + + private static String strictUtf8(byte[] bytes) { + CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT); + try { + return decoder.decode(ByteBuffer.wrap(bytes)).toString(); + } catch (CharacterCodingException e) { + return null; + } + } + + private static String encodeCursor(Cursor cursor) { + JsonObject object = new JsonObject(); + object.addProperty("root", cursor.root); + object.addProperty("recursive", cursor.recursive); + object.addProperty("offset", cursor.offset); + return Base64.getEncoder().encodeToString(GSON.toJson(object).getBytes(StandardCharsets.UTF_8)); + } + + private static Cursor decodeCursor(String continuation) { + String json = new String(Base64.getDecoder().decode(continuation), StandardCharsets.UTF_8); + JsonObject object = JsonParser.parseString(json).getAsJsonObject(); + return new Cursor( + object.get("root").getAsString(), + object.get("recursive").getAsBoolean(), + object.get("offset").getAsInt() + ); + } + + private static List listDirectChildren(CuratorFramework active, String root) throws Exception { + List paths = new ArrayList<>(); + List children = new ArrayList<>(active.getChildren().forPath(root)); + Collections.sort(children); + for (String child : children) { + paths.add(childPath(root, child)); + } + return paths; + } + + private static List listRecursive(CuratorFramework active, String root) throws Exception { + List result = new ArrayList<>(); + collectRecursive(active, root, result); + return result; + } + + private static void collectRecursive(CuratorFramework active, String path, List result) throws Exception { + List children; + try { + children = new ArrayList<>(active.getChildren().forPath(path)); + } catch (KeeperException.NoNodeException e) { + return; + } + Collections.sort(children); + for (String child : children) { + String childPath = childPath(path, child); + result.add(childPath); + collectRecursive(active, childPath, result); + } + } + + private static int countSubtree(CuratorFramework active, String path) throws Exception { + try { + if (active.checkExists().forPath(path) == null) { + return 0; + } + int count = 1; + for (String child : active.getChildren().forPath(path)) { + count += countSubtree(active, childPath(path, child)); + } + return count; + } catch (KeeperException.NoNodeException e) { + return 0; + } + } + + private static boolean shouldHideFromRootListing(String root, String path) { + return "/".equals(root) && ("/zookeeper".equals(path) || path.startsWith("/zookeeper/")); + } + + private static boolean hasTlsOptions(JsonObject connection) { + return boolOrDefault(connection, "ssl", false) + || firstNonBlank( + stringOrNull(connection, "ca_cert_path"), + stringOrNull(connection, "client_cert_path"), + stringOrNull(connection, "client_key_path"), + stringOrNull(connection, "cert_path"), + stringOrNull(connection, "key_path") + ) != null; + } + + private static String trimSlashes(String value) { + String trimmed = value.trim(); + while (trimmed.startsWith("/")) { + trimmed = trimmed.substring(1); + } + while (trimmed.endsWith("/")) { + trimmed = trimmed.substring(0, trimmed.length() - 1); + } + return trimmed; + } + + private static String stringOrNull(JsonObject object, String key) { + JsonElement element = object.get(key); + return element == null || element.isJsonNull() ? null : element.getAsString(); + } + + private static String stringOrEmpty(JsonObject object, String key) { + return stringOrDefault(object, key, ""); + } + + private static String stringOrDefault(JsonObject object, String key, String fallback) { + String value = stringOrNull(object, key); + return value == null ? fallback : value; + } + + private static int intOrDefault(JsonObject object, String key, int fallback) { + JsonElement element = object.get(key); + return element == null || element.isJsonNull() ? fallback : element.getAsInt(); + } + + private static boolean boolOrDefault(JsonObject object, String key, boolean fallback) { + JsonElement element = object.get(key); + return element == null || element.isJsonNull() ? fallback : element.getAsBoolean(); + } + + private static String firstNonBlank(String... values) { + for (String value : values) { + if (value != null && !value.isBlank()) { + return value; + } + } + return null; + } + + public static void main(String[] args) throws Exception { + // Entry point for the standalone DBX agent process. + System.setProperty("org.slf4j.simpleLogger.logFile", "System.err"); + System.out.println("{\"ready\":true}"); + System.out.flush(); + + BufferedReader reader = new BufferedReader(new InputStreamReader(System.in)); + while (true) { + String line = reader.readLine(); + if (line == null) { + break; + } + + System.out.println(handleRequest(line)); + System.out.flush(); + } + } + + private static final class Cursor { + private final String root; + private final boolean recursive; + private final int offset; + + private Cursor(String root, boolean recursive, int offset) { + this.root = root; + this.recursive = recursive; + this.offset = offset; + } + } + + private static final class HandshakeResult { + private final int protocolVersion; + private final int agentProtocolVersion; + private final List capabilities; + + private HandshakeResult(int protocolVersion, int agentProtocolVersion, List capabilities) { + this.protocolVersion = protocolVersion; + this.agentProtocolVersion = agentProtocolVersion; + this.capabilities = capabilities; + } + } +} diff --git a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java new file mode 100644 index 0000000000..d14aac9110 --- /dev/null +++ b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java @@ -0,0 +1,396 @@ +package com.dbx.agent.zookeeper; + +import com.google.gson.JsonArray; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.util.ArrayList; +import java.util.List; +import org.apache.curator.test.TestingServer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +final class ZooKeeperAgentTest { + @AfterEach + void disconnect() { + ZooKeeperAgent.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":999,\"method\":\"disconnect\",\"params\":{}}"); + } + + @Test + void handshakeAdvertisesKvCapability() { + String response = ZooKeeperAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"handshake\",\"params\":{}}" + ); + + JsonObject result = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("result"); + + Assertions.assertEquals(1, result.get("protocolVersion").getAsInt()); + Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"kv\""))); + Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"connect\""))); + Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"test_connection\""))); + } + + @Test + void kvMethodDispatchReturnsJsonRpcErrorWhenDisconnected() { + String response = ZooKeeperAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"kv_get\",\"params\":{\"key\":\"/app/name\"}}" + ); + + JsonObject error = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("error"); + + Assertions.assertEquals(-1, error.get("code").getAsInt()); + Assertions.assertEquals("Not connected", error.get("message").getAsString()); + } + + @Test + void unknownMethodReturnsJsonRpcError() { + String response = ZooKeeperAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"missing_method\",\"params\":{}}" + ); + + JsonObject error = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("error"); + + Assertions.assertEquals(-1, error.get("code").getAsInt()); + Assertions.assertEquals("Unknown method: missing_method", error.get("message").getAsString()); + } + + @Test + void connectStringUsesConfiguredValuesAndHostPortFallback() { + Assertions.assertEquals( + "zk-1:2181,zk-2:2181", + ZooKeeperAgent.connectString( + JsonParser.parseString("{\"connect_string\":\"zk-1:2181,zk-2:2181\"}").getAsJsonObject() + ) + ); + Assertions.assertEquals( + "zk-main:2181", + ZooKeeperAgent.connectString( + JsonParser.parseString("{\"zookeeper_connect_string\":\"zk-main:2181\"}").getAsJsonObject() + ) + ); + Assertions.assertEquals( + "127.0.0.1:2181", + ZooKeeperAgent.connectString(JsonParser.parseString("{}").getAsJsonObject()) + ); + Assertions.assertEquals( + "zk.local:2281", + ZooKeeperAgent.connectString(JsonParser.parseString("{\"host\":\"zk.local\",\"port\":2281}").getAsJsonObject()) + ); + } + + @Test + void buildClientRejectsTlsOptionsUntilTlsIsSupported() { + JsonObject connection = JsonParser.parseString( + "{\"ssl\":true,\"ca_cert_path\":\"ca.pem\",\"client_cert_path\":\"client.pem\",\"client_key_path\":\"client.key\"}" + ).getAsJsonObject(); + + IllegalArgumentException error = Assertions.assertThrows( + IllegalArgumentException.class, + () -> ZooKeeperAgent.buildClient(connection) + ); + + Assertions.assertEquals("ZooKeeper TLS is not supported", error.getMessage()); + } + + @Test + void normalizePathStandardizesPublicPaths() { + Assertions.assertEquals("/", ZooKeeperAgent.normalizePath("")); + Assertions.assertEquals("/", ZooKeeperAgent.normalizePath("/")); + Assertions.assertEquals("/app", ZooKeeperAgent.normalizePath("app")); + Assertions.assertEquals("/app", ZooKeeperAgent.normalizePath("/app/")); + Assertions.assertEquals("/app/name", ZooKeeperAgent.normalizePath("/app/name")); + } + + @Test + void connectAndTestConnectionWorkAgainstTestingServer() throws Exception { + try (TestingServer server = new TestingServer()) { + JsonObject connect = result(request( + 1, + "connect", + "{\"connection\":{\"connect_string\":\"" + server.getConnectString() + "\"}}" + )); + Assertions.assertTrue(connect.get("ok").getAsBoolean()); + + JsonObject test = result(request( + 2, + "test_connection", + "{\"connection\":{\"connect_string\":\"" + server.getConnectString() + "\"}}" + )); + Assertions.assertTrue(test.get("ok").getAsBoolean()); + } + } + + @Test + void testConnectionDoesNotReplaceActiveClient() throws Exception { + try (TestingServer active = new TestingServer(); TestingServer probe = new TestingServer()) { + result(request(1, "connect", "{\"connection\":{\"connect_string\":\"" + active.getConnectString() + "\"}}")); + result(request( + 2, + "kv_put", + "{\"key\":\"/active\",\"value\":{\"encoding\":\"utf8\",\"data\":\"still-here\"}}" + )); + + result(request(3, "test_connection", "{\"connection\":{\"connect_string\":\"" + probe.getConnectString() + "\"}}")); + + JsonObject get = result(request(4, "kv_get", "{\"key\":\"/active\"}")); + Assertions.assertTrue(get.get("found").getAsBoolean()); + Assertions.assertEquals("still-here", get.getAsJsonObject("value").get("data").getAsString()); + } + } + + @Test + void kvPutCreatesMissingParentsAndGetReadsUtf8Value() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + JsonObject put = result(request( + 2, + "kv_put", + "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}" + )); + Assertions.assertTrue(put.has("version")); + + JsonObject get = result(request(3, "kv_get", "{\"key\":\"/app/name\"}")); + + Assertions.assertTrue(get.get("found").getAsBoolean()); + Assertions.assertEquals("utf8", get.getAsJsonObject("value").get("encoding").getAsString()); + Assertions.assertEquals("dbx", get.getAsJsonObject("value").get("data").getAsString()); + JsonObject metadata = get.getAsJsonObject("metadata"); + Assertions.assertEquals(0, metadata.get("numChildren").getAsInt()); + Assertions.assertEquals(metadata.get("czxid").getAsLong(), metadata.get("createRevision").getAsLong()); + Assertions.assertEquals(metadata.get("mzxid").getAsLong(), metadata.get("modRevision").getAsLong()); + Assertions.assertEquals(metadata.get("dataLength").getAsInt(), metadata.get("valueSize").getAsInt()); + } + } + + @Test + void kvPutUpdatesExistingZnode() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + result(request( + 3, + "kv_put", + "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx-updated\"}}" + )); + + JsonObject value = result(request(4, "kv_get", "{\"key\":\"/app/name\"}")).getAsJsonObject("value"); + Assertions.assertEquals("dbx-updated", value.get("data").getAsString()); + } + } + + @Test + void kvGetReturnsBase64ForBinaryData() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + result(request(2, "kv_put", "{\"key\":\"/bin\",\"value\":{\"encoding\":\"base64\",\"data\":\"/wAB\"}}")); + + JsonObject value = result(request(3, "kv_get", "{\"key\":\"/bin\"}")).getAsJsonObject("value"); + Assertions.assertEquals("base64", value.get("encoding").getAsString()); + Assertions.assertEquals("/wAB", value.get("data").getAsString()); + } + } + + @Test + void kvPutRejectsRootZnode() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + JsonObject error = error(request( + 2, + "kv_put", + "{\"key\":\"/\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}" + )); + + Assertions.assertEquals("Root znode cannot be modified", error.get("message").getAsString()); + } + } + + @Test + void kvDeleteReturnsZeroForMissingZnode() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + JsonObject delete = result(request(4, "kv_delete", "{\"key\":\"/missing\"}")); + + Assertions.assertEquals(0, delete.get("deleted").getAsInt()); + } + } + + @Test + void kvDeleteRemovesLeafZnode() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/leaf\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + + JsonObject delete = result(request(3, "kv_delete", "{\"key\":\"/leaf\"}")); + JsonObject get = result(request(4, "kv_get", "{\"key\":\"/leaf\"}")); + + Assertions.assertEquals(1, delete.get("deleted").getAsInt()); + Assertions.assertFalse(get.get("found").getAsBoolean()); + } + } + + @Test + void kvDeleteRejectsNonRecursiveDeleteOfNonEmptyZnode() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + + JsonObject error = error(request(3, "kv_delete", "{\"key\":\"/app\"}")); + + Assertions.assertTrue(error.get("message").getAsString().contains("not empty")); + } + } + + @Test + void kvDeleteRecursiveDeletesFullSubtreeAndCountsZnodes() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + result(request(3, "kv_put", "{\"key\":\"/app/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); + + JsonObject delete = result(request(4, "kv_delete", "{\"key\":\"/app\",\"recursive\":true}")); + JsonObject get = result(request(5, "kv_get", "{\"key\":\"/app/name\"}")); + + Assertions.assertEquals(3, delete.get("deleted").getAsInt()); + Assertions.assertFalse(get.get("found").getAsBoolean()); + } + } + + @Test + void kvDeleteRejectsRootZnode() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + JsonObject error = error(request(2, "kv_delete", "{\"key\":\"/\",\"recursive\":true}")); + + Assertions.assertEquals("Root znode cannot be deleted", error.get("message").getAsString()); + } + } + + @Test + void listPrefixListsDirectChildrenOfRootWhenRecursiveIsFalse() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); + + JsonObject list = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\",\"recursive\":false}")); + + Assertions.assertEquals(List.of("/app", "/config"), listedKeys(list)); + Assertions.assertTrue(list.get("continuation").isJsonNull()); + } + } + + @Test + void listPrefixDefaultsToRecursiveForDbxKvBrowser() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); + + JsonObject list = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\"}")); + + Assertions.assertEquals(List.of("/app", "/app/name", "/config"), listedKeys(list)); + Assertions.assertTrue(list.get("continuation").isJsonNull()); + } + } + + @Test + void listPrefixListsRecursiveDescendantsOfRoot() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); + + JsonObject list = result(request(6, "kv_list_prefix", "{\"prefix\":\"/\",\"recursive\":true}")); + + Assertions.assertEquals(List.of("/app", "/app/name", "/config"), listedKeys(list)); + } + } + + @Test + void listPrefixReturnsEmptyResultForMissingPrefix() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + + JsonObject list = result(request(7, "kv_list_prefix", "{\"prefix\":\"/missing\"}")); + + Assertions.assertEquals(List.of(), listedKeys(list)); + Assertions.assertTrue(list.get("continuation").isJsonNull()); + } + } + + @Test + void listPrefixPaginatesBySortedPath() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/c\",\"value\":{\"encoding\":\"utf8\",\"data\":\"c\"}}")); + result(request(3, "kv_put", "{\"key\":\"/a\",\"value\":{\"encoding\":\"utf8\",\"data\":\"a\"}}")); + result(request(4, "kv_put", "{\"key\":\"/b\",\"value\":{\"encoding\":\"utf8\",\"data\":\"b\"}}")); + + JsonObject first = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\",\"limit\":2}")); + String continuation = first.get("continuation").getAsString(); + JsonObject second = result(request( + 6, + "kv_list_prefix", + "{\"prefix\":\"/\",\"limit\":2,\"continuation\":\"" + continuation + "\"}" + )); + + Assertions.assertEquals(List.of("/a", "/b"), listedKeys(first)); + Assertions.assertFalse(continuation.isBlank()); + Assertions.assertEquals(List.of("/c"), listedKeys(second)); + Assertions.assertTrue(second.get("continuation").isJsonNull()); + } + } + + @Test + void listPrefixRejectsContinuationForDifferentRequest() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/a\",\"value\":{\"encoding\":\"utf8\",\"data\":\"a\"}}")); + result(request(3, "kv_put", "{\"key\":\"/app/b\",\"value\":{\"encoding\":\"utf8\",\"data\":\"b\"}}")); + + JsonObject first = result(request(4, "kv_list_prefix", "{\"prefix\":\"/app\",\"limit\":1}")); + String continuation = first.get("continuation").getAsString(); + JsonObject error = error(request( + 5, + "kv_list_prefix", + "{\"prefix\":\"/\",\"limit\":1,\"continuation\":\"" + continuation + "\"}" + )); + + Assertions.assertEquals("Continuation does not match request", error.get("message").getAsString()); + } + } + + private static JsonObject connect(TestingServer server) { + return result(request(1, "connect", "{\"connection\":{\"connect_string\":\"" + server.getConnectString() + "\"}}")); + } + + private static String request(int id, String method, String params) { + return ZooKeeperAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":" + id + ",\"method\":\"" + method + "\",\"params\":" + params + "}" + ); + } + + private static JsonObject result(String response) { + return JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("result"); + } + + private static JsonObject error(String response) { + return JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("error"); + } + + private static List listedKeys(JsonObject listResult) { + List keys = new ArrayList<>(); + JsonArray rows = listResult.getAsJsonArray("keys"); + for (JsonElement row : rows) { + keys.add(row.getAsJsonObject().get("key").getAsString()); + } + return keys; + } +} diff --git a/agents/gradle/wrapper/gradle-wrapper.jar b/agents/gradle/wrapper/gradle-wrapper.jar new file mode 100644 index 0000000000..b1b8ef56b4 Binary files /dev/null and b/agents/gradle/wrapper/gradle-wrapper.jar differ diff --git a/agents/gradle/wrapper/gradle-wrapper.properties b/agents/gradle/wrapper/gradle-wrapper.properties new file mode 100644 index 0000000000..b52fb7e713 --- /dev/null +++ b/agents/gradle/wrapper/gradle-wrapper.properties @@ -0,0 +1,9 @@ +distributionBase=GRADLE_USER_HOME +distributionPath=wrapper/dists +distributionUrl=https\://services.gradle.org/distributions/gradle-9.5.0-bin.zip +networkTimeout=10000 +retries=0 +retryBackOffMs=500 +validateDistributionUrl=true +zipStoreBase=GRADLE_USER_HOME +zipStorePath=wrapper/dists diff --git a/agents/gradlew b/agents/gradlew new file mode 100755 index 0000000000..b9bb139f79 --- /dev/null +++ b/agents/gradlew @@ -0,0 +1,248 @@ +#!/bin/sh + +# +# Copyright © 2015 the original authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 +# + +############################################################################## +# +# Gradle start up script for POSIX generated by Gradle. +# +# Important for running: +# +# (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is +# noncompliant, but you have some other compliant shell such as ksh or +# bash, then to run this script, type that shell name before the whole +# command line, like: +# +# ksh Gradle +# +# Busybox and similar reduced shells will NOT work, because this script +# requires all of these POSIX shell features: +# * functions; +# * expansions «$var», «${var}», «${var:-default}», «${var+SET}», +# «${var#prefix}», «${var%suffix}», and «$( cmd )»; +# * compound commands having a testable exit status, especially «case»; +# * various built-in commands including «command», «set», and «ulimit». +# +# Important for patching: +# +# (2) This script targets any POSIX shell, so it avoids extensions provided +# by Bash, Ksh, etc; in particular arrays are avoided. +# +# The "traditional" practice of packing multiple parameters into a +# space-separated string is a well documented source of bugs and security +# problems, so this is (mostly) avoided, by progressively accumulating +# options in "$@", and eventually passing that to Java. +# +# Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS, +# and GRADLE_OPTS) rely on word-splitting, this is performed explicitly; +# see the in-line comments for details. +# +# There are tweaks for specific operating systems such as AIX, CygWin, +# Darwin, MinGW, and NonStop. +# +# (3) This script is generated from the Groovy template +# https://github.com/gradle/gradle/blob/3d91ce3b8caaf77ad09f381f43615b715b53f72c/platforms/jvm/plugins-application/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt +# within the Gradle project. +# +# You can find Gradle at https://github.com/gradle/gradle/. +# +############################################################################## + +# Attempt to set APP_HOME + +# Resolve links: $0 may be a link +app_path=$0 + +# Need this for daisy-chained symlinks. +while + APP_HOME=${app_path%"${app_path##*/}"} # leaves a trailing /; empty if no leading path + [ -h "$app_path" ] +do + ls=$( ls -ld "$app_path" ) + link=${ls#*' -> '} + case $link in #( + /*) app_path=$link ;; #( + *) app_path=$APP_HOME$link ;; + esac +done + +# This is normally unused +# shellcheck disable=SC2034 +APP_BASE_NAME=${0##*/} +# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036) +APP_HOME=$( cd -P "${APP_HOME:-./}" > /dev/null && printf '%s\n' "$PWD" ) || exit + +# Use the maximum available, or set MAX_FD != -1 to use that value. +MAX_FD=maximum + +warn () { + echo "$*" +} >&2 + +die () { + echo + echo "$*" + echo + exit 1 +} >&2 + +# OS specific support (must be 'true' or 'false'). +cygwin=false +msys=false +darwin=false +nonstop=false +case "$( uname )" in #( + CYGWIN* ) cygwin=true ;; #( + Darwin* ) darwin=true ;; #( + MSYS* | MINGW* ) msys=true ;; #( + NONSTOP* ) nonstop=true ;; +esac + + + +# Determine the Java command to use to start the JVM. +if [ -n "$JAVA_HOME" ] ; then + if [ -x "$JAVA_HOME/jre/sh/java" ] ; then + # IBM's JDK on AIX uses strange locations for the executables + JAVACMD=$JAVA_HOME/jre/sh/java + else + JAVACMD=$JAVA_HOME/bin/java + fi + if [ ! -x "$JAVACMD" ] ; then + die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME + +Please set the JAVA_HOME variable in your environment to match the +location of your Java installation." + fi +else + JAVACMD=java + if ! command -v java >/dev/null 2>&1 + then + die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. + +Please set the JAVA_HOME variable in your environment to match the +location of your Java installation." + fi +fi + +# Increase the maximum file descriptors if we can. +if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then + case $MAX_FD in #( + max*) + # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + MAX_FD=$( ulimit -H -n ) || + warn "Could not query maximum file descriptor limit" + esac + case $MAX_FD in #( + '' | soft) :;; #( + *) + # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + ulimit -n "$MAX_FD" || + warn "Could not set maximum file descriptor limit to $MAX_FD" + esac +fi + +# Collect all arguments for the java command, stacking in reverse order: +# * args from the command line +# * the main class name +# * -classpath +# * -D...appname settings +# * --module-path (only if needed) +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables. + +# For Cygwin or MSYS, switch paths to Windows format before running java +if "$cygwin" || "$msys" ; then + APP_HOME=$( cygpath --path --mixed "$APP_HOME" ) + + JAVACMD=$( cygpath --unix "$JAVACMD" ) + + # Now convert the arguments - kludge to limit ourselves to /bin/sh + for arg do + if + case $arg in #( + -*) false ;; # don't mess with options #( + /?*) t=${arg#/} t=/${t%%/*} # looks like a POSIX filepath + [ -e "$t" ] ;; #( + *) false ;; + esac + then + arg=$( cygpath --path --ignore --mixed "$arg" ) + fi + # Roll the args list around exactly as many times as the number of + # args, so each arg winds up back in the position where it started, but + # possibly modified. + # + # NB: a `for` loop captures its iteration list before it begins, so + # changing the positional parameters here affects neither the number of + # iterations, nor the values presented in `arg`. + shift # remove old arg + set -- "$@" "$arg" # push replacement arg + done +fi + + +# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"' + +# Collect all arguments for the java command: +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments, +# and any embedded shellness will be escaped. +# * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be +# treated as '${Hostname}' itself on the command line. + +set -- \ + "-Dorg.gradle.appname=$APP_BASE_NAME" \ + -jar "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" \ + "$@" + +# Stop when "xargs" is not available. +if ! command -v xargs >/dev/null 2>&1 +then + die "xargs is not available" +fi + +# Use "xargs" to parse quoted args. +# +# With -n1 it outputs one arg per line, with the quotes and backslashes removed. +# +# In Bash we could simply go: +# +# readarray ARGS < <( xargs -n1 <<<"$var" ) && +# set -- "${ARGS[@]}" "$@" +# +# but POSIX shell has neither arrays nor command substitution, so instead we +# post-process each arg (as a line of input to sed) to backslash-escape any +# character that might be a shell metacharacter, then use eval to reverse +# that process (while maintaining the separation between arguments), and wrap +# the whole thing up as a single "set" statement. +# +# This will of course break if any of these variables contains a newline or +# an unmatched quote. +# + +eval "set -- $( + printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" | + xargs -n1 | + sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' | + tr '\n' ' ' + )" '"$@"' + +exec "$JAVACMD" "$@" diff --git a/agents/gradlew.bat b/agents/gradlew.bat new file mode 100644 index 0000000000..aa5f10b069 --- /dev/null +++ b/agents/gradlew.bat @@ -0,0 +1,82 @@ +@rem +@rem Copyright 2015 the original author or authors. +@rem +@rem Licensed under the Apache License, Version 2.0 (the "License"); +@rem you may not use this file except in compliance with the License. +@rem You may obtain a copy of the License at +@rem +@rem https://www.apache.org/licenses/LICENSE-2.0 +@rem +@rem Unless required by applicable law or agreed to in writing, software +@rem distributed under the License is distributed on an "AS IS" BASIS, +@rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +@rem See the License for the specific language governing permissions and +@rem limitations under the License. +@rem +@rem SPDX-License-Identifier: Apache-2.0 +@rem + +@if "%DEBUG%"=="" @echo off +@rem ########################################################################## +@rem +@rem Gradle startup script for Windows +@rem +@rem ########################################################################## + +@rem Set local scope for the variables, and ensure extensions are enabled +setlocal EnableExtensions + +set DIRNAME=%~dp0 +if "%DIRNAME%"=="" set DIRNAME=. +@rem This is normally unused +set APP_BASE_NAME=%~n0 +set APP_HOME=%DIRNAME% + +@rem Resolve any "." and ".." in APP_HOME to make it shorter. +for %%i in ("%APP_HOME%") do set APP_HOME=%%~fi + +@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +set DEFAULT_JVM_OPTS="-Xmx64m" "-Xms64m" + +@rem Find java.exe +if defined JAVA_HOME goto findJavaFromJavaHome + +set JAVA_EXE=java.exe +%JAVA_EXE% -version >NUL 2>&1 +if %ERRORLEVEL% equ 0 goto execute + +echo. 1>&2 +echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 + +"%COMSPEC%" /c exit 1 + +:findJavaFromJavaHome +set JAVA_HOME=%JAVA_HOME:"=% +set JAVA_EXE=%JAVA_HOME%/bin/java.exe + +if exist "%JAVA_EXE%" goto execute + +echo. 1>&2 +echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 + +"%COMSPEC%" /c exit 1 + +:execute +@rem Setup the command line + + + +@rem Execute Gradle +@rem endlocal doesn't take effect until after the line is parsed and variables are expanded +@rem which allows us to clear the local environment before executing the java command +endlocal & "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -jar "%APP_HOME%\gradle\wrapper\gradle-wrapper.jar" %* & call :exitWithErrorLevel + +:exitWithErrorLevel +@rem Use "%COMSPEC%" /c exit to allow operators to work properly in scripts +"%COMSPEC%" /c exit %ERRORLEVEL% diff --git a/agents/opengauss/build.gradle b/agents/opengauss/build.gradle new file mode 100644 index 0000000000..b8b3af713b --- /dev/null +++ b/agents/opengauss/build.gradle @@ -0,0 +1,10 @@ +dependencies { + implementation fileTree(dir: 'libs', include: ['*.jar']) + implementation 'org.opengauss:opengauss-jdbc:7.0.0-RC3-og' +} + +tasks.named('shadowJar') { + manifest { + attributes('Agent-Label': 'openGauss', 'Main-Class': 'com.dbx.agent.opengauss.OpenGaussAgent') + } +} diff --git a/agents/opengauss/src/main/java/com/dbx/agent/opengauss/OpenGaussAgent.java b/agents/opengauss/src/main/java/com/dbx/agent/opengauss/OpenGaussAgent.java new file mode 100644 index 0000000000..985b964c3c --- /dev/null +++ b/agents/opengauss/src/main/java/com/dbx/agent/opengauss/OpenGaussAgent.java @@ -0,0 +1,20 @@ +package com.dbx.agent.opengauss; + +import com.dbx.agent.JsonRpcServer; +import com.dbx.agent.PostgresLikeAgent; +import com.dbx.agent.PostgresLikeAgentProfile; + +public final class OpenGaussAgent extends PostgresLikeAgent { + public static final PostgresLikeAgentProfile OPENGAUSS_PROFILE = new PostgresLikeAgentProfile( + "org.opengauss.Driver", + "jdbc:opengauss://{host}:{port}/{database}" + ); + + public OpenGaussAgent() { + super(OPENGAUSS_PROFILE); + } + + public static void main(String[] args) { + new JsonRpcServer(new OpenGaussAgent()).run(); + } +} diff --git a/agents/opengauss/src/test/java/com/dbx/agent/opengauss/OpenGaussAgentTest.java b/agents/opengauss/src/test/java/com/dbx/agent/opengauss/OpenGaussAgentTest.java new file mode 100644 index 0000000000..8b8452d873 --- /dev/null +++ b/agents/opengauss/src/test/java/com/dbx/agent/opengauss/OpenGaussAgentTest.java @@ -0,0 +1,21 @@ +package com.dbx.agent.opengauss; + +import com.dbx.agent.ConnectParams; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class OpenGaussAgentTest { + @Test + void buildsOpenGaussUrl() { + ConnectParams params = new ConnectParams(); + params.setHost("db.example.com"); + params.setPort(15432); + params.setDatabase("postgres"); + + assertEquals( + "jdbc:opengauss://db.example.com:15432/postgres", + OpenGaussAgent.OPENGAUSS_PROFILE.buildUrl(params) + ); + } +} diff --git a/agents/scripts/build_offline_zip.sh b/agents/scripts/build_offline_zip.sh new file mode 100755 index 0000000000..90393240dc --- /dev/null +++ b/agents/scripts/build_offline_zip.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Build offline ZIP bundles for each platform. +# Usage: ./scripts/build_offline_zip.sh +# The release-dir must contain agent-registry.json, dbx-agent-*.jar, and dbx-jre-*.tar.gz + +RELEASE_DIR="$(cd "${1:?Usage: build_offline_zip.sh }" && pwd)" + +if [ ! -f "$RELEASE_DIR/agent-registry.json" ]; then + echo "ERROR: agent-registry.json not found in $RELEASE_DIR" + exit 1 +fi + +PLATFORMS=( + "macos-aarch64" + "macos-x64" + "linux-x64" + "linux-aarch64" + "windows-x64" + "windows-aarch64" +) + +STAGING=$(mktemp -d) +trap 'rm -rf "$STAGING"' EXIT + +for platform in "${PLATFORMS[@]}"; do + WORK="$STAGING/$platform" + mkdir -p "$WORK/jre" "$WORK/drivers" + + cp "$RELEASE_DIR/agent-registry.json" "$WORK/" + + # Copy all JRE archives for this platform + JRE_COUNT=0 + for jre_file in "$RELEASE_DIR"/dbx-jre-*-"$platform".tar.gz; do + [ -f "$jre_file" ] || continue + cp "$jre_file" "$WORK/jre/" + JRE_COUNT=$((JRE_COUNT + 1)) + done + + if [ "$JRE_COUNT" -eq 0 ]; then + echo "SKIP $platform (no JRE found)" + rm -rf "$WORK" + continue + fi + + # Copy all driver JARs (platform-independent) + for jar_file in "$RELEASE_DIR"/dbx-agent-*.jar; do + [ -f "$jar_file" ] || continue + cp "$jar_file" "$WORK/drivers/" + done + + for native_file in "$RELEASE_DIR"/dbx-agent-*-"$platform" "$RELEASE_DIR"/dbx-agent-*-"$platform".exe; do + [ -f "$native_file" ] || continue + cp "$native_file" "$WORK/drivers/" + done + + ZIP_NAME="dbx-agents-offline-${platform}.zip" + (cd "$WORK" && zip -r "$RELEASE_DIR/$ZIP_NAME" agent-registry.json jre/ drivers/) + SIZE=$(du -h "$RELEASE_DIR/$ZIP_NAME" | cut -f1) + echo "Created $ZIP_NAME ($SIZE)" +done + +echo "Done." diff --git a/agents/scripts/release.sh b/agents/scripts/release.sh new file mode 100755 index 0000000000..5885068229 --- /dev/null +++ b/agents/scripts/release.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +set -euo pipefail + +# 自动打 tag 并推送触发发布 +# 用法: +# ./scripts/release.sh # 自动递增 patch 版本 +# ./scripts/release.sh minor # 递增 minor 版本 +# ./scripts/release.sh major # 递增 major 版本 +# ./scripts/release.sh agents-v0.3.0 # 手动指定版本号 + +BUMP="${1:-patch}" + +# 获取最新的 agents-v 开头的 tag +latest_tag=$(git tag --sort=-v:refname | grep '^agents-v' | head -1) + +if [ -z "$latest_tag" ]; then + echo "未找到任何 Agent 版本 tag，从 agents-v0.1.0 开始" + latest_tag="agents-v0.0.0" +fi + +echo "当前最新 tag: $latest_tag" + +# 去掉 agents-v 前缀 +version="${latest_tag#agents-v}" +IFS='.' read -r major minor patch <<< "$version" + +# 计算新版本号 +case "$BUMP" in + major) + new_tag="agents-v$((major + 1)).0.0" + ;; + minor) + new_tag="agents-v${major}.$((minor + 1)).0" + ;; + patch) + new_tag="agents-v${major}.${minor}.$((patch + 1))" + ;; + agents-v*) + new_tag="$BUMP" + ;; + *) + echo "错误: 未知的版本递增方式 '$BUMP'，请使用 patch/minor/major 或 agents-vX.Y.Z" + exit 1 + ;; +esac + +echo "新版本: $new_tag" + +# 确认 +read -r -p "确认创建并推送 tag ${new_tag}? [y/N] " confirm +if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then + echo "已取消" + exit 0 +fi + +# 创建并推送 tag +git tag "$new_tag" +git push origin "$new_tag" + +echo "完成! tag $new_tag 已推送，CI 将自动触发发布流程。" diff --git a/agents/scripts/validate_agent_jars.py b/agents/scripts/validate_agent_jars.py new file mode 100644 index 0000000000..a51e3295ac --- /dev/null +++ b/agents/scripts/validate_agent_jars.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +import sys +import zipfile +from email.parser import Parser +from pathlib import Path + + +def agent_modules(root: Path) -> list[str]: + return sorted( + path.parents[2].name + for path in list(root.glob("*/build/libs/dbx-agent-*.jar")) + + list(root.glob("drivers/*/build/libs/dbx-agent-*.jar")) + if path.is_file() + ) + + +def module_dir(root: Path, module: str) -> Path: + nested = root / "drivers" / module + if nested.exists(): + return nested + return root / module + + +def validate_agent_jars(root: Path) -> list[str]: + problems: list[str] = [] + for module in agent_modules(root): + jar = module_dir(root, module) / "build/libs" / f"dbx-agent-{module}.jar" + relative = jar.relative_to(root) + try: + with zipfile.ZipFile(jar) as archive: + manifest = read_manifest(archive) + main_class = manifest.get("Main-Class") + if not main_class: + problems.append(f"{relative}: missing Main-Class manifest attribute") + continue + class_entry = main_class.replace(".", "/") + ".class" + if class_entry not in archive.namelist(): + problems.append(f"{relative}: Main-Class class not found: {class_entry}") + if not manifest.get("Agent-Label"): + problems.append(f"{relative}: missing Agent-Label manifest attribute") + except zipfile.BadZipFile: + problems.append(f"{relative}: invalid jar") + return problems + + +def read_manifest(archive: zipfile.ZipFile) -> dict[str, str]: + try: + raw = archive.read("META-INF/MANIFEST.MF").decode("utf-8") + except KeyError: + return {} + return dict(Parser().parsestr(raw)) + + +def main() -> int: + root = Path.cwd() + problems = validate_agent_jars(root) + if problems: + print("Agent jar validation failed:") + for problem in problems: + print(f"- {problem}") + return 1 + + print("Agent jar validation passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/agents/scripts/validate_agents.py b/agents/scripts/validate_agents.py new file mode 100644 index 0000000000..8d2398dcd4 --- /dev/null +++ b/agents/scripts/validate_agents.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +import json +import re +import sys +from pathlib import Path +from typing import Optional + + +INFRA_MODULES = {"common", "test-support"} +SOURCE_GLOBS = ("*/src/main/**/*.java", "drivers/*/src/main/**/*.java") +KOTLIN_FILE_SUFFIXES = (".kt", ".kts") +KOTLIN_SCAN_EXCLUDED_PARTS = {".git", ".gradle", "build"} +DEFAULT_AGENT_JRE_KEY = "21" +LEGACY_ORACLE_JRE_KEY = "8" +NON_JDBC_AGENT_MODULES = {"mongodb", "etcd", "zookeeper"} +NATIVE_ONLY_AGENT_MODULES = {"xugu"} +JDBC_ARCHITECTURE_ALLOWLIST = { + "access": "custom Access metadata and URL behavior pending migration", + "dameng": "custom Dameng metadata and DDL pending migration", + "db2": "custom DB2 metadata pending migration", + "gaussdb": "custom GaussDB metadata pending migration", + "gbase8s": "custom GBase 8s metadata pending migration", + "goldendb": "custom GoldenDB metadata pending migration", + "informix": "custom Informix metadata pending migration", + "neo4j": "custom Neo4j transaction/query behavior pending migration", + "oracle": "custom Oracle metadata and connection properties pending migration", + "oracle-legacy": "custom Oracle legacy metadata and connection properties pending migration", + "oracle-10g": "custom Oracle 10g metadata and Java 8 runtime pending migration", + "sundb": "custom SunDB metadata pending migration", + "tdengine": "custom TDengine WebSocket JDBC behavior pending migration", +} +APPROVED_JDBC_BASES = { + "AbstractJdbcAgent", + "ConfiguredJdbcAgent", + "PostgresLikeAgent", +} + +FORBIDDEN_PATTERNS = [ + (re.compile(r"private\s+val\s+QUERY_PREFIXES"), "forbidden local SQL prefix classifier"), + (re.compile(r"private\s+const\s+val\s+MAX_ROWS"), "forbidden local result row cap"), + (re.compile(r"executeUpdate\s*\(\s*trimmedSql\s*\)"), "forbidden executeUpdate(trimmedSql) in query execution"), + (re.compile(r"truncated\s*=\s*rows\.size\s*>=\s*MAX_ROWS"), "forbidden inclusive truncation check"), + (re.compile(r"val\s+truncated\s*=\s*rs\.next\s*\(\s*\)"), "forbidden extra ResultSet next() truncation probe"), +] + +COPIED_JDBC_INFRASTRUCTURE_PATTERNS = [ + (re.compile(r"\bClass\.forName\s*\("), "copied driver loading"), + (re.compile(r"\bDriverManager\.getConnection\s*\("), "copied JDBC connection creation"), + (re.compile(r"public\s+void\s+disconnect\s*\("), "copied JDBC disconnect lifecycle"), + (re.compile(r"JdbcExecutor\.INSTANCE\.execute\s*\("), "copied JDBC query execution"), + (re.compile(r"private\s+Object\s+(?:getResultValue|resultValue|stringResultValue)\s*\("), "copied result value reader"), +] + + +def included_agent_modules(root: Path) -> set[str]: + settings = (root / "settings.gradle").read_text(encoding="utf-8") + include_args = "\n".join( + parenthesized or bare + for parenthesized, bare in re.findall( + r"\binclude\b\s*(?:\((.*?)\)|([^\n]+))", settings, flags=re.S + ) + ) + included = set(re.findall(r"""['"]([^'"]+)['"]""", include_args)) + for variable in ("infrastructureModules", "driverModules"): + match = re.search(rf"\bdef\s+{variable}\s*=\s*\[(.*?)\]", settings, flags=re.S) + if match: + included.update(re.findall(r"""['"]([^'"]+)['"]""", match.group(1))) + return included - INFRA_MODULES + + +def agent_modules(root: Path) -> set[str]: + native = {name for name in NATIVE_ONLY_AGENT_MODULES if module_dir(root, name).exists()} + return included_agent_modules(root) | native + + +def module_dir(root: Path, module: str) -> Path: + nested = root / "drivers" / module + if nested.exists(): + return nested + return root / module + + +def module_relative_path(root: Path, module: str, *parts: str) -> Path: + return module_dir(root, module).joinpath(*parts) + + +def validate_versions(root: Path) -> list[str]: + included = agent_modules(root) + versions = set(json.loads((root / "versions.json").read_text(encoding="utf-8"))) + return ( + [f"included module missing version: {name}" for name in sorted(included - versions)] + + [f"versions key not included: {name}" for name in sorted(versions - included)] + ) + + +def validate_source_patterns(root: Path) -> list[str]: + problems: list[str] = [] + sources = sorted(source for glob in SOURCE_GLOBS for source in root.glob(glob)) + for source in sources: + relative = source.relative_to(root) + text = source.read_text(encoding="utf-8") + for line_number, line in enumerate(text.splitlines(), start=1): + for pattern, message in FORBIDDEN_PATTERNS: + if pattern.search(line): + problems.append(f"{relative}:{line_number}: {message}") + return problems + + +def validate_no_kotlin_residue(root: Path) -> list[str]: + problems: list[str] = [] + for path in sorted(root.rglob("*")): + if not path.is_file(): + continue + relative = path.relative_to(root) + if any(part in KOTLIN_SCAN_EXCLUDED_PARTS for part in relative.parts): + continue + if path.suffix in KOTLIN_FILE_SUFFIXES: + problems.append(f"{relative}: forbidden Kotlin file") + return problems + + +def validate_manifest_fields(root: Path, modules: set[str]) -> list[str]: + problems: list[str] = [] + root_build_text = (root / "build.gradle").read_text(encoding="utf-8") if (root / "build.gradle").exists() else "" + archive_convention = re.search( + r"archiveBaseName\s*=\s*['\"]dbx-agent-\$\{project\.name\}['\"]", + root_build_text, + ) + for module in sorted(modules): + if module in NATIVE_ONLY_AGENT_MODULES: + continue + build_file = module_relative_path(root, module, "build.gradle") + relative = build_file.relative_to(root) + if not build_file.exists(): + problems.append(f"{module}: missing build.gradle") + continue + + text = build_file.read_text(encoding="utf-8") + expected_archive = f"archiveBaseName = 'dbx-agent-{module}'" + archive_pattern = re.compile( + rf"archiveBaseName\s*=\s*['\"]dbx-agent-{re.escape(module)}['\"]" + ) + if not archive_convention and not archive_pattern.search(text): + problems.append(f"{relative}: missing {expected_archive}") + attrs = manifest_attributes(text) + if "Agent-Label" not in attrs: + problems.append(f"{relative}: missing Agent-Label manifest attribute") + main_class = attrs.get("Main-Class") + if main_class is None: + problems.append(f"{relative}: missing Main-Class manifest attribute") + else: + main_source = module_relative_path( + root, + module, + "src/main/java", + str(Path(*main_class.split(".")).with_suffix(".java")), + ) + if not main_source.exists(): + problems.append(f"{relative}: Main-Class source not found: {main_source.relative_to(root)}") + return problems + + +def validate_jdbc_architecture(root: Path, modules: set[str]) -> list[str]: + problems: list[str] = [] + for module in sorted(modules - NON_JDBC_AGENT_MODULES - NATIVE_ONLY_AGENT_MODULES): + source = main_class_source(root, module) + if source is None or not source.exists(): + continue + text = source.read_text(encoding="utf-8") + base_match = re.search(r"\bextends\s+([A-Za-z0-9_]+)", text) + base = base_match.group(1) if base_match else "" + if module not in JDBC_ARCHITECTURE_ALLOWLIST and base not in APPROVED_JDBC_BASES: + problems.append( + f"{source.relative_to(root)}: JDBC agents must extend AbstractJdbcAgent, ConfiguredJdbcAgent, or PostgresLikeAgent" + ) + if module in JDBC_ARCHITECTURE_ALLOWLIST: + continue + for line_number, line in enumerate(text.splitlines(), start=1): + for pattern, message in COPIED_JDBC_INFRASTRUCTURE_PATTERNS: + if pattern.search(line): + problems.append(f"{source.relative_to(root)}:{line_number}: {message}; use shared JDBC foundation") + return problems + + +def validate_authoring_template(root: Path) -> list[str]: + template = root / "docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java" + if not template.exists(): + return [] + text = template.read_text(encoding="utf-8") + problems: list[str] = [] + if "extends BaseDatabaseAgent" in text: + problems.append(f"{template.relative_to(root)}: template must use shared JDBC foundation") + for pattern, message in COPIED_JDBC_INFRASTRUCTURE_PATTERNS: + if pattern.search(text): + problems.append(f"{template.relative_to(root)}: template contains {message}") + return problems + + +def main_class_source(root: Path, module: str) -> Optional[Path]: + build_file = module_relative_path(root, module, "build.gradle") + if not build_file.exists(): + return None + attrs = manifest_attributes(build_file.read_text(encoding="utf-8")) + main_class = attrs.get("Main-Class") + if main_class is None: + return None + return module_relative_path(root, module, "src/main/java", str(Path(*main_class.split(".")).with_suffix(".java"))) + + +def validate_release_runtime_keys(root: Path) -> list[str]: + workflow_candidates = [ + root.parent / ".github/workflows/agents-release.yml", + root / ".github/workflows/agents-release.yml", + root / ".github/workflows/release.yml", + ] + workflow = next((candidate for candidate in workflow_candidates if candidate.exists()), workflow_candidates[-1]) + if not workflow.exists(): + return [] + text = workflow.read_text(encoding="utf-8") + problems: list[str] = [] + required_patterns = [ + ( + rf'jre-key:\s*"{DEFAULT_AGENT_JRE_KEY}"', + f"release workflow must build the default JRE with key {DEFAULT_AGENT_JRE_KEY}", + ), + ( + rf'java-version:\s*"{DEFAULT_AGENT_JRE_KEY}"', + f"release workflow must build the default JRE from Java {DEFAULT_AGENT_JRE_KEY}", + ), + ( + rf'oracle-10g\)\s*echo\s*"{LEGACY_ORACLE_JRE_KEY}"', + f"oracle-10g must keep JRE key {LEGACY_ORACLE_JRE_KEY}", + ), + ( + rf'\*\)\s*echo\s*"{DEFAULT_AGENT_JRE_KEY}"', + f"non-legacy agents must use JRE key {DEFAULT_AGENT_JRE_KEY}", + ), + ( + rf'"{DEFAULT_AGENT_JRE_KEY}":\s*\{{\s*"version":\s*"{DEFAULT_AGENT_JRE_KEY}\.', + f"registry must publish Java {DEFAULT_AGENT_JRE_KEY} under JRE key {DEFAULT_AGENT_JRE_KEY}", + ), + ( + r"legacy-placeholder\.jar", + "native-only registry entries must publish a legacy jar placeholder for older DBX clients", + ), + ] + for pattern, message in required_patterns: + if not re.search(pattern, text, flags=re.S): + problems.append(message) + forbidden_patterns = [ + (r'jre-key:\s*"17"', "release workflow must not build Java 21 under JRE key 17"), + (r'\*\)\s*echo\s*"17"', "non-legacy agents must not use JRE key 17"), + (r'"17":\s*\{\s*"version":\s*"21\.', "registry must not publish Java 21 under JRE key 17"), + ] + for pattern, message in forbidden_patterns: + if re.search(pattern, text, flags=re.S): + problems.append(message) + return problems + + +def manifest_attributes(build_text: str) -> dict[str, str]: + attrs: dict[str, str] = {} + for key, value in re.findall( + r"['\"]([^'\"]+)['\"]\s*:\s*['\"]([^'\"]+)['\"]", + build_text, + ): + attrs[key] = value + return attrs + + +def validate(root: Path) -> list[str]: + modules = agent_modules(root) + return ( + validate_versions(root) + + validate_source_patterns(root) + + validate_manifest_fields(root, modules) + + validate_jdbc_architecture(root, modules) + + validate_authoring_template(root) + + validate_release_runtime_keys(root) + + validate_no_kotlin_residue(root) + ) + + +def main() -> int: + root = Path.cwd() + problems = validate(root) + if problems: + print("Agent validation failed:") + for problem in problems: + print(f"- {problem}") + return 1 + + print("Agent validation passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/agents/scripts/validate_agents_test.py b/agents/scripts/validate_agents_test.py new file mode 100644 index 0000000000..8c6dedbb15 --- /dev/null +++ b/agents/scripts/validate_agents_test.py @@ -0,0 +1,437 @@ +import json +import tempfile +import textwrap +import unittest +import zipfile +from pathlib import Path + +import validate_agents +import validate_agent_jars + + +class ValidateAgentsTest(unittest.TestCase): + def test_versions_must_match_included_agent_modules(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "settings.gradle").write_text( + "include 'common', 'test-support', 'h2', 'oracle'\n", + encoding="utf-8", + ) + (root / "versions.json").write_text( + json.dumps({"h2": "0.1.0", "mongodb": "0.1.0"}), + encoding="utf-8", + ) + + problems = validate_agents.validate_versions(root) + + self.assertEqual( + [ + "included module missing version: oracle", + "versions key not included: mongodb", + ], + problems, + ) + + def test_versions_support_driver_modules_variable(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "settings.gradle").write_text( + "def infrastructureModules = ['common', 'test-support']\n" + "def driverModules = ['h2', 'oracle']\n" + "include(*(infrastructureModules + driverModules))\n", + encoding="utf-8", + ) + (root / "versions.json").write_text( + json.dumps({"h2": "0.1.0", "oracle": "0.1.0"}), + encoding="utf-8", + ) + + self.assertEqual([], validate_agents.validate_versions(root)) + + def test_source_scan_rejects_old_execute_query_patterns(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + source = root / "h2/src/main/java/com/dbx/agent/h2/H2Agent.java" + source.parent.mkdir(parents=True) + source.write_text( + textwrap.dedent( + """ + private val QUERY_PREFIXES = listOf("SELECT") + fun executeQuery(sql: String) { + val trimmedSql = sql.trim() + stmt.executeUpdate(trimmedSql) + } + """ + ), + encoding="utf-8", + ) + + problems = validate_agents.validate_source_patterns(root) + + self.assertEqual( + [ + "h2/src/main/java/com/dbx/agent/h2/H2Agent.java:2: forbidden local SQL prefix classifier", + "h2/src/main/java/com/dbx/agent/h2/H2Agent.java:5: forbidden executeUpdate(trimmedSql) in query execution", + ], + problems, + ) + + def test_jdbc_architecture_requires_shared_base_for_new_agents(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + module = root / "example" + source = module / "src/main/java/com/dbx/agent/example/ExampleAgent.java" + source.parent.mkdir(parents=True) + (module / "build.gradle").write_text( + textwrap.dedent( + """ + tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'Example', + 'Main-Class': 'com.dbx.agent.example.ExampleAgent' + ) + } + } + """ + ), + encoding="utf-8", + ) + source.write_text( + textwrap.dedent( + """ + package com.dbx.agent.example; + + import com.dbx.agent.BaseDatabaseAgent; + import java.sql.DriverManager; + + public final class ExampleAgent extends BaseDatabaseAgent { + public void connect() throws Exception { + Class.forName("example.Driver"); + DriverManager.getConnection("jdbc:example"); + } + } + """ + ), + encoding="utf-8", + ) + + problems = validate_agents.validate_jdbc_architecture(root, {"example"}) + + self.assertEqual( + [ + "example/src/main/java/com/dbx/agent/example/ExampleAgent.java: JDBC agents must extend AbstractJdbcAgent, ConfiguredJdbcAgent, or PostgresLikeAgent", + "example/src/main/java/com/dbx/agent/example/ExampleAgent.java:9: copied driver loading; use shared JDBC foundation", + "example/src/main/java/com/dbx/agent/example/ExampleAgent.java:10: copied JDBC connection creation; use shared JDBC foundation", + ], + problems, + ) + + def test_jdbc_architecture_allows_documented_migration_exceptions(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + module = root / "oracle" + source = module / "src/main/java/com/dbx/agent/oracle/OracleAgent.java" + source.parent.mkdir(parents=True) + (module / "build.gradle").write_text( + textwrap.dedent( + """ + tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'Oracle', + 'Main-Class': 'com.dbx.agent.oracle.OracleAgent' + ) + } + } + """ + ), + encoding="utf-8", + ) + source.write_text( + "package com.dbx.agent.oracle; public final class OracleAgent extends BaseDatabaseAgent {}\n", + encoding="utf-8", + ) + + self.assertEqual([], validate_agents.validate_jdbc_architecture(root, {"oracle"})) + + def test_authoring_template_must_use_shared_foundation(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + template = root / "docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java" + template.parent.mkdir(parents=True) + template.write_text( + textwrap.dedent( + """ + package com.dbx.agent.template; + import com.dbx.agent.BaseDatabaseAgent; + import java.sql.DriverManager; + public final class TemplateAgent extends BaseDatabaseAgent { + void connect() throws Exception { + DriverManager.getConnection("jdbc:template"); + } + } + """ + ), + encoding="utf-8", + ) + + self.assertEqual( + [ + "docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java: template must use shared JDBC foundation", + "docs/examples/jdbc-agent-template/src/main/java/com/dbx/agent/template/TemplateAgent.java: template contains copied JDBC connection creation", + ], + validate_agents.validate_authoring_template(root), + ) + + def test_manifest_validation_requires_registry_fields(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + module = root / "h2" + module.mkdir() + (root / "build.gradle").write_text( + textwrap.dedent( + """ + configure(agentProjects) { + tasks.named('shadowJar') { + archiveBaseName = "dbx-agent-${project.name}" + } + } + """ + ), + encoding="utf-8", + ) + (module / "build.gradle").write_text( + textwrap.dedent( + """ + tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'H2' + ) + } + } + """ + ), + encoding="utf-8", + ) + + problems = validate_agents.validate_manifest_fields(root, {"h2"}) + + self.assertEqual( + ["h2/build.gradle: missing Main-Class manifest attribute"], + problems, + ) + + def test_manifest_validation_requires_main_class_source(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + module = root / "h2" + module.mkdir() + (root / "build.gradle").write_text( + 'archiveBaseName = "dbx-agent-${project.name}"\n', + encoding="utf-8", + ) + (module / "build.gradle").write_text( + textwrap.dedent( + """ + tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'H2', + 'Main-Class': 'com.dbx.agent.h2.H2Agent' + ) + } + } + """ + ), + encoding="utf-8", + ) + + problems = validate_agents.validate_manifest_fields(root, {"h2"}) + + self.assertEqual( + ["h2/build.gradle: Main-Class source not found: h2/src/main/java/com/dbx/agent/h2/H2Agent.java"], + problems, + ) + + source = module / "src/main/java/com/dbx/agent/h2/H2Agent.java" + source.parent.mkdir(parents=True) + source.write_text("package com.dbx.agent.h2; public final class H2Agent {}", encoding="utf-8") + + self.assertEqual([], validate_agents.validate_manifest_fields(root, {"h2"})) + + def test_manifest_validation_supports_drivers_directory(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + module = root / "drivers/h2" + source = module / "src/main/java/com/dbx/agent/h2/H2Agent.java" + source.parent.mkdir(parents=True) + source.write_text("package com.dbx.agent.h2; public final class H2Agent {}", encoding="utf-8") + (root / "build.gradle").write_text( + 'archiveBaseName = "dbx-agent-${project.name}"\n', + encoding="utf-8", + ) + (module / "build.gradle").write_text( + textwrap.dedent( + """ + tasks.named('shadowJar') { + manifest { + attributes( + 'Agent-Label': 'H2', + 'Main-Class': 'com.dbx.agent.h2.H2Agent' + ) + } + } + """ + ), + encoding="utf-8", + ) + + self.assertEqual([], validate_agents.validate_manifest_fields(root, {"h2"})) + + def test_kotlin_residue_scan_rejects_kt_and_kts_files_outside_build_dirs(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + forbidden_source = root / "h2/src/main/kotlin/com/dbx/agent/h2/H2Agent.kt" + forbidden_build = root / "settings.gradle.kts" + ignored_build_output = root / "h2/build/tmp/Generated.kt" + ignored_gradle_cache = root / ".gradle/caches/init.gradle.kts" + for path in ( + forbidden_source, + forbidden_build, + ignored_build_output, + ignored_gradle_cache, + ): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("", encoding="utf-8") + + problems = validate_agents.validate_no_kotlin_residue(root) + + self.assertEqual( + [ + "h2/src/main/kotlin/com/dbx/agent/h2/H2Agent.kt: forbidden Kotlin file", + "settings.gradle.kts: forbidden Kotlin file", + ], + problems, + ) + + def test_release_runtime_keys_match_java_21_default_and_oracle_8_exception(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + workflow = root / ".github/workflows/release.yml" + workflow.parent.mkdir(parents=True) + workflow.write_text( + textwrap.dedent( + ''' + strategy: + matrix: + include: + - jre-key: "21" + java-version: "21" + detect_jre_key() { + case "$name" in + oracle-10g) echo "8" ;; + *) echo "21" ;; + esac + } + cat > release/agent-registry.json < release/agent-registry.json < + project(":${name}").projectDir = file("drivers/${name}") +} diff --git a/agents/test-support/build.gradle b/agents/test-support/build.gradle new file mode 100644 index 0000000000..0769fc7fd6 --- /dev/null +++ b/agents/test-support/build.gradle @@ -0,0 +1,4 @@ +dependencies { + api 'org.junit.jupiter:junit-jupiter-api:5.11.4' + implementation project(':common') +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcAgentFake.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcAgentFake.java new file mode 100644 index 0000000000..7665fea330 --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcAgentFake.java @@ -0,0 +1,177 @@ +package com.dbx.agent.test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.sql.Types; +import java.util.ArrayList; +import java.util.List; + +public final class JdbcAgentFake { + public static final JdbcAgentFake INSTANCE = new JdbcAgentFake(); + public static final List calls = new ArrayList(); + + private JdbcAgentFake() { + } + + public static Connection connection() { + calls.clear(); + Statement statement = statement(); + return proxy(Connection.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("createStatement".equals(name)) { + return statement; + } + if ("getAutoCommit".equals(name)) { + return true; + } + if ("setAutoCommit".equals(name) || "commit".equals(name) || "rollback".equals(name) || "close".equals(name)) { + return null; + } + if ("isClosed".equals(name)) { + return false; + } + return defaultValue(method.getReturnType()); + } + }); + } + + public List getCalls() { + return calls; + } + + private static Statement statement() { + ResultSet resultSet = resultSet(); + return proxy(Statement.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("execute".equals(name)) { + calls.add("execute"); + return true; + } + if ("executeQuery".equals(name)) { + calls.add("executeQuery"); + return resultSet; + } + if ("executeUpdate".equals(name)) { + calls.add("executeUpdate"); + return 0; + } + if ("getResultSet".equals(name)) { + return resultSet; + } + if ("getUpdateCount".equals(name)) { + return 0; + } + if ("setMaxRows".equals(name)) { + calls.add("setMaxRows:" + args[0]); + return null; + } + if ("setFetchSize".equals(name)) { + calls.add("setFetchSize:" + args[0]); + return null; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet resultSet() { + final int[] index = {-1}; + ResultSetMetaData metadata = metadata(); + return proxy(ResultSet.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + index[0] += 1; + return index[0] == 0; + } + if ("getMetaData".equals(name)) { + return metadata; + } + if ("getObject".equals(name) || "getString".equals(name)) { + return "row-value"; + } + if ("wasNull".equals(name)) { + return false; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSetMetaData metadata() { + return proxy(ResultSetMetaData.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("getColumnCount".equals(name)) { + return 1; + } + if ("getColumnLabel".equals(name)) { + return "VALUE"; + } + if ("getColumnType".equals(name)) { + return Types.VARCHAR; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static T proxy(Class type, final MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcConnectedAgentTest.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcConnectedAgentTest.java new file mode 100644 index 0000000000..664e3cb71b --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcConnectedAgentTest.java @@ -0,0 +1,18 @@ +package com.dbx.agent.test; + +import com.dbx.agent.DatabaseAgent; + +import java.util.function.Consumer; + +public abstract class JdbcConnectedAgentTest { + protected abstract DatabaseAgent createConnectedAgent(String databaseName); + + protected void withAgent(String databaseName, Consumer block) { + DatabaseAgent agent = createConnectedAgent(databaseName); + try { + block.accept(agent); + } finally { + agent.disconnect(); + } + } +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcExecutionBehaviorTest.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcExecutionBehaviorTest.java new file mode 100644 index 0000000000..c15d377dcc --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcExecutionBehaviorTest.java @@ -0,0 +1,148 @@ +package com.dbx.agent.test; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.JdbcExecutor; +import com.dbx.agent.QueryPageOptions; +import com.dbx.agent.QueryPageResult; +import com.dbx.agent.QueryResult; +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public abstract class JdbcExecutionBehaviorTest extends JdbcConnectedAgentTest { + protected abstract String resultSetSql(); + + protected abstract List expectedResultSetColumns(); + + protected abstract List> expectedResultSetRows(); + + protected abstract String rowsSql(int rowCount); + + @Test + void executesStatementsThatReturnResultSetsWithoutSelectPrefix() { + withAgent("dbx-agent-result-set", agent -> { + QueryResult result = agent.executeQuery(resultSetSql(), null, new ExecuteQueryOptions()); + + assertEquals(expectedResultSetColumns(), result.getColumns()); + assertEquals(expectedResultSetRows(), result.getRows()); + assertEquals(0, result.getAffected_rows()); + assertFalse(result.getTruncated()); + }); + } + + @Test + void doesNotMarkExactlyMaxRowsAsTruncated() { + withAgent("dbx-agent-limit", agent -> { + QueryResult result = agent.executeQuery(rowsSql(JdbcExecutor.DEFAULT_MAX_ROWS), null, new ExecuteQueryOptions()); + + assertEquals(JdbcExecutor.DEFAULT_MAX_ROWS, result.getRows().size()); + assertFalse(result.getTruncated()); + }); + } + + @Test + void marksResultsBeyondMaxRowsAsTruncated() { + withAgent("dbx-agent-truncated", agent -> { + QueryResult result = agent.executeQuery(rowsSql(JdbcExecutor.DEFAULT_MAX_ROWS + 1), null, new ExecuteQueryOptions()); + + assertEquals(JdbcExecutor.DEFAULT_MAX_ROWS, result.getRows().size()); + assertTrue(result.getTruncated()); + }); + } + + @Test + void fetchesSubsequentResultPagesFromTheSameJdbcSession() { + withAgent("dbx-agent-session-pages", agent -> { + QueryPageResult first = agent.executeQueryPage(rowsSql(5), null, new QueryPageOptions(2, 2, 10)); + + assertEquals(Arrays.asList(Collections.singletonList(1L), Collections.singletonList(2L)), first.getRows()); + assertTrue(first.getHas_more()); + assertFalse(first.getTruncated()); + String sessionId = first.getSession_id(); + assertNotNull(sessionId); + + QueryPageResult second = agent.fetchQueryPage(sessionId, 2); + + assertEquals(Arrays.asList(Collections.singletonList(3L), Collections.singletonList(4L)), second.getRows()); + assertTrue(second.getHas_more()); + assertFalse(second.getTruncated()); + + QueryPageResult third = agent.fetchQueryPage(sessionId, 2); + + assertEquals(Collections.singletonList(Collections.singletonList(5L)), third.getRows()); + assertFalse(third.getHas_more()); + assertFalse(third.getTruncated()); + }); + } + + @Test + void closesJdbcResultSessionsExplicitly() { + withAgent("dbx-agent-session-close", agent -> { + QueryPageResult first = agent.executeQueryPage(rowsSql(5), null, new QueryPageOptions(2, 2, 10)); + String sessionId = first.getSession_id(); + assertNotNull(sessionId); + + assertTrue(agent.closeQuerySession(sessionId)); + assertFalse(agent.closeQuerySession(sessionId)); + assertThrows(IllegalArgumentException.class, () -> agent.fetchQueryPage(sessionId, 2)); + }); + } + + @Test + void expiresIdleJdbcResultSessions() { + withAgent("dbx-agent-session-expiry", agent -> { + QueryPageResult first = agent.executeQueryPage(rowsSql(5), null, new QueryPageOptions(2, 2, 10)); + String sessionId = first.getSession_id(); + assertNotNull(sessionId); + + assertEquals(1, JdbcExecutor.INSTANCE.expireIdleQuerySessions(System.currentTimeMillis(), 0L)); + assertThrows(IllegalArgumentException.class, () -> agent.fetchQueryPage(sessionId, 2)); + }); + } + + @Test + void handlesTransactionControlStatementsConsistently() { + withAgent("dbx-agent-transaction", agent -> { + Connection connection = agent.getConnection(); + assertNotNull(connection); + + try { + QueryResult begin = agent.executeQuery("BEGIN", null, new ExecuteQueryOptions()); + assertFalse(connection.getAutoCommit()); + assertEquals(Collections.emptyList(), begin.getColumns()); + assertEquals(Collections.emptyList(), begin.getRows()); + assertEquals(0, begin.getAffected_rows()); + assertFalse(begin.getTruncated()); + + QueryResult commit = agent.executeQuery("COMMIT", null, new ExecuteQueryOptions()); + assertTrue(connection.getAutoCommit()); + assertEquals(Collections.emptyList(), commit.getColumns()); + assertEquals(Collections.emptyList(), commit.getRows()); + assertEquals(0, commit.getAffected_rows()); + assertFalse(commit.getTruncated()); + + agent.executeQuery("BEGIN", null, new ExecuteQueryOptions()); + assertFalse(connection.getAutoCommit()); + + QueryResult rollback = agent.executeQuery("ROLLBACK", null, new ExecuteQueryOptions()); + assertTrue(connection.getAutoCommit()); + assertEquals(Collections.emptyList(), rollback.getColumns()); + assertEquals(Collections.emptyList(), rollback.getRows()); + assertEquals(0, rollback.getAffected_rows()); + assertFalse(rollback.getTruncated()); + } catch (Exception e) { + throw new AssertionError(e); + } + }); + } +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcFakeExecutionBehaviorTest.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcFakeExecutionBehaviorTest.java new file mode 100644 index 0000000000..0f74d12cb1 --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcFakeExecutionBehaviorTest.java @@ -0,0 +1,47 @@ +package com.dbx.agent.test; + +import com.dbx.agent.DatabaseAgent; +import com.dbx.agent.ExecuteQueryOptions; +import com.dbx.agent.QueryResult; +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public abstract class JdbcFakeExecutionBehaviorTest { + protected abstract DatabaseAgent createAgent(); + + protected abstract String resultSetSql(); + + @Test + void executesNonSelectStatementsThatReturnResultSets() { + DatabaseAgent agent = createAgent(); + TestSupport.setPrivateConnection(agent, JdbcAgentFake.connection()); + + QueryResult result = agent.executeQuery(resultSetSql(), null, new ExecuteQueryOptions()); + + assertEquals(Collections.singletonList("VALUE"), result.getColumns()); + assertEquals(Collections.singletonList(Collections.singletonList("row-value")), result.getRows()); + assertEquals(asList("setMaxRows:10001", "execute"), JdbcAgentFake.calls); + } + + @Test + void appliesExecutionRowAndFetchLimitsToJdbcStatements() { + DatabaseAgent agent = createAgent(); + TestSupport.setPrivateConnection(agent, JdbcAgentFake.connection()); + + QueryResult result = agent.executeQuery(resultSetSql(), null, new ExecuteQueryOptions(1, 25)); + + assertEquals(Collections.singletonList(Collections.singletonList("row-value")), result.getRows()); + assertEquals(asList("setMaxRows:2", "setFetchSize:25", "execute"), JdbcAgentFake.calls); + } + + private static java.util.List asList(String first, String second) { + return java.util.Arrays.asList(first, second); + } + + private static java.util.List asList(String first, String second, String third) { + return java.util.Arrays.asList(first, second, third); + } +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataBehaviorTest.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataBehaviorTest.java new file mode 100644 index 0000000000..9ef0b26669 --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataBehaviorTest.java @@ -0,0 +1,67 @@ +package com.dbx.agent.test; + +import com.dbx.agent.ColumnInfo; +import com.dbx.agent.TableInfo; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public abstract class JdbcMetadataBehaviorTest extends JdbcConnectedAgentTest { + protected abstract List metadataFixtureSql(); + + protected abstract String metadataSchema(); + + protected abstract List expectedTablesInOrder(); + + protected abstract String metadataColumnsTable(); + + protected abstract List expectedColumnsInOrder(); + + @Test + void returnsMetadataInStableOrder() { + withAgent("dbx-agent-metadata", agent -> { + for (String sql : metadataFixtureSql()) { + agent.executeQuery(sql, metadataSchema(), new com.dbx.agent.ExecuteQueryOptions()); + } + + List expectedTables = expectedTablesInOrder(); + List tableNames = new ArrayList(); + for (TableInfo table : agent.listTables(metadataSchema())) { + if (expectedTables.contains(table.getName())) { + tableNames.add(table.getName()); + } + } + assertEquals(expectedTables, tableNames); + + List expectedColumns = expectedColumnsInOrder(); + List columnNames = new ArrayList(); + for (ColumnInfo column : agent.getColumns(metadataSchema(), metadataColumnsTable())) { + if (expectedColumns.contains(column.getName())) { + columnNames.add(column.getName()); + } + } + assertEquals(expectedColumns, columnNames); + }); + } + + @Test + void buildsTableDdlFromMetadata() { + withAgent("dbx-agent-ddl", agent -> { + for (String sql : metadataFixtureSql()) { + agent.executeQuery(sql, metadataSchema(), new com.dbx.agent.ExecuteQueryOptions()); + } + + String ddl = agent.getTableDdl(metadataSchema(), metadataColumnsTable()); + + for (String column : expectedColumnsInOrder()) { + assertTrue(ddl.contains("\"" + column + "\""), "DDL should include column " + column + ": " + ddl); + } + assertTrue(ddl.contains("CREATE TABLE"), "DDL should include CREATE TABLE: " + ddl); + assertTrue(ddl.contains(metadataColumnsTable()), "DDL should include table name: " + ddl); + }); + } +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java new file mode 100644 index 0000000000..872a4c53e7 --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/JdbcMetadataSqlFake.java @@ -0,0 +1,172 @@ +package com.dbx.agent.test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; + +public final class JdbcMetadataSqlFake { + public static final JdbcMetadataSqlFake INSTANCE = new JdbcMetadataSqlFake(); + public static final List statements = new ArrayList(); + + private JdbcMetadataSqlFake() { + } + + public static Connection connection() { + statements.clear(); + return proxy(Connection.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("createStatement".equals(name)) { + return statement(); + } + if ("prepareStatement".equals(name)) { + statements.add((String) args[0]); + return preparedStatement(); + } + if ("getAutoCommit".equals(name)) { + return true; + } + if ("setAutoCommit".equals(name) || "commit".equals(name) || "rollback".equals(name) || "close".equals(name)) { + return null; + } + if ("isClosed".equals(name)) { + return false; + } + return defaultValue(method.getReturnType()); + } + }); + } + + public List getStatements() { + return statements; + } + + private static Statement statement() { + ResultSet resultSet = emptyResultSet(); + return proxy(Statement.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("execute".equals(name) || "executeQuery".equals(name)) { + statements.add((String) args[0]); + return "execute".equals(name) ? false : resultSet; + } + if ("getResultSet".equals(name)) { + return resultSet; + } + if ("getUpdateCount".equals(name)) { + return 0; + } + if ("setMaxRows".equals(name) || "close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static PreparedStatement preparedStatement() { + ResultSet resultSet = emptyResultSet(); + return proxy(PreparedStatement.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("execute".equals(name) || "executeQuery".equals(name)) { + return "execute".equals(name) ? false : resultSet; + } + if ("getResultSet".equals(name)) { + return resultSet; + } + if ("getUpdateCount".equals(name)) { + return 0; + } + if ("setString".equals(name)) { + statements.add("param:" + args[0] + "=" + args[1]); + return null; + } + if ("setMaxRows".equals(name) || "close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static ResultSet emptyResultSet() { + ResultSetMetaData metadata = proxy(ResultSetMetaData.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + if ("getColumnCount".equals(method.getName())) { + return 0; + } + return defaultValue(method.getReturnType()); + } + }); + return proxy(ResultSet.class, new MethodHandler() { + @Override + public Object handle(Method method, Object[] args) { + String name = method.getName(); + if ("next".equals(name)) { + return false; + } + if ("getMetaData".equals(name)) { + return metadata; + } + if ("close".equals(name)) { + return null; + } + return defaultValue(method.getReturnType()); + } + }); + } + + private static T proxy(Class type, final MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return handler.handle(method, args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (Boolean.TYPE.equals(type)) { + return false; + } + if (Byte.TYPE.equals(type)) { + return (byte) 0; + } + if (Short.TYPE.equals(type)) { + return (short) 0; + } + if (Integer.TYPE.equals(type)) { + return 0; + } + if (Long.TYPE.equals(type)) { + return 0L; + } + if (Float.TYPE.equals(type)) { + return 0f; + } + if (Double.TYPE.equals(type)) { + return 0.0d; + } + if (Character.TYPE.equals(type)) { + return '\0'; + } + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args); + } +} diff --git a/agents/test-support/src/main/java/com/dbx/agent/test/TestSupport.java b/agents/test-support/src/main/java/com/dbx/agent/test/TestSupport.java new file mode 100644 index 0000000000..aaddefa504 --- /dev/null +++ b/agents/test-support/src/main/java/com/dbx/agent/test/TestSupport.java @@ -0,0 +1,31 @@ +package com.dbx.agent.test; + +import java.lang.reflect.Field; +import java.sql.Connection; + +public final class TestSupport { + private TestSupport() { + } + + public static void setPrivateConnection(Object target, Connection connection) { + Field field = findConnectionField(target.getClass()); + field.setAccessible(true); + try { + field.set(target, connection); + } catch (IllegalAccessException e) { + throw new IllegalStateException("Unable to set connection", e); + } + } + + private static Field findConnectionField(Class type) { + Class current = type; + while (current != null) { + try { + return current.getDeclaredField("connection"); + } catch (NoSuchFieldException ignored) { + current = current.getSuperclass(); + } + } + throw new IllegalStateException(new NoSuchFieldException("connection")); + } +} diff --git a/agents/versions.json b/agents/versions.json new file mode 100644 index 0000000000..0f7aa57453 --- /dev/null +++ b/agents/versions.json @@ -0,0 +1,41 @@ +{ + "access": "0.1.19", + "bigquery": "0.1.23", + "cassandra": "0.1.22", + "dameng": "0.1.25", + "databend": "0.1.10", + "databricks": "0.1.17", + "db2": "0.1.22", + "exasol": "0.1.17", + "firebird": "0.1.17", + "gaussdb": "0.1.23", + "gbase8a": "0.1.17", + "gbase8s": "0.1.11", + "goldendb": "0.1.22", + "h2": "0.1.22", + "hive": "0.1.22", + "informix": "0.1.23", + "kingbase": "0.1.22", + "kylin": "0.1.22", + "neo4j": "0.1.22", + "oceanbase-oracle": "0.1.17", + "oracle": "0.1.26", + "oracle-legacy": "0.1.13", + "oracle-10g": "0.1.25", + "snowflake": "0.1.22", + "saphana": "0.1.17", + "sundb": "0.1.22", + "teradata": "0.1.17", + "tdengine": "0.1.20", + "trino": "0.1.22", + "vastbase": "0.1.22", + "vertica": "0.1.17", + "mongodb": "0.1.21", + "highgo": "0.1.20", + "yashandb": "0.1.18", + "xugu": "0.1.13", + "iris": "0.1.13", + "iotdb": "0.1.10", + "etcd": "0.1.10", + "zookeeper": "0.1.0" +} diff --git a/crates/dbx-core/src/agent_service.rs b/crates/dbx-core/src/agent_service.rs index 6c9de673f3..95764adaea 100644 --- a/crates/dbx-core/src/agent_service.rs +++ b/crates/dbx-core/src/agent_service.rs @@ -128,7 +128,7 @@ fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result } } -const REGISTRY_PATH: &str = "https://github.com/t8y2/dbx-agents/releases/latest/download/agent-registry.json"; +const REGISTRY_PATH: &str = "https://github.com/t8y2/dbx/releases/download/agents-latest/agent-registry.json"; const REGISTRY_R2_PATH: &str = "agents/agent-registry.json"; static REGISTRY_CACHE: std::sync::LazyLock>> = @@ -255,6 +255,9 @@ pub fn jre_needs_install(am: &AgentManager, registry: &AgentRegistry, jre_key: & pub fn local_agent_jar_candidates(db_type: &str) -> Vec { let jar_name = format!("dbx-agent-{db_type}.jar"); + let monorepo_driver = + PathBuf::from("agents").join("drivers").join(db_type).join("build").join("libs").join(&jar_name); + let monorepo_legacy = PathBuf::from("agents").join(db_type).join("build").join("libs").join(&jar_name); let relative_driver = PathBuf::from("..").join("dbx-agents").join("drivers").join(db_type).join("build").join("libs").join(&jar_name); let nested_driver = @@ -262,7 +265,7 @@ pub fn local_agent_jar_candidates(db_type: &str) -> Vec { let relative_legacy = PathBuf::from("..").join("dbx-agents").join(db_type).join("build").join("libs").join(&jar_name); let nested_legacy = PathBuf::from("dbx-agents").join(db_type).join("build").join("libs").join(&jar_name); - vec![relative_driver, nested_driver, relative_legacy, nested_legacy] + vec![monorepo_driver, monorepo_legacy, relative_driver, nested_driver, relative_legacy, nested_legacy] } pub fn find_local_agent_jar(db_type: &str) -> Option { diff --git a/crates/dbx-core/tests/agent_service.rs b/crates/dbx-core/tests/agent_service.rs index 5aff015287..a1f6f2d8e0 100644 --- a/crates/dbx-core/tests/agent_service.rs +++ b/crates/dbx-core/tests/agent_service.rs @@ -228,9 +228,10 @@ fn jre_needs_install_when_managed_runtime_version_differs() { } #[test] -fn local_agent_jar_candidates_include_sibling_build_output() { +fn local_agent_jar_candidates_include_monorepo_and_legacy_build_output() { let candidates = local_agent_jar_candidates("tdengine"); + assert!(candidates.iter().any(|path| path.ends_with("agents/drivers/tdengine/build/libs/dbx-agent-tdengine.jar"))); assert!(candidates.iter().any(|path| path.ends_with("dbx-agents/tdengine/build/libs/dbx-agent-tdengine.jar"))); } @@ -244,6 +245,10 @@ fn github_agent_asset_urls_map_to_r2_paths_by_category() { github_url_to_r2_path("https://github.com/t8y2/dbx-agents/releases/download/v1/dbx-agent-h2.jar", "driver"), "agents/drivers/dbx-agent-h2.jar" ); + assert_eq!( + github_url_to_r2_path("https://github.com/t8y2/dbx/releases/download/agents-v0.3.0/dbx-agent-h2.jar", "driver"), + "agents/drivers/dbx-agent-h2.jar" + ); } fn test_path(name: &str) -> std::path::PathBuf { diff --git a/docs/lib/agentRegistry.ts b/docs/lib/agentRegistry.ts index 0d11e26f27..ceb205d6db 100644 --- a/docs/lib/agentRegistry.ts +++ b/docs/lib/agentRegistry.ts @@ -34,7 +34,7 @@ export interface AgentRegistry { } const AGENT_REGISTRY_R2_URL = "https://dl.dbxio.com/agents/agent-registry.json"; -const AGENT_REGISTRY_GITHUB_URL = "https://github.com/t8y2/dbx-agents/releases/latest/download/agent-registry.json"; +const AGENT_REGISTRY_GITHUB_URL = "https://github.com/t8y2/dbx/releases/download/agents-latest/agent-registry.json"; export async function fetchAgentRegistry(): Promise { const urls = [AGENT_REGISTRY_R2_URL, AGENT_REGISTRY_GITHUB_URL];